[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 01deba3048bcabf4ea2bfd110afa693a84f9a53e 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        |  5 ++
 bolt/test/X86/jt-symbol-disambiguation-4.s | 63 ++++++++++++++++++++++
 2 files changed, 68 insertions(+)
 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..7e599fc11e600a 100644
--- a/bolt/lib/Passes/ValidateMemRefs.cpp
+++ b/bolt/lib/Passes/ValidateMemRefs.cpp
@@ -34,6 +34,11 @@ bool ValidateMemRefs::checkAndFixJTReference(BinaryFunction &BF, MCInst &Inst,
   if (!JT)
     return false;
 
+  // If the operand does not refer to the jump table symbol, then we
+  // don't need to update the reference.
+  if (JT->getFirstLabel() != Sym)
+    return true;
+
   const bool IsLegitAccess = llvm::is_contained(JT->Parents, &BF);
   if (IsLegitAccess)
     return true;
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..ee87a0231800eb
--- /dev/null
+++ b/bolt/test/X86/jt-symbol-disambiguation-4.s
@@ -0,0 +1,63 @@
+# 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: llvm-mc -filetype=obj -triple x86_64-unknown-unknown %s -o %t.o
+# RUN: %clang -no-pie %t.o -o %t.exe -Wl,-q
+
+# RUN: %t.exe | FileCheck %s -check-prefix=CHECK
+# RUN: llvm-bolt -funcs=main,foo/1 %t.exe -o %t.exe.bolt -jump-tables=move
+# RUN: %t.exe.bolt | FileCheck %s -check-prefix=CHECK-AFTERBOLT
+
+# CHECK: {{^}}FF{{$}}
+# CHECK-AFTERBOLT: {{^}}FF{{$}}
+	.text
+	.globl	main
+	.p2align	4, 0x90
+	.type	main, at function
+main:
+	movq	$-16, %rax
+	movl	c+17(%rax), %edx
+	movl	%edx, %esi
+	movl	$.L.str, %edi
+	movl	$0, %eax
+	callq	printf
+	xorl	%eax, %eax
+	retq
+	.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
+	.p2align	3, 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
+	.type	c, at object
+	.data
+	.globl	c
+	.p2align	4, 0x0
+	.type	.L.str, at object
+	.section	.rodata.str1.1,"aMS", at progbits,1
+.L.str:
+	.asciz	"%X\n"
+	.size	.L.str, 4