[llvm] [BOLT] Avoid reference updates for non-JT symbol operands (PR #88838)

[via llvm-commits]

https://github.com/linsinan1995 updated https://github.com/llvm/llvm-project/pull/88838

>From 999bddf16896c6f9332b1deb039e3f719ce9c1a6 Mon Sep 17 00:00:00 2001
From: Sinan Lin <sinan.lin at linux.alibaba.com>
Date: Tue, 16 Apr 2024 10:35:48 +0800
Subject: [PATCH] [BOLT] Avoid reference updates for non-JT symbol operands

Add a check to skip updating references for operands that do not directly
refer to jump table symbols but fall within a jump table's address
range to prevent unintended modifications.
---
 bolt/lib/Passes/ValidateMemRefs.cpp        |  8 +--
 bolt/test/X86/jt-symbol-disambiguation-4.s | 60 ++++++++++++++++++++++
 2 files changed, 64 insertions(+), 4 deletions(-)
 create mode 100644 bolt/test/X86/jt-symbol-disambiguation-4.s

diff --git a/bolt/lib/Passes/ValidateMemRefs.cpp b/bolt/lib/Passes/ValidateMemRefs.cpp
index f29a97c43f497c..ca58493b279c9e 100644
--- a/bolt/lib/Passes/ValidateMemRefs.cpp
+++ b/bolt/lib/Passes/ValidateMemRefs.cpp
@@ -29,8 +29,7 @@ bool ValidateMemRefs::checkAndFixJTReference(BinaryFunction &BF, MCInst &Inst,
   if (!BD)
     return false;
 
-  const uint64_t TargetAddress = BD->getAddress() + Offset;
-  JumpTable *JT = BC.getJumpTableContainingAddress(TargetAddress);
+  JumpTable *JT = BC.getJumpTableContainingAddress(BD->getAddress());
   if (!JT)
     return false;
 
@@ -43,8 +42,9 @@ bool ValidateMemRefs::checkAndFixJTReference(BinaryFunction &BF, MCInst &Inst,
   // the jump table label with a regular rodata reference. Get a
   // non-JT reference by fetching the symbol 1 byte before the JT
   // label.
-  MCSymbol *NewSym = BC.getOrCreateGlobalSymbol(TargetAddress - 1, "DATAat");
-  BC.MIB->setOperandToSymbolRef(Inst, OperandNum, NewSym, 1, &*BC.Ctx, 0);
+  MCSymbol *NewSym = BC.getOrCreateGlobalSymbol(BD->getAddress() - 1, "DATAat");
+  BC.MIB->setOperandToSymbolRef(Inst, OperandNum, NewSym, Offset + 1, &*BC.Ctx,
+                                0);
   LLVM_DEBUG(dbgs() << "BOLT-DEBUG: replaced reference @" << BF.getPrintName()
                     << " from " << BD->getName() << " to " << NewSym->getName()
                     << " + 1\n");
diff --git a/bolt/test/X86/jt-symbol-disambiguation-4.s b/bolt/test/X86/jt-symbol-disambiguation-4.s
new file mode 100644
index 00000000000000..816b0fb501a590
--- /dev/null
+++ b/bolt/test/X86/jt-symbol-disambiguation-4.s
@@ -0,0 +1,60 @@
+# If the operand references a symbol that differs from the jump table label,
+# no reference updating is required even if its target address resides within
+# the jump table's range.
+# In this test case, consider the second instruction within the main function,
+# where the address resulting from 'c + 17' corresponds to one byte beyond the
+# address of the .LJTI2_0 jump table label. However, this operand represents
+# an offset calculation related to the global variable 'c' and should remain
+# unaffected by the jump table.
+
+# REQUIRES: system-linux
+
+
+# RUN: %clang -no-pie  %s -o %t.exe -Wl,-q
+
+# RUN: %t.exe
+# RUN: llvm-bolt -funcs=main,foo/1 %t.exe -o %t.exe.bolt -jump-tables=move
+# RUN: %t.exe.bolt
+
+	.text
+	.globl	main
+	.type	main, at function
+main:
+	pushq   %rbp
+	movq	%rsp, %rbp
+	movq	$-16, %rax
+	movl	c+17(%rax), %edx
+	cmpl	$255, %edx
+	je	.LCorrect
+	movl	$1, %eax
+	popq	%rbp
+	ret
+.LCorrect:
+	movl	$0, %eax
+	popq	%rbp
+	ret
+	.p2align	4, 0x90
+	.type	foo, at function
+foo:
+	movq	$0, %rax
+	jmpq	*.LJTI2_0(,%rax,8)
+	addl	$-36, %eax
+.LBB2_2:
+	addl	$-16, %eax
+	retq
+	.section	.rodata,"a", at progbits
+	.type	c, at object
+	.data
+	.globl	c
+	.p2align	4, 0x0
+c:
+	.byte 1
+  .byte 0xff
+	.zero	14
+	.size	c, 16
+.LJTI2_0:
+	.quad	.LBB2_2
+	.quad	.LBB2_2
+	.quad	.LBB2_2
+	.quad	.LBB2_2
+