From c2102259a04dd9ede1e7faf01daa0dfcb948c214 Mon Sep 17 00:00:00 2001
From: Bastian Kersting <bkersting@google.com>
Date: Fri, 8 Nov 2024 08:51:19 +0000
Subject: [PATCH] CFI: Append debug location to CFI blocks

---
 compiler/rustc_codegen_gcc/src/debuginfo.rs   |  4 ++++
 compiler/rustc_codegen_llvm/src/builder.rs    |  7 +++++++
 .../rustc_codegen_llvm/src/debuginfo/mod.rs   |  4 ++++
 compiler/rustc_codegen_llvm/src/llvm/ffi.rs   |  1 +
 .../rustc_codegen_ssa/src/traits/debuginfo.rs |  1 +
 .../cfi/dbg-location-on-cfi-blocks.rs         | 19 +++++++++++++++++++
 6 files changed, 36 insertions(+)
 create mode 100644 tests/codegen/sanitizer/cfi/dbg-location-on-cfi-blocks.rs

diff --git a/compiler/rustc_codegen_gcc/src/debuginfo.rs b/compiler/rustc_codegen_gcc/src/debuginfo.rs
index 9d62ccc95d5..5d8c5c199b1 100644
--- a/compiler/rustc_codegen_gcc/src/debuginfo.rs
+++ b/compiler/rustc_codegen_gcc/src/debuginfo.rs
@@ -52,6 +52,10 @@ impl<'a, 'gcc, 'tcx> DebugInfoBuilderMethods for Builder<'a, 'gcc, 'tcx> {
     fn clear_dbg_loc(&mut self) {
         self.location = None;
     }
+
+    fn get_dbg_loc(&self) -> Option<Self::DILocation> {
+        self.location
+    }
 }
 
 /// Generate the `debug_context` in an MIR Body.
diff --git a/compiler/rustc_codegen_llvm/src/builder.rs b/compiler/rustc_codegen_llvm/src/builder.rs
index 751b2235dc8..ac76b781218 100644
--- a/compiler/rustc_codegen_llvm/src/builder.rs
+++ b/compiler/rustc_codegen_llvm/src/builder.rs
@@ -1574,6 +1574,7 @@ impl<'a, 'll, 'tcx> Builder<'a, 'll, 'tcx> {
                 cfi::typeid_for_fnabi(self.tcx, fn_abi, options)
             };
             let typeid_metadata = self.cx.typeid_metadata(typeid).unwrap();
+            let dbg_loc = self.get_dbg_loc();
 
             // Test whether the function pointer is associated with the type identifier.
             let cond = self.type_test(llfn, typeid_metadata);
@@ -1582,10 +1583,16 @@ impl<'a, 'll, 'tcx> Builder<'a, 'll, 'tcx> {
             self.cond_br(cond, bb_pass, bb_fail);
 
             self.switch_to_block(bb_fail);
+            if let Some(dbg_loc) = dbg_loc {
+                self.set_dbg_loc(dbg_loc);
+            }
             self.abort();
             self.unreachable();
 
             self.switch_to_block(bb_pass);
+            if let Some(dbg_loc) = dbg_loc {
+                self.set_dbg_loc(dbg_loc);
+            }
         }
     }
 
diff --git a/compiler/rustc_codegen_llvm/src/debuginfo/mod.rs b/compiler/rustc_codegen_llvm/src/debuginfo/mod.rs
index 9e1e5127e80..89492e4b9fe 100644
--- a/compiler/rustc_codegen_llvm/src/debuginfo/mod.rs
+++ b/compiler/rustc_codegen_llvm/src/debuginfo/mod.rs
@@ -206,6 +206,10 @@ impl<'ll> DebugInfoBuilderMethods for Builder<'_, 'll, '_> {
         }
     }
 
