From 507efbcce03d8c2c5dbea3028bc39f02c88fea79 Mon Sep 17 00:00:00 2001
From: Fangrui Song <i@maskray.me>
Date: Sun, 18 Jun 2023 17:18:38 -0700
Subject: [PATCH] [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
---
 llvm/lib/MC/MCExpr.cpp                        | 23 +++++++++++++++----
 llvm/test/MC/ARM/directive-if-subtraction.s   |  3 +--
 .../test/MC/AsmParser/assembler-expressions.s | 12 +++++++++-
 llvm/test/MC/RISCV/cfi-advance.s              | 10 ++++----
 4 files changed, 36 insertions(+), 12 deletions(-)

diff --git a/llvm/lib/MC/MCExpr.cpp b/llvm/lib/MC/MCExpr.cpp
index 3d47a52d8a3b1..d5b5aa8fba4e0 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();
     }
   }
 }
diff --git a/llvm/test/MC/ARM/directive-if-subtraction.s b/llvm/test/MC/ARM/directive-if-subtraction.s
index edb386593ba63..e67fe0ba572c9 100644
--- a/llvm/test/MC/ARM/directive-if-subtraction.s
+++ b/llvm/test/MC/ARM/directive-if-subtraction.s
@@ -1,5 +1,5 @@
 // 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: not llvm-mc -triple armv7a-linux-gnueabihf %s -o /dev/null 2>&1 | FileCheck --check-prefix=ASM %s --implicit-check-not=error:
 // RUN: llvm-mc -triple armv7a-linux-gnueabihf %s -filetype=obj -o - | llvm-objdump -d - | FileCheck --check-prefix=DISASM %s
 
 nop
@@ -32,7 +32,6 @@ orr r1, r1, #2
       .space 4
       nop
 .if . - 9997b == 4
-// ARM-ERR:[[@LINE-1]]:5: error: expected absolute expression
 .endif
 
 9997:
diff --git a/llvm/test/MC/AsmParser/assembler-expressions.s b/llvm/test/MC/AsmParser/assembler-expressions.s
index 4e3c507c9a475..eb967adcabfcd 100644
--- a/llvm/test/MC/AsmParser/assembler-expressions.s
+++ b/llvm/test/MC/AsmParser/assembler-expressions.s
@@ -4,7 +4,7 @@
 .data
 
 # OBJDATA: Contents of section .data
-# OBJDATA-NEXT: 0000 aa01aa05 06ff
+# OBJDATA-NEXT: 0000 aa01aa05 06000000 08ff
 
 foo1:
 # ASM-ERR: :[[#@LINE+1]]:5: error: expected absolute expression
@@ -32,6 +32,16 @@ foo3:
    .byte 7
 .endif
 
+foo4:
+.byte 0
+.space 2
+# ASM-ERR: :[[#@LINE+1]]:5: error: expected absolute expression
+.if . - foo4 == 3
+    .byte 8
+.else
+    .byte 9
+.endif
+
 .byte 0xff
 
 # nop is a fixed size instruction so this should pass.
diff --git a/llvm/test/MC/RISCV/cfi-advance.s b/llvm/test/MC/RISCV/cfi-advance.s
index b599b50692dbc..5520c7ad330df 100644
--- a/llvm/test/MC/RISCV/cfi-advance.s
+++ b/llvm/test/MC/RISCV/cfi-advance.s
@@ -3,11 +3,11 @@
 # RUN: llvm-dwarfdump --debug-frame %t.o 2>&1 \
 # RUN:     | FileCheck -check-prefix=CHECK-DWARFDUMP %s
 
-# CHECK: 0x26 R_RISCV_SET8 - 0x0
-# CHECK-NEXT: 0x26 R_RISCV_SUB8 - 0x0
-# CHECK-NEXT: 0x2A R_RISCV_SET16 - 0x0
-# CHECK-NEXT: 0x2A R_RISCV_SUB16 - 0x0
-# CHECK-NEXT: 0x2F R_RISCV_SET32 - 0x0
+# CHECK:      .rela.eh_frame {
+# CHECK-NEXT:   0x1C R_RISCV_32_PCREL - 0x0
+# CHECK-NEXT:   0x20 R_RISCV_ADD32 - 0x0
+# CHECK-NEXT:   0x20 R_RISCV_SUB32 - 0x0
+# CHECK-NEXT: }
 # CHECK-DWARFDUMP: DW_CFA_advance_loc1
 # CHECK-DWARFDUMP-NEXT: DW_CFA_def_cfa_offset
 # CHECK-DWARFDUMP-NEXT: DW_CFA_advance_loc2