[BOLT][X86] Fix getTargetSymbol()

[maksfb]

In 96e5ee2, I inadvertently broke the way non-trivial symbol references got updated from non-optimized code. The breakage was a consequence of getTargetSymbol(MCExpr *) not returning a symbol when the parameter was a binary expression. Fix getTargetSymbol() to cover such cases.

[llvmbot]

@llvm/pr-subscribers-bolt

Author: Maksim Panchenko (maksfb)

Changes

In 96e5ee2, I inadvertently broke the way non-trivial symbol references got updated from non-optimized code. The breakage was a consequence of getTargetSymbol(MCExpr *) not returning a symbol when the parameter was a binary expression. Fix getTargetSymbol() to cover such cases.

---

Full diff: https://github.com/llvm/llvm-project/pull/133834.diff

5 Files Affected:

• (modified) bolt/include/bolt/Core/MCPlusBuilder.h (+5-1)
• (modified) bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp (+2-10)
• (modified) bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp (+1-9)
• (modified) bolt/lib/Target/X86/X86MCPlusBuilder.cpp (+1-5)
• (added) bolt/test/X86/lite-mode-target-expr.s (+33)

diff --git a/bolt/include/bolt/Core/MCPlusBuilder.h b/bolt/include/bolt/Core/MCPlusBuilder.h
index 6860f021eb849..1458d36d4813a 100644
--- a/bolt/include/bolt/Core/MCPlusBuilder.h
+++ b/bolt/include/bolt/Core/MCPlusBuilder.h
@@ -1266,7 +1266,11 @@ class MCPlusBuilder {
 
   /// Return MCSymbol extracted from the expression.
   virtual const MCSymbol *getTargetSymbol(const MCExpr *Expr) const {
-    if (auto *SymbolRefExpr = dyn_cast<const MCSymbolRefExpr>(Expr))
+    if (auto *BinaryExpr = dyn_cast<const MCBinaryExpr>(Expr))
+      return getTargetSymbol(BinaryExpr->getLHS());
+
+    auto *SymbolRefExpr = dyn_cast<const MCSymbolRefExpr>(Expr);
+    if (SymbolRefExpr && SymbolRefExpr->getKind() == MCSymbolRefExpr::VK_None)
       return &SymbolRefExpr->getSymbol();
 
     return nullptr;
diff --git a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
index d238a1df5c7d7..0fd127bfeba41 100644
--- a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
+++ b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
@@ -862,20 +862,12 @@ class AArch64MCPlusBuilder : public MCPlusBuilder {
     if (AArchExpr && AArchExpr->getSubExpr())
       return getTargetSymbol(AArchExpr->getSubExpr());
 
-    auto *BinExpr = dyn_cast<MCBinaryExpr>(Expr);
-    if (BinExpr)
-      return getTargetSymbol(BinExpr->getLHS());
-
-    auto *SymExpr = dyn_cast<MCSymbolRefExpr>(Expr);
-    if (SymExpr && SymExpr->getKind() == MCSymbolRefExpr::VK_None)
-      return &SymExpr->getSymbol();
-
-    return nullptr;
+    return MCPlusBuilder::getTargetSymbol(Expr);
   }
 
   const MCSymbol *getTargetSymbol(const MCInst &Inst,
                                   unsigned OpNum = 0) const override {
-    if (!getSymbolRefOperandNum(Inst, OpNum))
+    if (!OpNum && !getSymbolRefOperandNum(Inst, OpNum))
       return nullptr;
 
     const MCOperand &Op = Inst.getOperand(OpNum);
diff --git a/bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp b/bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp
index edbbdc491aee4..4320c679acd54 100644
--- a/bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp
+++ b/bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp
@@ -338,15 +338,7 @@ class RISCVMCPlusBuilder : public MCPlusBuilder {
     if (RISCVExpr && RISCVExpr->getSubExpr())
       return getTargetSymbol(RISCVExpr->getSubExpr());
 
-    auto *BinExpr = dyn_cast<MCBinaryExpr>(Expr);
-    if (BinExpr)
-      return getTargetSymbol(BinExpr->getLHS());
-
-    auto *SymExpr = dyn_cast<MCSymbolRefExpr>(Expr);
-    if (SymExpr && SymExpr->getKind() == MCSymbolRefExpr::VK_None)
-      return &SymExpr->getSymbol();
-
-    return nullptr;
+    return MCPlusBuilder::getTargetSymbol(Expr);
   }
 
   const MCSymbol *getTargetSymbol(const MCInst &Inst,
diff --git a/bolt/lib/Target/X86/X86MCPlusBuilder.cpp b/bolt/lib/Target/X86/X86MCPlusBuilder.cpp
index 4016ffe18dc02..0b2617600f5c0 100644
--- a/bolt/lib/Target/X86/X86MCPlusBuilder.cpp
+++ b/bolt/lib/Target/X86/X86MCPlusBuilder.cpp
@@ -1796,11 +1796,7 @@ class X86MCPlusBuilder : public MCPlusBuilder {
     if (!Op.isExpr())
       return nullptr;
 
-    auto *SymExpr = dyn_cast<MCSymbolRefExpr>(Op.getExpr());
-    if (!SymExpr || SymExpr->getKind() != MCSymbolRefExpr::VK_None)
-      return nullptr;
-
-    return &SymExpr->getSymbol();
+    return MCPlusBuilder::getTargetSymbol(Op.getExpr());
   }
 
   bool analyzeBranch(InstructionIterator Begin, InstructionIterator End,
diff --git a/bolt/test/X86/lite-mode-target-expr.s b/bolt/test/X86/lite-mode-target-expr.s
new file mode 100644
index 0000000000000..5480748c20f05
--- /dev/null
+++ b/bolt/test/X86/lite-mode-target-expr.s
@@ -0,0 +1,33 @@
+## Check that llvm-bolt properly updates references in unoptimized code when
+## such references are non-trivial expressions.
+
+# RUN: %clang %cflags %s -o %t.exe -Wl,-q -no-pie
+# RUN: llvm-bolt %t.exe -o %t.bolt --funcs=_start
+# RUN: llvm-objdump -d --disassemble-symbols=_start %t.bolt > %t.out
+# RUN: llvm-objdump -d --disassemble-symbols=cold %t.bolt >> %t.out
+# RUN: FileCheck %s < %t.out
+
+## _start() will be optimized and assigned a new address.
+# CHECK: [[#%x,ADDR:]] <_start>:
+
+## cold() is not optimized, but references to _start are updated.
+# CHECK-LABEL: <cold>:
+# CHECK-NEXT: movl $0x[[#ADDR - 1]], %ecx
+# CHECK-NEXT: movl $0x[[#ADDR]], %ecx
+# CHECK-NEXT: movl $0x[[#ADDR + 1]], %ecx
+
+  .text
+  .globl cold
+  .type cold, %function
+cold:
+	movl $_start-1, %ecx
+	movl $_start, %ecx
+	movl $_start+1, %ecx
+  ret
+  .size cold, .-cold
+
+  .globl _start
+  .type _start, %function
+_start:
+	ret
+  .size _start, .-_start