+    fn get_dbg_loc(&self) -> Option<&'ll DILocation> {
+        unsafe { llvm::LLVMGetCurrentDebugLocation2(self.llbuilder) }
+    }
+
     fn insert_reference_to_gdb_debug_scripts_section_global(&mut self) {
         gdb::insert_reference_to_gdb_debug_scripts_section_global(self)
     }
diff --git a/compiler/rustc_codegen_llvm/src/llvm/ffi.rs b/compiler/rustc_codegen_llvm/src/llvm/ffi.rs
index d84ae8d8836..9ea1f469813 100644
--- a/compiler/rustc_codegen_llvm/src/llvm/ffi.rs
+++ b/compiler/rustc_codegen_llvm/src/llvm/ffi.rs
@@ -1063,6 +1063,7 @@ unsafe extern "C" {
 
     // Metadata
     pub fn LLVMSetCurrentDebugLocation2<'a>(Builder: &Builder<'a>, Loc: *const Metadata);
+    pub fn LLVMGetCurrentDebugLocation2<'a>(Builder: &Builder<'a>) -> Option<&'a Metadata>;
 
     // Terminators
     pub fn LLVMBuildRetVoid<'a>(B: &Builder<'a>) -> &'a Value;
diff --git a/compiler/rustc_codegen_ssa/src/traits/debuginfo.rs b/compiler/rustc_codegen_ssa/src/traits/debuginfo.rs
index 670433a6c33..fe135e911fb 100644
--- a/compiler/rustc_codegen_ssa/src/traits/debuginfo.rs
+++ b/compiler/rustc_codegen_ssa/src/traits/debuginfo.rs
@@ -81,6 +81,7 @@ pub trait DebugInfoBuilderMethods: BackendTypes {
     );
     fn set_dbg_loc(&mut self, dbg_loc: Self::DILocation);
     fn clear_dbg_loc(&mut self);
+    fn get_dbg_loc(&self) -> Option<Self::DILocation>;
     fn insert_reference_to_gdb_debug_scripts_section_global(&mut self);
     fn set_var_name(&mut self, value: Self::Value, name: &str);
 }
diff --git a/tests/codegen/sanitizer/cfi/dbg-location-on-cfi-blocks.rs b/tests/codegen/sanitizer/cfi/dbg-location-on-cfi-blocks.rs
new file mode 100644
index 00000000000..df65960dfe0
--- /dev/null
+++ b/tests/codegen/sanitizer/cfi/dbg-location-on-cfi-blocks.rs
@@ -0,0 +1,19 @@
+// Verifies that the parent block's debug information are assigned to the inserted cfi block.
+//
+//@ needs-sanitizer-cfi
+//@ compile-flags: -Clto -Cno-prepopulate-passes -Copt-level=0 -Zsanitizer=cfi -Ctarget-feature=-crt-static -Cdebuginfo=1
+
+#![crate_type = "lib"]
+
+pub fn foo(f: fn(i32) -> i32, arg: i32) -> i32 {
+    // CHECK-LABEL: define{{.*}}foo{{.*}}!dbg !{{[0-9]+}} !type !{{[0-9]+}} !type !{{[0-9]+}} !type !{{[0-9]+}} !type !{{[0-9]+}}
+    // CHECK:       start:
+    // CHECK:       [[TT:%.+]] = call i1 @llvm.type.test(ptr {{%f|%0}}, metadata !"{{[[:print:]]+}}"), !dbg !{{[0-9]+}}
+    // CHECK-NEXT:  br i1 [[TT]], label %type_test.pass, label %type_test.fail, !dbg !{{[0-9]+}}
+    // CHECK:       type_test.pass:                                   ; preds = %start
+    // CHECK-NEXT:  {{%.+}} = call i32 %f(i32{{.*}} %arg), !dbg !{{[0-9]+}}
+    // CHECK:       type_test.fail:                                   ; preds = %start
+    // CHECK-NEXT:  call void @llvm.trap(), !dbg !{{[0-9]+}}
+    // CHECK-NEXT:  unreachable, !dbg !{{[0-9]+}}
+    f(arg)
+}