Rollup merge of rust-lang#158278 - ZuseZ4:typetrees-for-enzyme2, r=oli-obk

autodiff - typtree cleanups

In [some recent experiments](https://github.com/rust-lang/rust/pull/157994/changes) I managed to lower more MIR information to LLVM to help Enzyme compile a test case that was previously failing.

That PR became a little bigger, and we might want to add even more code. To get started I split it out of mod.rs, since that file is already 2.4k loc. In the process, I also fix some checks and deleted `typetree_from_ty_impl` and `typetree_from_ty_inner`, since they don't do much.

r? oli-obk

(noticed a regression in a test, marking as draft for now.)

Author: JonathanBrouwer

compiler/rustc_codegen_llvm/src/intrinsic.rs

@@ -43,6 +43,7 @@ use crate::errors::{
     AutoDiffWithoutEnable, AutoDiffWithoutLto, IntrinsicSignatureMismatch, IntrinsicWrongArch,
     OffloadWithoutEnable, OffloadWithoutFatLTO, UnknownIntrinsic,
 };
+use crate::intrinsic::ty::typetree::fnc_typetrees;
 use crate::llvm::{self, Type, Value};
 use crate::type_of::LayoutLlvmExt;
 use crate::va_arg::emit_va_arg;
@@ -1787,7 +1788,7 @@ fn codegen_autodiff<'ll, 'tcx>(
         &mut diff_attrs.input_activity,
     );
 
-    let fnc_tree = rustc_middle::ty::fnc_typetrees(tcx, source_fn_ptr_ty);
+    let fnc_tree = fnc_typetrees(tcx, source_fn_ptr_ty);
 
     // Build body
     generate_enzyme_call(

compiler/rustc_middle/src/ty/mod.rs

@@ -27,7 +27,6 @@ pub use intrinsic::IntrinsicDef;
 use rustc_abi::{
     Align, FieldIdx, Integer, IntegerType, ReprFlags, ReprOptions, ScalableElt, VariantIdx,
 };
-use rustc_ast::expand::typetree::{FncTree, Kind, Type, TypeTree};
 use rustc_ast::node_id::NodeMap;
 use rustc_ast::{self as ast};
 pub use rustc_ast_ir::{Movability, Mutability, try_visit};
@@ -67,7 +66,7 @@ pub use rustc_type_ir::solve::{CandidatePreferenceMode, SizedTraitKind, VisibleF
 pub use rustc_type_ir::*;
 #[allow(hidden_glob_reexports, unused_imports)]
 use rustc_type_ir::{InferCtxtLike, Interner};
-use tracing::{debug, instrument, trace};
+use tracing::{debug, instrument};
 pub use vtable::*;
 
 pub use self::closure::{
@@ -145,6 +144,7 @@ pub mod print;
 pub mod relate;
 pub mod significant_drop_order;
 pub mod trait_def;
+pub mod typetree;
 pub mod util;
 pub mod vtable;
 
@@ -2378,228 +2378,3 @@ pub struct DestructuredAdtConst<'tcx> {
     pub variant: VariantIdx,
     pub fields: &'tcx [ty::Const<'tcx>],
 }
-
-/// Generate TypeTree information for autodiff.
-/// This function creates TypeTree metadata that describes the memory layout
-/// of function parameters and return types for Enzyme autodiff.
-pub fn fnc_typetrees<'tcx>(tcx: TyCtxt<'tcx>, fn_ty: Ty<'tcx>) -> FncTree {
-    // Check if TypeTrees are disabled via NoTT flag
-    if tcx.sess.opts.unstable_opts.autodiff.contains(&rustc_session::config::AutoDiff::NoTT) {
-        return FncTree { args: vec![], ret: TypeTree::new() };
-    }
-
-    // Check if this is actually a function type
-    if !fn_ty.is_fn() {
-        return FncTree { args: vec![], ret: TypeTree::new() };
-    }
-
-    // Get the function signature
-    let fn_sig = fn_ty.fn_sig(tcx);
-    let sig = tcx.instantiate_bound_regions_with_erased(fn_sig);
-
-    // Create TypeTrees for each input parameter
-    let mut args = vec![];
-    for ty in sig.inputs().iter() {
-        let type_tree = typetree_from_ty(tcx, *ty);
-        args.push(type_tree);
-    }
-
-    // Create TypeTree for return type
-    let ret = typetree_from_ty(tcx, sig.output());
-
-    FncTree { args, ret }
-}
-
-/// Generate TypeTree for a specific type.
-/// This function analyzes a Rust type and creates appropriate TypeTree metadata.
-pub fn typetree_from_ty<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> TypeTree {
-    let mut visited = Vec::new();
-    typetree_from_ty_inner(tcx, ty, 0, &mut visited)
-}
-
-/// Maximum recursion depth for TypeTree generation to prevent stack overflow
-/// from pathological deeply nested types. Combined with cycle detection.
-const MAX_TYPETREE_DEPTH: usize = 6;
-
-/// Internal recursive function for TypeTree generation with cycle detection and depth limiting.
-fn typetree_from_ty_inner<'tcx>(
-    tcx: TyCtxt<'tcx>,
-    ty: Ty<'tcx>,
-    depth: usize,
-    visited: &mut Vec<Ty<'tcx>>,
-) -> TypeTree {
-    if depth >= MAX_TYPETREE_DEPTH {
-        trace!("typetree depth limit {} reached for type: {}", MAX_TYPETREE_DEPTH, ty);
-        return TypeTree::new();
-    }
-
-    if visited.contains(&ty) {
-        return TypeTree::new();
-    }
-
-    visited.push(ty);
-    let result = typetree_from_ty_impl(tcx, ty, depth, visited);
-    visited.pop();
-    result
-}
-
-/// Implementation of TypeTree generation logic.
-fn typetree_from_ty_impl<'tcx>(
-    tcx: TyCtxt<'tcx>,
-    ty: Ty<'tcx>,
-    depth: usize,
-    visited: &mut Vec<Ty<'tcx>>,
-) -> TypeTree {
-    typetree_from_ty_impl_inner(tcx, ty, depth, visited, false)
-}
-
-/// Internal implementation with context about whether this is for a reference target.
-fn typetree_from_ty_impl_inner<'tcx>(
-    tcx: TyCtxt<'tcx>,
-    ty: Ty<'tcx>,
-    depth: usize,
-    visited: &mut Vec<Ty<'tcx>>,
-    is_reference_target: bool,
-) -> TypeTree {
-    if ty.is_scalar() {
-        let (kind, size) = if ty.is_integral() || ty.is_char() || ty.is_bool() {
-            (Kind::Integer, ty.primitive_size(tcx).bytes_usize())
-        } else if ty.is_floating_point() {
-            match ty {
-                x if x == tcx.types.f16 => (Kind::Half, 2),
-                x if x == tcx.types.f32 => (Kind::Float, 4),
-                x if x == tcx.types.f64 => (Kind::Double, 8),
-                x if x == tcx.types.f128 => (Kind::F128, 16),
-                _ => (Kind::Integer, 0),
-            }
-        } else {
-            (Kind::Integer, 0)
-        };
-
-        // Use offset 0 for scalars that are direct targets of references (like &f64)
-        // Use offset -1 for scalars used directly (like function return types)
-        let offset = if is_reference_target && !ty.is_array() { 0 } else { -1 };
-        return TypeTree(vec![Type { offset, size, kind, child: TypeTree::new() }]);
-    }
-
-    if ty.is_ref() || ty.is_raw_ptr() || ty.is_box() {
-        let Some(inner_ty) = ty.builtin_deref(true) else {
-            return TypeTree::new();
-        };
-
-        let child = typetree_from_ty_impl_inner(tcx, inner_ty, depth + 1, visited, true);
-        return TypeTree(vec![Type {
-            offset: -1,
-            size: tcx.data_layout.pointer_size().bytes_usize(),
-            kind: Kind::Pointer,
-            child,
-        }]);
-    }
-
-    if ty.is_array() {
-        if let ty::Array(element_ty, len_const) = ty.kind() {
-            let len = len_const.try_to_target_usize(tcx).unwrap_or(0);
-            if len == 0 {
-                return TypeTree::new();
-            }
-            let element_tree =
-                typetree_from_ty_impl_inner(tcx, *element_ty, depth + 1, visited, false);
-            let mut types = Vec::new();
-            for elem_type in &element_tree.0 {
-                types.push(Type {
-                    offset: -1,
-                    size: elem_type.size,
-                    kind: elem_type.kind,
-                    child: elem_type.child.clone(),
-                });
-            }
-
-            return TypeTree(types);
-        }
-    }
-
-    if ty.is_slice() {
-        if let ty::Slice(element_ty) = ty.kind() {
-            let element_tree =
-                typetree_from_ty_impl_inner(tcx, *element_ty, depth + 1, visited, false);
-            return element_tree;
-        }
-    }
-
-    if let ty::Tuple(tuple_types) = ty.kind() {
-        if tuple_types.is_empty() {
-            return TypeTree::new();
-        }
-
-        let mut types = Vec::new();
-        let mut current_offset = 0;
-
-        for tuple_ty in tuple_types.iter() {
-            let element_tree =
-                typetree_from_ty_impl_inner(tcx, tuple_ty, depth + 1, visited, false);
-
-            let element_layout = tcx
-                .layout_of(ty::TypingEnv::fully_monomorphized().as_query_input(tuple_ty))
-                .ok()
-                .map(|layout| layout.size.bytes_usize())
-                .unwrap_or(0);
-
-            for elem_type in &element_tree.0 {
-                types.push(Type {
-                    offset: if elem_type.offset == -1 {
-                        current_offset as isize
-                    } else {
-                        current_offset as isize + elem_type.offset
-                    },
-                    size: elem_type.size,
-                    kind: elem_type.kind,
-                    child: elem_type.child.clone(),
-                });
-            }
-
-            current_offset += element_layout;
-        }
-
-        return TypeTree(types);
-    }
-
-    if let ty::Adt(adt_def, args) = ty.kind() {
-        if adt_def.is_struct() {
-            let struct_layout =
-                tcx.layout_of(ty::TypingEnv::fully_monomorphized().as_query_input(ty));
-            if let Ok(layout) = struct_layout {
-                let mut types = Vec::new();
-
-                for (field_idx, field_def) in adt_def.all_fields().enumerate() {
-                    let field_ty = field_def.ty(tcx, args);
-                    let field_tree = typetree_from_ty_impl_inner(
-                        tcx,
-                        field_ty.skip_norm_wip(),
-                        depth + 1,
-                        visited,
-                        false,
-                    );
-
-                    let field_offset = layout.fields.offset(field_idx).bytes_usize();
-
-                    for elem_type in &field_tree.0 {
-                        types.push(Type {
-                            offset: if elem_type.offset == -1 {
-                                field_offset as isize
-                            } else {
-                                field_offset as isize + elem_type.offset
-                            },
-                            size: elem_type.size,
-                            kind: elem_type.kind,
-                            child: elem_type.child.clone(),
-                        });
-                    }
-                }
-
-                return TypeTree(types);
-            }
-        }
-    }
-
-    TypeTree::new()
-}