From b2d272ccfb9407ded7e54d1eabd5b5743aa9dd1b Mon Sep 17 00:00:00 2001
From: Maksim Panchenko <maks@fb.com>
Date: Mon, 31 Mar 2025 18:31:33 -0700
Subject: [PATCH] [BOLT][X86] Fix getTargetSymbol() (#133834)

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.
---
 bolt/include/bolt/Core/MCPlusBuilder.h        |  6 +++-
 .../Target/AArch64/AArch64MCPlusBuilder.cpp   | 12 ++-----
 bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp  | 10 +-----
 bolt/lib/Target/X86/X86MCPlusBuilder.cpp      |  6 +---
 bolt/test/X86/lite-mode-target-expr.s         | 33 +++++++++++++++++++
 5 files changed, 42 insertions(+), 25 deletions(-)
 create mode 100644 bolt/test/X86/lite-mode-target-expr.s

diff --git a/bolt/include/bolt/Core/MCPlusBuilder.h b/bolt/include/bolt/Core/MCPlusBuilder.h
index 6860f021eb84..1458d36d4813 100644
--- a/bolt/include/bolt/Core/MCPlusBuilder.h
+++ b/bolt/include/bolt/Core/MCPlusBuilder.h
@@ -1266,7 +1266,11 @@ public:
 
   /// 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 d238a1df5c7d..0fd127bfeba4 100644
--- a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
+++ b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
@@ -862,20 +862,12 @@ public:
     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 edbbdc491aee..4320c679acd5 100644
--- a/bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp
+++ b/bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp
@@ -338,15 +338,7 @@ public:
     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 4016ffe18dc0..0b2617600f5c 100644
--- a/bolt/lib/Target/X86/X86MCPlusBuilder.cpp
+++ b/bolt/lib/Target/X86/X86MCPlusBuilder.cpp
@@ -1796,11 +1796,7 @@ public:
     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 000000000000..5480748c20f0
--- /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