From 2c2c4734b0f4ba16fac58e91e635991bb732de59 Mon Sep 17 00:00:00 2001
From: Sergei Barannikov <barannikov88@gmail.com>
Date: Mon, 14 Apr 2025 05:59:57 +0300
Subject: [PATCH 1/2] [AsmPrinter] Link .section_sizes to the correct section

AsmPrinter may switch the current section when e.g., emitting a jump
table for a switch. `.stack_sizes` should still be linked to the
function section. If the section is wrong, objdump emits a warning
"relocation symbol is not in the expected section".
---
 llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp    |  2 +-
 .../CodeGen/SystemZ/stack-size-section.ll     | 30 +++++++++++++++++++
 2 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp b/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
index cf8f1c878ea5a..821879a1f7c36 100644
--- a/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
+++ b/llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp
@@ -1586,7 +1586,7 @@ void AsmPrinter::emitStackSizeSection(const MachineFunction &MF) {
     return;
 
   MCSection *StackSizeSection =
-      getObjFileLowering().getStackSizesSection(*getCurrentSection());
+      getObjFileLowering().getStackSizesSection(*MF.getSection());
   if (!StackSizeSection)
     return;
 
diff --git a/llvm/test/CodeGen/SystemZ/stack-size-section.ll b/llvm/test/CodeGen/SystemZ/stack-size-section.ll
index 024f20bfc2dc8..f2429bc4bfbf8 100644
--- a/llvm/test/CodeGen/SystemZ/stack-size-section.ll
+++ b/llvm/test/CodeGen/SystemZ/stack-size-section.ll
@@ -38,4 +38,34 @@ define void @dynalloc(i32 %N) #0 {
   ret void
 }
 
+; Check that .stack_sizes section is linked to the function's section (.text),
+; and not to the section containing the jump table (.rodata).
+; CHECK-LABEL: .section .stack_sizes,"o",@progbits,.text{{$}}
+; CHECK-NEXT: .quad .Lfunc_begin4
+; CHECK-NEXT: .ascii "\260!"
+define i32 @jump_table(i32 %x) {
+  %arr = alloca [1024 x i32]
+  switch i32 %x, label %sw.epilog [
+    i32 0, label %sw.bb0
+    i32 1, label %sw.bb1
+    i32 2, label %sw.bb2
+    i32 3, label %sw.bb3
+  ]
+
+sw.bb0:
+  ret i32 0
+
+sw.bb1:
+  ret i32 1
+
+sw.bb2:
+  ret i32 2
+
+sw.bb3:
+  ret i32 3
+
+sw.epilog:
+  ret i32 -1
+}
+
 attributes #0 = { "frame-pointer"="all" }

From b612daf3bf01884a6d5d8043b3e632bdf369db53 Mon Sep 17 00:00:00 2001
From: Sergei Barannikov <barannikov88@gmail.com>
Date: Mon, 14 Apr 2025 11:27:00 +0300
Subject: [PATCH 2/2] Adapt the test to AArch64

---
 .../CodeGen/AArch64/stack-size-section.ll     | 63 +++++++++++++++++++
 .../CodeGen/SystemZ/stack-size-section.ll     |  2 +-
 2 files changed, 64 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/CodeGen/AArch64/stack-size-section.ll

diff --git a/llvm/test/CodeGen/AArch64/stack-size-section.ll b/llvm/test/CodeGen/AArch64/stack-size-section.ll
new file mode 100644
index 0000000000000..e7429032da16a
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/stack-size-section.ll
@@ -0,0 +1,63 @@
+; RUN: llc -mtriple=aarch64 -aarch64-min-jump-table-entries=4 -stack-size-section %s -o - | FileCheck %s
+
+; CHECK-LABEL: .section .stack_sizes,"o",@progbits,.text{{$}}
+; CHECK-NEXT: .xword .Lfunc_begin0
+; CHECK-NEXT: .byte 0
+define void @empty() {
+  ret void
+}
+
+; CHECK-LABEL: .section .stack_sizes,"o",@progbits,.text{{$}}
+; CHECK-NEXT: .xword .Lfunc_begin1
+; CHECK-NEXT: .ascii "\200\001"
+define void @non_empty() #0 {
+  alloca [32 x i32]
+  ret void
+}
+
+; CHECK-LABEL: dynalloc:
+; CHECK-NOT: .section .stack_sizes
+define void @dynalloc(i32 %n) #0 {
+  alloca i32, i32 %n
+  ret void
+}
+
+; Check that .stack_sizes section is linked to the function's section (.text),
+; and not to the section containing the jump table (.rodata).
+; CHECK-LABEL: linked_section:
+; CHECK: .section .rodata,"a",@progbits
+; CHECK: .section .stack_sizes,"o",@progbits,.text
+; CHECK-NEXT: .xword .Lfunc_begin3
+; CHECK-NEXT: .ascii "\220\001"
+declare void @case0()
+declare void @case1()
+declare void @case2()
+declare void @case3()
+define void @linked_section(i32 %x) {
+  %arr = alloca [32 x i32]
+  switch i32 %x, label %sw.epilog [
+    i32 0, label %sw.bb0
+    i32 1, label %sw.bb1
+    i32 2, label %sw.bb2
+    i32 3, label %sw.bb3
+  ]
+
+sw.bb0:
+  call void @case0()
+  ret void
+
+sw.bb1:
+  call void @case1()
+  ret void
+
+sw.bb2:
+  call void @case2()
+  ret void
+
+sw.bb3:
+  call void @case3()
+  ret void
+
+sw.epilog:
+  ret void
+}
diff --git a/llvm/test/CodeGen/SystemZ/stack-size-section.ll b/llvm/test/CodeGen/SystemZ/stack-size-section.ll
index f2429bc4bfbf8..d9fb98cbb8944 100644
--- a/llvm/test/CodeGen/SystemZ/stack-size-section.ll
+++ b/llvm/test/CodeGen/SystemZ/stack-size-section.ll
@@ -43,7 +43,7 @@ define void @dynalloc(i32 %N) #0 {
 ; CHECK-LABEL: .section .stack_sizes,"o",@progbits,.text{{$}}
 ; CHECK-NEXT: .quad .Lfunc_begin4
 ; CHECK-NEXT: .ascii "\260!"
-define i32 @jump_table(i32 %x) {
+define i32 @linked_section(i32 %x) {
   %arr = alloca [1024 x i32]
   switch i32 %x, label %sw.epilog [
     i32 0, label %sw.bb0