[MC] Fold A-B when A is a pending label or A/B are separated by a MCFillFragment

When the MCAssembler is non-null and the MCAsmLayout is null, we can fold A-B
in these additional cases:

* when A is a pending label (will be reassigned to a real fragment in flushPendingLabels())
* A and B are separated by a MCFillFragment with a constant size

1 files changed, 19 insertions, 4 deletions

diff --git a/llvm/lib/MC/MCExpr.cpp b/llvm/lib/MC/MCExpr.cpp
index 3d47a52d..d5b5aa8 100644
--- a/llvm/lib/MC/MCExpr.cpp
+++ b/llvm/lib/MC/MCExpr.cpp
@@ -658,15 +658,30 @@ static void AttemptToFoldSymbolOffsetDifference(
     // 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();
+    bool Found = false;
     for (auto FI = FB->getIterator(), FE = SecA.end(); FI != FE; ++FI) {
+      auto DF = dyn_cast<MCDataFragment>(FI);
       if (&*FI == FA) {
-        Addend += Displacement;
-        return FinalizeFolding();
+        Found = true;
+        break;
       }
 
-      if (FI->getKind() != MCFragment::FT_Data)
+      int64_t Num;
+      if (DF) {
+        Displacement += DF->getContents().size();
+      } else if (auto *FF = dyn_cast<MCFillFragment>(FI);
+                 FF && FF->getNumValues().evaluateAsAbsolute(Num)) {
+        Displacement += Num * FF->getValueSize();
+      } else {
         return;
-      Displacement += cast<MCDataFragment>(FI)->getContents().size();
+      }
+    }
+    // If FA is found or if FA is a dummy fragment not in the fragment list,
+    // (which means SA is a pending label (see flushPendingLabels)), we can
+    // resolve the difference.
+    if (Found || isa<MCDummyFragment>(FA)) {
+      Addend += Displacement;
+      FinalizeFolding();
     }
   }
 }