1 files changed, 127 insertions, 5 deletions

diff --git a/llvm/lib/MC/MCAssembler.cpp b/llvm/lib/MC/MCAssembler.cpp
index dd563d8..e6e77d3 100644
--- a/llvm/lib/MC/MCAssembler.cpp
+++ b/llvm/lib/MC/MCAssembler.cpp
@@ -383,6 +383,16 @@ uint64_t MCAssembler::computeFragmentSize(const MCAsmLayout &Layout,
   llvm_unreachable("invalid fragment kind");
 }
 
+MCAsmLayout::MCAsmLayout(MCAssembler &Asm) : Assembler(Asm) {
+  // Compute the section layout order. Virtual sections must go last.
+  for (MCSection &Sec : Asm)
+    if (!Sec.isVirtualSection())
+      SectionOrder.push_back(&Sec);
+  for (MCSection &Sec : Asm)
+    if (Sec.isVirtualSection())
+      SectionOrder.push_back(&Sec);
+}
+
 // Compute the amount of padding required before the fragment \p F to
 // obey bundling restrictions, where \p FOffset is the fragment's offset in
 // its section and \p FSize is the fragment's size.
@@ -464,11 +474,6 @@ void MCAsmLayout::layoutBundle(MCFragment *Prev, MCFragment *F) {
       DF->Offset = EF->Offset;
 }
 
-uint64_t MCAsmLayout::getFragmentOffset(const MCFragment *F) const {
-  ensureValid(F);
-  return F->Offset;
-}
-
 void MCAsmLayout::ensureValid(const MCFragment *Frag) const {
   MCSection &Sec = *Frag->getParent();
   if (Sec.hasLayout())
@@ -487,6 +492,123 @@ void MCAsmLayout::ensureValid(const MCFragment *Frag) const {
   }
 }
 
+uint64_t MCAsmLayout::getFragmentOffset(const MCFragment *F) const {
+  ensureValid(F);
+  return F->Offset;
+}
+
+// Simple getSymbolOffset helper for the non-variable case.
+static bool getLabelOffset(const MCAsmLayout &Layout, const MCSymbol &S,
+                           bool ReportError, uint64_t &Val) {
+  if (!S.getFragment()) {
+    if (ReportError)
+      report_fatal_error("unable to evaluate offset to undefined symbol '" +
+                         S.getName() + "'");
+    return false;
+  }
+  Val = Layout.getFragmentOffset(S.getFragment()) + S.getOffset();
+  return true;
+}
+
+static bool getSymbolOffsetImpl(const MCAsmLayout &Layout, const MCSymbol &S,
+                                bool ReportError, uint64_t &Val) {
+  if (!S.isVariable())
+    return getLabelOffset(Layout, S, ReportError, Val);
+
+  // If SD is a variable, evaluate it.
+  MCValue Target;
+  if (!S.getVariableValue()->evaluateAsValue(Target, Layout))
+    report_fatal_error("unable to evaluate offset for variable '" +
+                       S.getName() + "'");
+
+  uint64_t Offset = Target.getConstant();
+
+  const MCSymbolRefExpr *A = Target.getSymA();
+  if (A) {
+    uint64_t ValA;
+    // FIXME: On most platforms, `Target`'s component symbols are labels from
+    // having been simplified during evaluation, but on Mach-O they can be
+    // variables due to PR19203. This, and the line below for `B` can be
+    // restored to call `getLabelOffset` when PR19203 is fixed.
+    if (!getSymbolOffsetImpl(Layout, A->getSymbol(), ReportError, ValA))
+      return false;
+    Offset += ValA;
+  }
+
+  const MCSymbolRefExpr *B = Target.getSymB();
+  if (B) {
+    uint64_t ValB;
+    if (!getSymbolOffsetImpl(Layout, B->getSymbol(), ReportError, ValB))
+      return false;
+    Offset -= ValB;
+  }
+
+  Val = Offset;
+  return true;
+}
+
+bool MCAsmLayout::getSymbolOffset(const MCSymbol &S, uint64_t &Val) const {
+  return getSymbolOffsetImpl(*this, S, false, Val);
+}
+
+uint64_t MCAsmLayout::getSymbolOffset(const MCSymbol &S) const {
+  uint64_t Val;
+  getSymbolOffsetImpl(*this, S, true, Val);
+  return Val;
+}
+
+const MCSymbol *MCAsmLayout::getBaseSymbol(const MCSymbol &Symbol) const {
+  if (!Symbol.isVariable())
+    return &Symbol;
+
+  const MCExpr *Expr = Symbol.getVariableValue();
+  MCValue Value;
+  if (!Expr->evaluateAsValue(Value, *this)) {
+    Assembler.getContext().reportError(Expr->getLoc(),
+                                       "expression could not be evaluated");
+    return nullptr;
+  }
+
+  const MCSymbolRefExpr *RefB = Value.getSymB();
+  if (RefB) {
+    Assembler.getContext().reportError(
+        Expr->getLoc(),
+        Twine("symbol '") + RefB->getSymbol().getName() +
+            "' could not be evaluated in a subtraction expression");
+    return nullptr;
+  }
+
+  const MCSymbolRefExpr *A = Value.getSymA();
+  if (!A)
+    return nullptr;
+
+  const MCSymbol &ASym = A->getSymbol();
+  const MCAssembler &Asm = getAssembler();
+  if (ASym.isCommon()) {
+    Asm.getContext().reportError(Expr->getLoc(),
+                                 "Common symbol '" + ASym.getName() +
+                                     "' cannot be used in assignment expr");
+    return nullptr;
+  }
+
+  return &ASym;
+}
+
+uint64_t MCAsmLayout::getSectionAddressSize(const MCSection *Sec) const {
+  // The size is the last fragment's end offset.
+  const MCFragment &F = *Sec->curFragList()->Tail;
+  return getFragmentOffset(&F) + getAssembler().computeFragmentSize(*this, F);
+}
+
+uint64_t MCAsmLayout::getSectionFileSize(const MCSection *Sec) const {
+  // Virtual sections have no file size.
+  if (Sec->isVirtualSection())
+    return 0;
+
+  // Otherwise, the file size is the same as the address space size.
+  return getSectionAddressSize(Sec);
+}
+
 bool MCAssembler::registerSymbol(const MCSymbol &Symbol) {
   bool Changed = !Symbol.isRegistered();
   if (Changed) {