Mark more LLVM FFI as safe

oli-obk · oli-obk · commit 553828c6f48a · 2025-02-24T15:11:29.000Z
diff --git a/compiler/rustc_codegen_llvm/src/asm.rs b/compiler/rustc_codegen_llvm/src/asm.rs
@@ -482,12 +482,13 @@ pub(crate) fn inline_asm_call<'ll>(
 
     debug!("Asm Output Type: {:?}", output);
     let fty = bx.cx.type_func(&argtys, output);
-    unsafe {
-        // Ask LLVM to verify that the constraints are well-formed.
-        let constraints_ok = llvm::LLVMRustInlineAsmVerify(fty, cons.as_c_char_ptr(), cons.len());
-        debug!("constraint verification result: {:?}", constraints_ok);
-        if constraints_ok {
-            let v = llvm::LLVMRustInlineAsm(
+    // Ask LLVM to verify that the constraints are well-formed.
+    let constraints_ok =
+        unsafe { llvm::LLVMRustInlineAsmVerify(fty, cons.as_c_char_ptr(), cons.len()) };
+    debug!("constraint verification result: {:?}", constraints_ok);
+    if constraints_ok {
+        let v = unsafe {
+            llvm::LLVMRustInlineAsm(
                 fty,
                 asm.as_c_char_ptr(),
                 asm.len(),
@@ -497,50 +498,50 @@ pub(crate) fn inline_asm_call<'ll>(
                 alignstack,
                 dia,
                 can_throw,
-            );
-
-            let call = if !labels.is_empty() {
-                assert!(catch_funclet.is_none());
-                bx.callbr(fty, None, None, v, inputs, dest.unwrap(), labels, None, None)
-            } else if let Some((catch, funclet)) = catch_funclet {
-                bx.invoke(fty, None, None, v, inputs, dest.unwrap(), catch, funclet, None)
-            } else {
-                bx.call(fty, None, None, v, inputs, None, None)
-            };
+            )
+        };
 
-            // Store mark in a metadata node so we can map LLVM errors
-            // back to source locations. See #17552.
-            let key = "srcloc";
-            let kind = bx.get_md_kind_id(key);
+        let call = if !labels.is_empty() {
+            assert!(catch_funclet.is_none());
+            bx.callbr(fty, None, None, v, inputs, dest.unwrap(), labels, None, None)
+        } else if let Some((catch, funclet)) = catch_funclet {
+            bx.invoke(fty, None, None, v, inputs, dest.unwrap(), catch, funclet, None)
+        } else {
+            bx.call(fty, None, None, v, inputs, None, None)
+        };
 
-            // `srcloc` contains one 64-bit integer for each line of assembly code,
-            // where the lower 32 bits hold the lo byte position and the upper 32 bits
-            // hold the hi byte position.
-            let mut srcloc = vec![];
-            if dia == llvm::AsmDialect::Intel && line_spans.len() > 1 {
-                // LLVM inserts an extra line to add the ".intel_syntax", so add
-                // a dummy srcloc entry for it.
-                //
-                // Don't do this if we only have 1 line span since that may be
-                // due to the asm template string coming from a macro. LLVM will
-                // default to the first srcloc for lines that don't have an
-                // associated srcloc.
-                srcloc.push(llvm::LLVMValueAsMetadata(bx.const_u64(0)));
-            }
-            srcloc.extend(line_spans.iter().map(|span| {
-                llvm::LLVMValueAsMetadata(bx.const_u64(
-                    u64::from(span.lo().to_u32()) | (u64::from(span.hi().to_u32()) << 32),
-                ))
-            }));
-            let md = llvm::LLVMMDNodeInContext2(bx.llcx, srcloc.as_ptr(), srcloc.len());
-            let md = bx.get_metadata_value(md);
-            llvm::LLVMSetMetadata(call, kind, md);
+        // Store mark in a metadata node so we can map LLVM errors
+        // back to source locations. See #17552.
+        let key = "srcloc";
+        let kind = bx.get_md_kind_id(key);
 
-            Some(call)
-        } else {
-            // LLVM has detected an issue with our constraints, bail out
-            None
+        // `srcloc` contains one 64-bit integer for each line of assembly code,
+        // where the lower 32 bits hold the lo byte position and the upper 32 bits
+        // hold the hi byte position.
+        let mut srcloc = vec![];
+        if dia == llvm::AsmDialect::Intel && line_spans.len() > 1 {
+            // LLVM inserts an extra line to add the ".intel_syntax", so add
+            // a dummy srcloc entry for it.
+            //
+            // Don't do this if we only have 1 line span since that may be
+            // due to the asm template string coming from a macro. LLVM will
+            // default to the first srcloc for lines that don't have an
+            // associated srcloc.
+            srcloc.push(llvm::LLVMValueAsMetadata(bx.const_u64(0)));
         }
+        srcloc.extend(line_spans.iter().map(|span| {
+            llvm::LLVMValueAsMetadata(
+                bx.const_u64(u64::from(span.lo().to_u32()) | (u64::from(span.hi().to_u32()) << 32)),
+            )
+        }));
+        let md = unsafe { llvm::LLVMMDNodeInContext2(bx.llcx, srcloc.as_ptr(), srcloc.len()) };
+        let md = bx.get_metadata_value(md);
+        llvm::LLVMSetMetadata(call, kind, md);
+
+        Some(call)
+    } else {
+        // LLVM has detected an issue with our constraints, bail out
+        None
     }
 }
 
diff --git a/compiler/rustc_codegen_llvm/src/builder.rs b/compiler/rustc_codegen_llvm/src/builder.rs
@@ -311,8 +311,8 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
         // This function handles switch instructions with more than 2 targets and it needs to
         // emit branch weights metadata instead of using the intrinsic.
         // The values 1 and 2000 are the same as the values used by the `llvm.expect` intrinsic.
-        let cold_weight = unsafe { llvm::LLVMValueAsMetadata(self.cx.const_u32(1)) };
-        let hot_weight = unsafe { llvm::LLVMValueAsMetadata(self.cx.const_u32(2000)) };
+        let cold_weight = llvm::LLVMValueAsMetadata(self.cx.const_u32(1));
+        let hot_weight = llvm::LLVMValueAsMetadata(self.cx.const_u32(2000));
         let weight =
             |is_cold: bool| -> &Metadata { if is_cold { cold_weight } else { hot_weight } };
 
diff --git a/compiler/rustc_codegen_llvm/src/context.rs b/compiler/rustc_codegen_llvm/src/context.rs
@@ -34,7 +34,7 @@ use crate::back::write::to_llvm_code_model;
 use crate::callee::get_fn;
 use crate::common::AsCCharPtr;
 use crate::debuginfo::metadata::apply_vcall_visibility_metadata;
-use crate::llvm::{Metadata, MetadataType};
+use crate::llvm::Metadata;
 use crate::type_::Type;
 use crate::value::Value;
 use crate::{attributes, coverageinfo, debuginfo, llvm, llvm_util};
@@ -664,7 +664,7 @@ impl<'ll, CX: Borrow<SCx<'ll>>> GenericCx<'ll, CX> {
         unsafe { llvm::LLVMGetNamedFunction((**self).borrow().llmod, name.as_ptr()) }
     }
 
