[WIP] TypeTree support in autodiff

compiler/rustc_ast/src/expand/autodiff_attrs.rs

@@ -9,6 +9,7 @@ use std::str::FromStr;
 use crate::expand::{Decodable, Encodable, HashStable_Generic};
 use crate::ptr::P;
 use crate::{Ty, TyKind};
+use crate::expand::typetree::TypeTree;
 
 /// Forward and Reverse Mode are well known names for automatic differentiation implementations.
 /// Enzyme does support both, but with different semantics, see DiffActivity. The First variants
@@ -85,6 +86,9 @@ pub struct AutoDiffItem {
     /// The name of the function being generated
     pub target: String,
     pub attrs: AutoDiffAttrs,
+    // --- TypeTree support ---
+    pub inputs: Vec<TypeTree>,
+    pub output: TypeTree,
 }
 
 #[derive(Clone, Eq, PartialEq, Encodable, Decodable, Debug, HashStable_Generic)]
@@ -112,6 +116,10 @@ impl AutoDiffAttrs {
     pub fn has_primal_ret(&self) -> bool {
         matches!(self.ret_activity, DiffActivity::Active | DiffActivity::Dual)
     }
+    /// New constructor for type tree support
+    pub fn into_item(self, source: String, target: String, inputs: Vec<TypeTree>, output: TypeTree) -> AutoDiffItem {
+        AutoDiffItem { source, target, attrs: self, inputs, output }
+    }
 }
 
 impl DiffMode {
@@ -284,6 +292,8 @@ impl AutoDiffAttrs {
 impl fmt::Display for AutoDiffItem {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         write!(f, "Differentiating {} -> {}", self.source, self.target)?;
-        write!(f, " with attributes: {:?}", self.attrs)
+        write!(f, " with attributes: {:?}", self.attrs)?;
+        write!(f, " with inputs: {:?}", self.inputs)?;
+        write!(f, " with output: {:?}", self.output)
     }
 }

compiler/rustc_builtin_macros/src/autodiff.rs

@@ -11,7 +11,9 @@ mod llvm_enzyme {
         AutoDiffAttrs, DiffActivity, DiffMode, valid_input_activity, valid_ret_activity,
         valid_ty_for_activity,
     };
+    use rustc_ast::expand::typetree::{TypeTree, Type, Kind};
     use rustc_ast::ptr::P;
+    use crate::typetree::construct_typetree_from_fnsig;
     use rustc_ast::token::{Lit, LitKind, Token, TokenKind};
     use rustc_ast::tokenstream::*;
     use rustc_ast::visit::AssocCtxt::*;
@@ -324,6 +326,17 @@ mod llvm_enzyme {
         }
         let span = ecx.with_def_site_ctxt(expand_span);
 
+        // Construct real type trees from function signature
+        let (inputs, output) = construct_typetree_from_fnsig(&sig);
+
+        // Use the new into_item method to construct the AutoDiffItem
+        let autodiff_item = x.clone().into_item(
+            primal.to_string(),
+            first_ident(&meta_item_vec[0]).to_string(),
+            inputs,
+            output,
+        );
+
         let n_active: u32 = x
             .input_activity
             .iter()
@@ -1045,5 +1058,3 @@ mod llvm_enzyme {
         (d_sig, new_inputs, idents, false)
     }
 }
-
-pub(crate) use llvm_enzyme::{expand_forward, expand_reverse};

compiler/rustc_builtin_macros/src/lib.rs

@@ -51,6 +51,7 @@ mod pattern_type;
 mod source_util;
 mod test;
 mod trace_macros;
+mod typetree;
 
 pub mod asm;
 pub mod cmdline_attrs;