rustc_ast/expand/
typetree.rs

1//! This module contains the definition of the `TypeTree` and `Type` structs.
2//! They are thin Rust wrappers around the TypeTrees used by Enzyme as the LLVM based autodiff
3//! backend. The Enzyme TypeTrees currently have various limitations and should be rewritten, so the
4//! Rust frontend obviously has the same limitations. The main motivation of TypeTrees is to
5//! represent how a type looks like "in memory". Enzyme can deduce this based on usage patterns in
6//! the user code, but this is extremely slow and not even always sufficient. As such we lower some
7//! information from rustc to help Enzyme. For a full explanation of their design it is necessary to
8//! analyze the implementation in Enzyme core itself. As a rough summary, `-1` in Enzyme speech means
9//! everywhere. That is `{0:-1: Float}` means at index 0 you have a ptr, if you dereference it it
10//! will be floats everywhere. Thus `* f32`. If you have `{-1:int}` it means int's everywhere,
11//! e.g. [i32; N]. `{0:-1:-1 float}` then means one pointer at offset 0, if you dereference it there
12//! will be only pointers, if you dereference these new pointers they will point to array of floats.
13//! Generally, it allows byte-specific descriptions.
14//! FIXME: This description might be partly inaccurate and should be extended, along with
15//! adding documentation to the corresponding Enzyme core code.
16//! FIXME: Rewrite the TypeTree logic in Enzyme core to reduce the need for the rustc frontend to
17//! provide typetree information.
18//! FIXME: We should also re-evaluate where we create TypeTrees from Rust types, since MIR
19//! representations of some types might not be accurate. For example a vector of floats might be
20//! represented as a vector of u8s in MIR in some cases.
21
22use std::fmt;
23
24use crate::expand::{Decodable, Encodable, HashStable_Generic};
25
26#[derive(Clone, Copy, Eq, PartialEq, Encodable, Decodable, Debug, HashStable_Generic)]
27pub enum Kind {
28    Anything,
29    Integer,
30    Pointer,
31    Half,
32    Float,
33    Double,
34    F128,
35    Unknown,
36}
37
38#[derive(Clone, Eq, PartialEq, Encodable, Decodable, Debug, HashStable_Generic)]
39pub struct TypeTree(pub Vec<Type>);
40
41impl TypeTree {
42    pub fn new() -> Self {
43        Self(Vec::new())
44    }
45    pub fn all_ints() -> Self {
46        Self(vec![Type { offset: -1, size: 1, kind: Kind::Integer, child: TypeTree::new() }])
47    }
48    pub fn int(size: usize) -> Self {
49        let mut ints = Vec::with_capacity(size);
50        for i in 0..size {
51            ints.push(Type {
52                offset: i as isize,
53                size: 1,
54                kind: Kind::Integer,
55                child: TypeTree::new(),
56            });
57        }
58        Self(ints)
59    }
60}
61
62#[derive(Clone, Eq, PartialEq, Encodable, Decodable, Debug, HashStable_Generic)]
63pub struct FncTree {
64    pub args: Vec<TypeTree>,
65    pub ret: TypeTree,
66}
67
68#[derive(Clone, Eq, PartialEq, Encodable, Decodable, Debug, HashStable_Generic)]
69pub struct Type {
70    pub offset: isize,
71    pub size: usize,
72    pub kind: Kind,
73    pub child: TypeTree,
74}
75
76impl Type {
77    pub fn add_offset(self, add: isize) -> Self {
78        let offset = match self.offset {
79            -1 => add,
80            x => add + x,
81        };
82
83        Self { size: self.size, kind: self.kind, child: self.child, offset }
84    }
85}
86
87impl fmt::Display for Type {
88    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
89        <Self as fmt::Debug>::fmt(self, f)
90    }
91}
rustc_ast/expand/typetree.rs

rustc_ast/expand/
typetree.rs
