[BOLT][X86] Fix getTargetSymbol()

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;

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);

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,

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,

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