llmcc_core/
ir.rs

1use strum_macros::{Display, EnumIter, EnumString, FromRepr};
2use tree_sitter::Node;
3
4use crate::context::CompileUnit;
5use crate::declare_arena;
6use crate::symbol::{Scope, Symbol};
7
8// Declare the arena with all HIR types
9declare_arena!([
10    hir_root: HirRoot<'tcx>,
11    hir_text: HirText<'tcx>,
12    hir_internal: HirInternal<'tcx>,
13    hir_scope: HirScope<'tcx>,
14    hir_file: HirFile<'tcx>,
15    hir_ident: HirIdent<'tcx>,
16    symbol: Symbol,
17    scope: Scope<'tcx>,
18]);
19
20#[derive(
21    Debug, Clone, Copy, PartialEq, Eq, Hash, EnumIter, EnumString, FromRepr, Display, Default,
22)]
23#[strum(serialize_all = "snake_case")]
24pub enum HirKind {
25    #[default]
26    Undefined,
27    Error,
28    File,
29    Scope,
30    Text,
31    Internal,
32    Comment,
33    Identifier,
34}
35
36#[derive(Debug, Clone, Copy, Default)]
37pub enum HirNode<'hir> {
38    #[default]
39    Undefined,
40    Root(&'hir HirRoot<'hir>),
41    Text(&'hir HirText<'hir>),
42    Internal(&'hir HirInternal<'hir>),
43    Scope(&'hir HirScope<'hir>),
44    File(&'hir HirFile<'hir>),
45    Ident(&'hir HirIdent<'hir>),
46}
47
48impl<'hir> HirNode<'hir> {
49    pub fn format_node(&self, unit: CompileUnit<'hir>) -> String {
50        let id = self.hir_id();
51        let kind = self.kind();
52        let mut f = format!("{}:{}", kind, id);
53
54        // if let Some(def) = unit.opt_defs(id) {
55        //     f.push_str(&format!("   d:{}", def.format_compact()));
56        // } else if let Some(sym) = unit.opt_uses(id) {
57        //     f.push_str(&format!("   u:{}", sym.format_compact()));
58        // }
59
60        if let Some(scope) = unit.opt_get_scope(id) {
61            f.push_str(&format!("   s:{}", scope.format_compact()));
62        }
63
64        f
65    }
66
67    /// Get the base information for any HIR node
68    pub fn base(&self) -> Option<&HirBase<'hir>> {
69        match self {
70            HirNode::Undefined => None,
71            HirNode::Root(node) => Some(&node.base),
72            HirNode::Text(node) => Some(&node.base),
73            HirNode::Internal(node) => Some(&node.base),
74            HirNode::Scope(node) => Some(&node.base),
75            HirNode::File(node) => Some(&node.base),
76            HirNode::Ident(node) => Some(&node.base),
77        }
78    }
79
80    /// Get the kind of this HIR node
81    pub fn kind(&self) -> HirKind {
82        self.base().map_or(HirKind::Undefined, |base| base.kind)
83    }
84
85    /// Check if this node is of a specific kind
86    pub fn is_kind(&self, kind: HirKind) -> bool {
87        self.kind() == kind
88    }
89
90    pub fn field_id(&self) -> u16 {
91        self.base().unwrap().field_id
92    }
93
94    /// Get children of this node
95    pub fn children(&self) -> &[HirId] {
96        self.base().map_or(&[], |base| &base.children)
97    }
98
99    pub fn kind_id(&self) -> u16 {
100        self.base().unwrap().node.kind_id()
101    }
102
103    pub fn hir_id(&self) -> HirId {
104        self.base().unwrap().hir_id
105    }
106
107    pub fn start_byte(&self) -> usize {
108        self.base().unwrap().node.start_byte()
109    }
110
111    pub fn end_byte(&self) -> usize {
112        self.base().unwrap().node.end_byte()
113    }
114
115    pub fn child_count(&self) -> usize {
116        self.children().len()
117    }
118
119    pub fn inner_ts_node(&self) -> Node<'hir> {
120        self.base().unwrap().node
121    }
122
123    pub fn parent(&self) -> Option<HirId> {
124        self.base().and_then(|base| base.parent)
125    }
126
127    pub fn opt_child_by_field(
128        &self,
129        unit: CompileUnit<'hir>,
130        field_id: u16,
131    ) -> Option<HirNode<'hir>> {
132        self.base().unwrap().opt_child_by_field(unit, field_id)
133    }
134
135    pub fn child_by_field(&self, unit: CompileUnit<'hir>, field_id: u16) -> HirNode<'hir> {
136        self.opt_child_by_field(unit, field_id)
137            .unwrap_or_else(|| panic!("no child with field_id {}", field_id))
138    }
139
140    pub fn expect_ident_child_by_field(
141        &self,
142        unit: CompileUnit<'hir>,
143        field_id: u16,
144    ) -> &'hir HirIdent<'hir> {
145        self.opt_child_by_field(unit, field_id)
146            .map(|child| child.expect_ident())
147            .unwrap_or_else(|| panic!("no child with field_id {}", field_id))
148    }
149
150    pub fn opt_child_by_kind(
151        &self,
152        unit: CompileUnit<'hir>,
153        kind_id: u16,
154    ) -> Option<HirNode<'hir>> {
155        self.children()
156            .iter()
157            .map(|id| unit.hir_node(*id))
158            .find(|child| child.kind_id() == kind_id)
159    }
160
161    pub fn child_by_kind(&self, unit: CompileUnit<'hir>, kind_id: u16) -> HirNode<'hir> {
162        self.opt_child_by_kind(unit, kind_id)
163            .unwrap_or_else(|| panic!("no child with kind_id {}", kind_id))
164    }
165
166    /// Recursively search for an identifier within this node.
167    ///
168    /// Useful for finding the actual identifier in complex AST nodes like generic_type
169    /// that wrap the identifier. For example, in `impl<'tcx> Holder<'tcx>`, the type
170    /// field points to a generic_type node, which contains the type_identifier "Holder".
171    pub fn find_ident(&self, unit: CompileUnit<'hir>) -> Option<&'hir HirIdent<'hir>> {
172        // Check if this node is already an identifier
173        if let Some(ident) = self.as_ident() {
174            return Some(ident);
175        }
176
177        // Otherwise, search through children of any node that has them
178        let children = match self {
179            HirNode::Root(r) => &r.base.children,
180            HirNode::Text(_) => return None,
181            HirNode::Internal(i) => &i.base.children,
182            HirNode::Scope(s) => &s.base.children,
183            HirNode::File(f) => &f.base.children,
184            HirNode::Ident(_) => return None,
185            HirNode::Undefined => return None,
186        };
187
188        // Recursively search all children
189        for child_id in children {
190            let child = unit.hir_node(*child_id);
191            if let Some(ident) = child.find_ident(unit) {
192                return Some(ident);
193            }
194        }
195
196        None
197    }
198}
199
200macro_rules! impl_getters {
201    ($($variant:ident => $type:ty),* $(,)?) => {
202        impl<'hir> HirNode<'hir> {
203            $(
204                paste::paste! {
205                    pub fn [<as_ $variant:lower>](&self) -> Option<$type> {
206                        match self {
207                            HirNode::$variant(r) => Some(r),
208                            _ => None,
209                        }
210                    }
211
212                    pub fn [<expect_ $variant:lower>](&self) -> $type {
213                        match self {
214                            HirNode::$variant(r) => r,
215                            _ => panic!("Expected {} variant", stringify!($variant)),
216                        }
217                    }
218
219                    pub fn [<is_ $variant:lower>](&self) -> bool {
220                        matches!(self, HirNode::$variant(_))
221                    }
222                }
223            )*
224        }
225    };
226}
227
228impl_getters! {
229    Root => &'hir HirRoot<'hir>,
230    Text => &'hir HirText<'hir>,
231    Internal => &'hir HirInternal<'hir>,
232    Scope => &'hir HirScope<'hir>,
233    File => &'hir HirFile<'hir>,
234    Ident => &'hir HirIdent<'hir>,
235}
236
237#[derive(Copy, Clone, PartialEq, Eq, Debug, Hash, Default)]
238pub struct HirId(pub u32);
239
240impl std::fmt::Display for HirId {
241    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
242        write!(f, "{}", self.0)
243    }
244}
245
246#[derive(Debug, Clone)]
247pub struct HirBase<'hir> {
248    pub hir_id: HirId,
249    pub parent: Option<HirId>,
250    pub node: Node<'hir>,
251    pub kind: HirKind,
252    pub field_id: u16,
253    pub children: Vec<HirId>,
254}
255
256impl<'hir> HirBase<'hir> {
257    pub fn opt_child_by_fields(
258        &self,
259        unit: CompileUnit<'hir>,
260        fields_id: &[u16],
261    ) -> Option<HirNode<'hir>> {
262        self.children
263            .iter()
264            .map(|id| unit.hir_node(*id))
265            .find(|child| fields_id.contains(&child.field_id()))
266    }
267
268    pub fn opt_child_by_field(
269        &self,
270        unit: CompileUnit<'hir>,
271        field_id: u16,
272    ) -> Option<HirNode<'hir>> {
273        self.children
274            .iter()
275            .map(|id| unit.hir_node(*id))
276            .find(|child| child.field_id() == field_id)
277    }
278}
279
280#[derive(Debug, Clone)]
281pub struct HirRoot<'hir> {
282    pub base: HirBase<'hir>,
283    pub file_name: Option<String>,
284}
285
286impl<'hir> HirRoot<'hir> {
287    pub fn new(base: HirBase<'hir>, file_name: Option<String>) -> Self {
288        Self { base, file_name }
289    }
290}
291
292#[derive(Debug, Clone)]
293pub struct HirText<'hir> {
294    pub base: HirBase<'hir>,
295    pub text: String,
296}
297
298impl<'hir> HirText<'hir> {
299    pub fn new(base: HirBase<'hir>, text: String) -> Self {
300        Self { base, text }
301    }
302}
303
304#[derive(Debug, Clone)]
305pub struct HirInternal<'hir> {
306    pub base: HirBase<'hir>,
307}
308
309impl<'hir> HirInternal<'hir> {
310    pub fn new(base: HirBase<'hir>) -> Self {
311        Self { base }
312    }
313}
314
315#[derive(Debug, Clone)]
316pub struct HirScope<'hir> {
317    pub base: HirBase<'hir>,
318    pub ident: Option<&'hir HirIdent<'hir>>,
319}
320
321impl<'hir> HirScope<'hir> {
322    pub fn new(base: HirBase<'hir>, ident: Option<&'hir HirIdent<'hir>>) -> Self {
323        Self { base, ident }
324    }
325
326    pub fn owner_name(&self) -> String {
327        if let Some(id) = self.ident {
328            id.name.clone()
329        } else {
330            "unamed_scope".to_string()
331        }
332    }
333}
334
335#[derive(Debug, Clone)]
336pub struct HirIdent<'hir> {
337    pub base: HirBase<'hir>,
338    pub name: String,
339}
340
341impl<'hir> HirIdent<'hir> {
342    pub fn new(base: HirBase<'hir>, name: String) -> Self {
343        Self { base, name }
344    }
345}
346
347#[derive(Debug, Clone)]
348pub struct HirFile<'hir> {
349    pub base: HirBase<'hir>,
350    pub file_path: String,
351}
352
353impl<'hir> HirFile<'hir> {
354    pub fn new(base: HirBase<'hir>, file_path: String) -> Self {
355        Self { base, file_path }
356    }
357}
llmcc_core/ir.rs

llmcc_core/
ir.rs
