fixing ffi issue

Author: KMJ-007

compiler/rustc_codegen_llvm/src/builder/autodiff.rs

@@ -16,7 +16,7 @@ use crate::context::SimpleCx;
 use crate::declare::declare_simple_fn;
 use crate::errors::{AutoDiffWithoutEnable, LlvmError};
 use crate::llvm::AttributePlace::Function;
-use crate::llvm::{Metadata, True};
+use crate::llvm::{Metadata, True, TypeTree};
 use crate::value::Value;
 use crate::{CodegenContext, LlvmCodegenBackend, ModuleLlvm, attributes, llvm};
 
@@ -520,13 +520,14 @@ pub(crate) fn differentiate<'ll>(
 /// This function takes a Rust-side TypeTree (from rustc_ast::expand::typetree)
 /// and converts it to Enzyme's internal C++ TypeTree representation that
 /// Enzyme can understand during differentiation analysis.
+#[cfg(llvm_enzyme)]
 fn to_enzyme_typetree(
     rust_typetree: RustTypeTree,
     data_layout: &str,
     llcx: &llvm::Context,
-) -> llvm::TypeTree {
+) -> TypeTree {
     // Start with an empty TypeTree
-    let mut enzyme_tt = llvm::TypeTree::new();
+    let mut enzyme_tt = TypeTree::new();
 
     // Convert each Type in the Rust TypeTree to Enzyme format
     for rust_type in rust_typetree.0 {
@@ -541,7 +542,7 @@ fn to_enzyme_typetree(
         };
 
         // Create a TypeTree for this specific type
-        let type_tt = llvm::TypeTree::from_type(concrete_type, llcx);
+        let type_tt = TypeTree::from_type(concrete_type, llcx);
 
         // Apply offset if specified
         let type_tt = if rust_type.offset == -1 {
@@ -558,7 +559,17 @@ fn to_enzyme_typetree(
     enzyme_tt
 }
 
+#[cfg(not(llvm_enzyme))]
+fn to_enzyme_typetree(
+    _rust_typetree: RustTypeTree,
+    _data_layout: &str,
+    _llcx: &llvm::Context,
+) -> ! {
+    unimplemented!("TypeTree conversion not available without llvm_enzyme support")
+}
+
 // Attaches TypeTree information to LLVM function as enzyme_type attributes.
+#[cfg(llvm_enzyme)]
 pub(crate) fn add_tt<'ll>(
     llmod: &'ll llvm::Module,
     llcx: &'ll llvm::Context,

compiler/rustc_codegen_llvm/src/llvm/enzyme_ffi.rs

@@ -1,6 +1,6 @@
 #![expect(dead_code)]
 
-use libc::{c_char, c_uint, size_t};
+use libc::{c_char, c_uint};
 
 use super::MetadataKindId;
 use super::ffi::{AttributeKind, BasicBlock, Context, Metadata, Module, Type, Value};
@@ -17,6 +17,7 @@ pub(crate) struct EnzymeTypeTree {
 
 #[repr(u32)]
 #[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)]
+#[allow(non_camel_case_types)]
 pub(crate) enum CConcreteType {
     DT_Anything = 0,
     DT_Integer = 1,
@@ -27,6 +28,10 @@ pub(crate) enum CConcreteType {
     DT_Unknown = 6,
 }
 
+pub(crate) struct TypeTree {
+    pub(crate) inner: CTypeTreeRef,
+}
+
 #[link(name = "llvm-wrapper", kind = "static")]
 unsafe extern "C" {
     // Enzyme
@@ -87,7 +92,7 @@ pub(crate) use self::Enzyme_AD::*;
 pub(crate) mod Enzyme_AD {
     use std::ffi::{CString, c_char};
 
-    use libc::{c_void, size_t};
+    use libc::c_void;
 
     use super::{CConcreteType, CTypeTreeRef, Context};
 
@@ -264,3 +269,81 @@ pub(crate) mod Fallback_AD {
         unimplemented!()
     }
 }
+
+impl TypeTree {
+    pub(crate) fn new() -> TypeTree {
+        let inner = unsafe { EnzymeNewTypeTree() };
+        TypeTree { inner }
+    }
+
+    pub(crate) fn from_type(t: CConcreteType, ctx: &Context) -> TypeTree {
+        let inner = unsafe { EnzymeNewTypeTreeCT(t, ctx) };
+        TypeTree { inner }
+    }
+
+    pub(crate) fn merge(self, other: Self) -> Self {
+        unsafe {
+            EnzymeMergeTypeTree(self.inner, other.inner);
+        }
+        drop(other);
+        self
+    }
+
+    #[must_use]
+    pub(crate) fn shift(
+        self,
+        layout: &str,
+        offset: isize,
+        max_size: isize,
+        add_offset: usize,
+    ) -> Self {
+        let layout = std::ffi::CString::new(layout).unwrap();
+
+        unsafe {
+            EnzymeTypeTreeShiftIndiciesEq(
+                self.inner,
+                layout.as_ptr(),
+                offset as i64,
+                max_size as i64,
+                add_offset as u64,
+            )
+        }
+
+        self
+    }
+}
+
+impl Clone for TypeTree {
+    fn clone(&self) -> Self {
+        let inner = unsafe { EnzymeNewTypeTreeTR(self.inner) };
+        TypeTree { inner }
+    }
+}
+
+impl std::fmt::Display for TypeTree {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let ptr = unsafe { EnzymeTypeTreeToString(self.inner) };
+        let cstr = unsafe { std::ffi::CStr::from_ptr(ptr) };
+        match cstr.to_str() {
+            Ok(x) => write!(f, "{}", x)?,
+            Err(err) => write!(f, "could not parse: {}", err)?,
+        }
+
+        // delete C string pointer
+        unsafe { EnzymeTypeTreeToStringFree(ptr) }
+
+        Ok(())
+    }
+}
+
+impl std::fmt::Debug for TypeTree {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        <Self as std::fmt::Display>::fmt(self, f)
+    }
+}
+
+impl Drop for TypeTree {
+    fn drop(&mut self) {
+        unsafe { EnzymeFreeTypeTree(self.inner) }
+    }
+}

compiler/rustc_codegen_llvm/src/llvm/ffi.rs

@@ -2678,145 +2678,4 @@ unsafe extern "C" {
     pub(crate) fn LLVMRustSetNoSanitizeAddress(Global: &Value);
     pub(crate) fn LLVMRustSetNoSanitizeHWAddress(Global: &Value);
 
-    // ========== ENZYME AUTODIFF FFI FUNCTIONS ==========
-
-    #[cfg(llvm_enzyme)]
-    // Enzyme Type Tree Functions (minimal set for TypeTree support)
-    pub(crate) fn EnzymeNewTypeTree() -> CTypeTreeRef;
-    #[cfg(llvm_enzyme)]
-    pub(crate) fn EnzymeFreeTypeTree(CTT: CTypeTreeRef);
-    #[cfg(llvm_enzyme)]
-    pub(crate) fn EnzymeNewTypeTreeCT(arg1: CConcreteType, ctx: &Context) -> CTypeTreeRef;
-    #[cfg(llvm_enzyme)]
-    pub(crate) fn EnzymeNewTypeTreeTR(arg1: CTypeTreeRef) -> CTypeTreeRef;
-    #[cfg(llvm_enzyme)]
-    pub(crate) fn EnzymeMergeTypeTree(arg1: CTypeTreeRef, arg2: CTypeTreeRef);
-    #[cfg(llvm_enzyme)]
-    pub(crate) fn EnzymeTypeTreeShiftIndiciesEq(
-        arg1: CTypeTreeRef,
-        data_layout: *const c_char,
-        offset: i64,
-        max_size: i64,
-        add_offset: u64,
-    );
-    #[cfg(llvm_enzyme)]
-    pub(crate) fn EnzymeTypeTreeToString(arg1: CTypeTreeRef) -> *const c_char;
-    #[cfg(llvm_enzyme)]
-    pub(crate) fn EnzymeTypeTreeToStringFree(arg1: *const c_char);
-}
-
-// ========== ENZYME TYPES AND ENUMS ==========
-
-// Type Tree Support for Autodiff
-
-#[cfg(llvm_enzyme)]
-#[repr(u32)]
-#[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)]
-pub(crate) enum CConcreteType {
-    DT_Anything = 0,
-    DT_Integer = 1,
-    DT_Pointer = 2,
-    DT_Half = 3,
-    DT_Float = 4,
-    DT_Double = 5,
-    DT_Unknown = 6,
-}
-
-#[cfg(llvm_enzyme)]
-pub(crate) type CTypeTreeRef = *mut EnzymeTypeTree;
-
-#[cfg(llvm_enzyme)]
-#[repr(C)]
-#[derive(Debug, Copy, Clone)]
-pub(crate) struct EnzymeTypeTree {
-    _unused: [u8; 0],
-}
-
-// TypeTree wrapper for Rust-side type safety and memory management
-#[cfg(llvm_enzyme)]
-pub(crate) struct TypeTree {
-    pub(crate) inner: CTypeTreeRef,
-}
-
-#[cfg(llvm_enzyme)]
-impl TypeTree {
-    pub(crate) fn new() -> TypeTree {
-        let inner = unsafe { EnzymeNewTypeTree() };
-        TypeTree { inner }
-    }
-
-    pub(crate) fn from_type(t: CConcreteType, ctx: &Context) -> TypeTree {
-        let inner = unsafe { EnzymeNewTypeTreeCT(t, ctx) };
-        TypeTree { inner }
-    }
-
-    pub(crate) fn merge(self, other: Self) -> Self {
-        unsafe {
-            EnzymeMergeTypeTree(self.inner, other.inner);
-        }
-        drop(other);
-        self
-    }
-
-    #[must_use]
-    pub(crate) fn shift(
-        self,
-        layout: &str,
-        offset: isize,
-        max_size: isize,
-        add_offset: usize,
-    ) -> Self {
-        let layout = std::ffi::CString::new(layout).unwrap();
-
-        unsafe {
-            EnzymeTypeTreeShiftIndiciesEq(
-                self.inner,
-                layout.as_ptr(),
-                offset as i64,
-                max_size as i64,
-                add_offset as u64,
-            )
-        }
-
-        self
-    }
-}
-
-#[cfg(llvm_enzyme)]
-impl Clone for TypeTree {
-    fn clone(&self) -> Self {
-        let inner = unsafe { EnzymeNewTypeTreeTR(self.inner) };
-        TypeTree { inner }
-    }
-}
-
-#[cfg(llvm_enzyme)]
-impl std::fmt::Display for TypeTree {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        let ptr = unsafe { EnzymeTypeTreeToString(self.inner) };
-        let cstr = unsafe { std::ffi::CStr::from_ptr(ptr) };
-        match cstr.to_str() {
-            Ok(x) => write!(f, "{}", x)?,
-            Err(err) => write!(f, "could not parse: {}", err)?,
-        }
-
-        // delete C string pointer
-        unsafe { EnzymeTypeTreeToStringFree(ptr) }
-
-        Ok(())
-    }
-}
-
-#[cfg(llvm_enzyme)]
-impl std::fmt::Debug for TypeTree {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        <Self as std::fmt::Display>::fmt(self, f)
-    }
-}
-
-#[cfg(llvm_enzyme)]
-impl Drop for TypeTree {
-    fn drop(&mut self) {
-        unsafe { EnzymeFreeTypeTree(self.inner) }
-    }
 }