-    pub(crate) fn get_md_kind_id(&self, name: &str) -> u32 {
+    pub(crate) fn get_md_kind_id(&self, name: &str) -> llvm::MetadataKindId {
         unsafe {
             llvm::LLVMGetMDKindIDInContext(
                 self.llcx(),
@@ -1228,13 +1228,11 @@ impl<'ll, CX: Borrow<SCx<'ll>>> GenericCx<'ll, CX> {
     pub(crate) fn set_metadata<'a>(
         &self,
         val: &'a Value,
-        kind_id: MetadataType,
+        kind_id: impl Into<llvm::MetadataKindId>,
         md: &'ll Metadata,
     ) {
         let node = self.get_metadata_value(md);
-        unsafe {
-            llvm::LLVMSetMetadata(val, kind_id as c_uint, node);
-        }
+        llvm::LLVMSetMetadata(val, kind_id.into(), node);
     }
 }
 
diff --git a/compiler/rustc_codegen_llvm/src/llvm/enzyme_ffi.rs b/compiler/rustc_codegen_llvm/src/llvm/enzyme_ffi.rs
@@ -3,13 +3,14 @@
 
 use libc::{c_char, c_uint};
 
+use super::MetadataKindId;
 use super::ffi::{BasicBlock, Metadata, Module, Type, Value};
 use crate::llvm::Bool;
 
 #[link(name = "llvm-wrapper", kind = "static")]
 unsafe extern "C" {
     // Enzyme
-    pub(crate) fn LLVMRustHasMetadata(I: &Value, KindID: c_uint) -> bool;
+    pub(crate) safe fn LLVMRustHasMetadata(I: &Value, KindID: MetadataKindId) -> bool;
     pub(crate) fn LLVMRustEraseInstUntilInclusive(BB: &BasicBlock, I: &Value);
     pub(crate) fn LLVMRustGetLastInstruction<'a>(BB: &BasicBlock) -> Option<&'a Value>;
     pub(crate) fn LLVMRustDIGetInstMetadata(I: &Value) -> Option<&Metadata>;
diff --git a/compiler/rustc_codegen_llvm/src/llvm/ffi.rs b/compiler/rustc_codegen_llvm/src/llvm/ffi.rs
@@ -976,14 +976,24 @@ pub type SelfProfileAfterPassCallback = unsafe extern "C" fn(*mut c_void);
 pub type GetSymbolsCallback = unsafe extern "C" fn(*mut c_void, *const c_char) -> *mut c_void;
 pub type GetSymbolsErrorCallback = unsafe extern "C" fn(*const c_char) -> *mut c_void;
 
+#[derive(Copy, Clone)]
+#[repr(transparent)]
+pub struct MetadataKindId(c_uint);
+
+impl From<MetadataType> for MetadataKindId {
+    fn from(value: MetadataType) -> Self {
+        Self(value as c_uint)
+    }
+}
+
 unsafe extern "C" {
     // Create and destroy contexts.
     pub(crate) fn LLVMContextDispose(C: &'static mut Context);
     pub(crate) fn LLVMGetMDKindIDInContext(
         C: &Context,
         Name: *const c_char,
         SLen: c_uint,
-    ) -> c_uint;
+    ) -> MetadataKindId;
 
     // Create modules.
     pub(crate) fn LLVMModuleCreateWithNameInContext(
@@ -1051,9 +1061,9 @@ unsafe extern "C" {
     pub(crate) fn LLVMGetValueName2(Val: &Value, Length: *mut size_t) -> *const c_char;
     pub(crate) fn LLVMSetValueName2(Val: &Value, Name: *const c_char, NameLen: size_t);
     pub(crate) fn LLVMReplaceAllUsesWith<'a>(OldVal: &'a Value, NewVal: &'a Value);
-    pub(crate) fn LLVMSetMetadata<'a>(Val: &'a Value, KindID: c_uint, Node: &'a Value);
+    pub(crate) safe fn LLVMSetMetadata<'a>(Val: &'a Value, KindID: MetadataKindId, Node: &'a Value);
     pub(crate) fn LLVMGlobalSetMetadata<'a>(Val: &'a Value, KindID: c_uint, Metadata: &'a Metadata);
-    pub(crate) fn LLVMValueAsMetadata(Node: &Value) -> &Metadata;
+    pub(crate) safe fn LLVMValueAsMetadata(Node: &Value) -> &Metadata;
 
     // Operations on constants of any type
     pub(crate) fn LLVMConstNull(Ty: &Type) -> &Value;
