[MC] Resolve the difference of symbols in consecutive MCDataFragements

Try to resolve the difference of two symbols in consecutive MCDataFragments.
This is important for an idiom like "foo:instr; .if . - foo; instr; .endif"
(https://bugs.llvm.org/show_bug.cgi?id=43795).

Reviewed By: nickdesaulniers

Differential Revision: https://reviews.llvm.org/D69411

llvm/include/llvm/MC/MCFragment.h

@@ -64,6 +64,10 @@ class MCFragment : public ilist_node_with_parent<MCFragment, MCSection> {
   /// The layout order of this fragment.
   unsigned LayoutOrder;
 
+  /// The subsection this fragment belongs to. This is 0 if the fragment is not
+  // in any subsection.
+  unsigned SubsectionNumber = 0;
+
   FragmentType Kind;
 
   /// Whether fragment is being laid out.
@@ -102,6 +106,9 @@ class MCFragment : public ilist_node_with_parent<MCFragment, MCSection> {
   bool hasInstructions() const { return HasInstructions; }
 
   void dump() const;
+
+  void setSubsectionNumber(unsigned Value) { SubsectionNumber = Value; }
+  unsigned getSubsectionNumber() const { return SubsectionNumber; }
 };
 
 class MCDummyFragment : public MCFragment {

llvm/lib/MC/MCExpr.cpp

@@ -588,12 +588,7 @@ static void AttemptToFoldSymbolOffsetDifference(
   if (!Asm->getWriter().isSymbolRefDifferenceFullyResolved(*Asm, A, B, InSet))
     return;
 
-  MCFragment *FA = SA.getFragment();
-  MCFragment *FB = SB.getFragment();
-  if (FA == FB && !SA.isVariable() && !SA.isUnset() && !SB.isVariable() &&
-      !SB.isUnset()) {
-    Addend += (SA.getOffset() - SB.getOffset());
-
+  auto FinalizeFolding = [&]() {
     // Pointers to Thumb symbols need to have their low-bit set to allow
     // for interworking.
     if (Asm->isThumbFunc(&SA))
@@ -607,42 +602,64 @@ static void AttemptToFoldSymbolOffsetDifference(
     // Clear the symbol expr pointers to indicate we have folded these
     // operands.
     A = B = nullptr;
-    return;
-  }
+  };
 
-  if (!Layout)
-    return;
+  const MCFragment *FA = SA.getFragment();
+  const MCFragment *FB = SB.getFragment();
+  // If both symbols are in the same fragment, return the difference of their
+  // offsets
+  if (FA == FB && !SA.isVariable() && !SA.isUnset() && !SB.isVariable() &&
+      !SB.isUnset()) {
+    Addend += SA.getOffset() - SB.getOffset();
+    return FinalizeFolding();
+  }
 
   const MCSection &SecA = *FA->getParent();
   const MCSection &SecB = *FB->getParent();
 
   if ((&SecA != &SecB) && !Addrs)
     return;
 
-  // One of the symbol involved is part of a fragment being laid out. Quit now
-  // to avoid a self loop.
-  if (!Layout->canGetFragmentOffset(FA) || !Layout->canGetFragmentOffset(FB))
-    return;
+  if (Layout) {
+    // One of the symbol involved is part of a fragment being laid out. Quit now
+    // to avoid a self loop.
+    if (!Layout->canGetFragmentOffset(FA) || !Layout->canGetFragmentOffset(FB))
+      return;
+
+    // Eagerly evaluate when layout is finalized.
+    Addend += Layout->getSymbolOffset(A->getSymbol()) -
+              Layout->getSymbolOffset(B->getSymbol());
+    if (Addrs && (&SecA != &SecB))
+      Addend += (Addrs->lookup(&SecA) - Addrs->lookup(&SecB));
+
+    FinalizeFolding();
+  } else {
+    // When layout is not finalized, our ability to resolve differences between
+    // symbols is limited to specific cases where the fragments between two
+    // symbols (including the fragments the symbols are defined in) are
+    // fixed-size fragments so the difference can be calculated. For example,
+    // this is important when the Subtarget is changed and a new MCDataFragment
+    // is created in the case of foo: instr; .arch_extension ext; instr .if . -
+    // foo.
+    if (SA.isVariable() || SA.isUnset() || SB.isVariable() || SB.isUnset() ||
+        FA->getKind() != MCFragment::FT_Data ||
+        FB->getKind() != MCFragment::FT_Data ||
+        FA->getSubsectionNumber() != FB->getSubsectionNumber())
+      return;
+    // Try to find a constant displacement from FA to FB, add the displacement
+    // between the offset in FA of SA and the offset in FB of SB.
+    int64_t Displacement = SA.getOffset() - SB.getOffset();
+    for (auto FI = FB->getIterator(), FE = SecA.end(); FI != FE; ++FI) {
+      if (&*FI == FA) {
+        Addend += Displacement;
+        return FinalizeFolding();
+      }
 
-  // Eagerly evaluate.
-  Addend += Layout->getSymbolOffset(A->getSymbol()) -
-            Layout->getSymbolOffset(B->getSymbol());
-  if (Addrs && (&SecA != &SecB))
-    Addend += (Addrs->lookup(&SecA) - Addrs->lookup(&SecB));
-
-  // Pointers to Thumb symbols need to have their low-bit set to allow
-  // for interworking.
-  if (Asm->isThumbFunc(&SA))
-    Addend |= 1;
-
-  // If symbol is labeled as micromips, we set low-bit to ensure
-  // correct offset in .gcc_except_table
-  if (Asm->getBackend().isMicroMips(&SA))
-    Addend |= 1;
-
-  // Clear the symbol expr pointers to indicate we have folded these
-  // operands.
-  A = B = nullptr;
+      if (FI->getKind() != MCFragment::FT_Data)
+        return;
+      Displacement += cast<MCDataFragment>(FI)->getContents().size();
+    }
+  }
 }
 
 static bool canFold(const MCAssembler *Asm, const MCSymbolRefExpr *A,

llvm/lib/MC/MCSection.cpp

@@ -82,6 +82,7 @@ MCSection::getSubsectionInsertionPoint(unsigned Subsection) {
     SubsectionFragmentMap.insert(MI, std::make_pair(Subsection, F));
     getFragmentList().insert(IP, F);
     F->setParent(this);
+    F->setSubsectionNumber(Subsection);
   }
 
   return IP;

llvm/test/MC/ARM/directive-if-subtraction.s

@@ -0,0 +1,52 @@
+// RUN: llvm-mc -triple armv7a-linux-gnueabihf %s -filetype=obj -o /dev/null 2>&1 | FileCheck --check-prefix=OBJ --allow-empty %s
+// RUN: not llvm-mc -triple armv7a-linux-gnueabihf %s -o /dev/null 2>&1 | FileCheck --check-prefix=ASM %s
+// RUN: llvm-mc -triple armv7a-linux-gnueabihf %s -filetype=obj -o - | llvm-objdump -d - | FileCheck --check-prefix=DISASM %s
+
+nop
+// Create a new MCDataFragment due to Subtarget change
+.arch_extension sec
+9997:nop
+.if . - 9997b == 0
+// OBJ-NOT:[[@LINE-1]]:5: error: expected absolute expression
+// ASM:[[@LINE-2]]:5: error: expected absolute expression
+// DISASM: orr	r1, r1, #2
+orr r1, r1, #1
+.else
+orr r1, r1, #2
+.endif
+
+
+
+@ RUN: not llvm-mc -filetype=obj -triple arm-linux-gnueabihf --defsym=ERR=1 %s -o /dev/null 2>&1 | FileCheck --check-prefix=ARM-ERR %s
+@ RUN: not llvm-mc -filetype=obj -triple thumbv7a-linux-gnueabihf --defsym=ERR=1 %s -o /dev/null 2>&1 | FileCheck --check-prefix=THUMB2-ERR %s
+
+.ifdef ERR
+9997: nop
+      .align 4
+      nop
+.if . - 9997b == 4
+// ARM-ERR:[[@LINE-1]]:5: error: expected absolute expression
+.endif
+
+9997: nop
+      .space 4
+      nop
+.if . - 9997b == 4
+// ARM-ERR:[[@LINE-1]]:5: error: expected absolute expression
+.endif
+
+9997:
+      ldr r0,=0x12345678
+      .ltorg
+      nop
+.if . - 9997b == 4
+// ARM-ERR:[[@LINE-1]]:5: error: expected absolute expression
+.endif
+
+9997: nop
+      b external
+      nop
+.if . - 9997b == 4
+// THUMB2-ERR:[[@LINE-1]]:5: error: expected absolute expression
+.endif
+.endif

llvm/test/MC/MachO/reloc-diff.s

@@ -22,9 +22,5 @@ Ltemp:
 // CHECK-NEXT:     0x0 0 2 n/a GENERIC_RELOC_PAIR 1 0x0
 // CHECK-NEXT:     0x8 0 2 n/a GENERIC_RELOC_LOCAL_SECTDIFF 1 0x0
 // CHECK-NEXT:     0x0 0 2 n/a GENERIC_RELOC_PAIR 1 0x0
-// CHECK-NEXT:     0x4 0 2 n/a GENERIC_RELOC_LOCAL_SECTDIFF 1 0x0
-// CHECK-NEXT:     0x0 0 2 n/a GENERIC_RELOC_PAIR 1 0x0
-// CHECK-NEXT:     0x0 0 2 n/a GENERIC_RELOC_SECTDIFF 1 0x0
-// CHECK-NEXT:     0x0 0 2 n/a GENERIC_RELOC_PAIR 1 0x0
 // CHECK-NEXT:   }
 // CHECK-NEXT: ]