llmcc_core/
ir_builder.rs

1use std::collections::HashMap;
2use std::marker::PhantomData;
3use std::sync::atomic::{AtomicU32, Ordering};
4
5use tree_sitter::Node;
6
7use crate::block::BlockKind;
8use crate::context::{CompileCtxt, ParentedNode};
9use crate::ir::{
10    Arena, HirBase, HirFile, HirId, HirIdent, HirInternal, HirKind, HirNode, HirScope, HirText,
11};
12use crate::lang_def::LanguageTrait;
13
14/// Global atomic counter for HIR ID allocation
15static HIR_ID_COUNTER: AtomicU32 = AtomicU32::new(0);
16
17#[derive(Debug, Clone, Copy)]
18pub struct IrBuildConfig {
19    pub compact: bool,
20}
21
22impl IrBuildConfig {
23    pub fn compact() -> Self {
24        Self { compact: true }
25    }
26}
27
28impl Default for IrBuildConfig {
29    fn default() -> Self {
30        Self { compact: false }
31    }
32}
33
34/// Builder that directly assigns HIR nodes to compile context
35struct HirBuilder<'a, Language> {
36    arena: &'a Arena<'a>,
37    hir_map: HashMap<HirId, ParentedNode<'a>>,
38    file_path: Option<String>,
39    file_content: String,
40    config: IrBuildConfig,
41    _language: PhantomData<Language>,
42}
43
44impl<'a, Language: LanguageTrait> HirBuilder<'a, Language> {
45    /// Create a new builder that directly assigns to context
46    fn new(
47        arena: &'a Arena<'a>,
48        file_path: Option<String>,
49        file_content: String,
50        config: IrBuildConfig,
51    ) -> Self {
52        Self {
53            arena,
54            hir_map: HashMap::new(),
55            file_path,
56            file_content,
57            config,
58            _language: PhantomData,
59        }
60    }
61
62    /// Reserve a new HIR ID
63    fn reserve_hir_id(&self) -> HirId {
64        let id = HIR_ID_COUNTER.fetch_add(1, Ordering::SeqCst);
65        HirId(id)
66    }
67
68    fn build(mut self, root: Node<'a>) -> (HirId, HashMap<HirId, ParentedNode<'a>>) {
69        let file_start_id = self.build_node(root, None);
70        (file_start_id, self.hir_map)
71    }
72
73    fn build_node(&mut self, node: Node<'a>, parent: Option<HirId>) -> HirId {
74        let hir_id = self.reserve_hir_id();
75        let kind_id = node.kind_id();
76        let kind = Language::hir_kind(kind_id);
77        let block_kind = Language::block_kind(kind_id);
78        let child_ids = if self.should_collect_children(kind, block_kind) {
79            self.collect_children(node, hir_id)
80        } else {
81            Vec::new()
82        };
83        let base = self.make_base(hir_id, parent, node, kind, child_ids);
84
85        let hir_node = match kind {
86            HirKind::File => {
87                let path = self.file_path.clone().unwrap_or_default();
88                let file_node = HirFile::new(base, path);
89                HirNode::File(self.arena.alloc(file_node))
90            }
91            HirKind::Text => {
92                let text = self.extract_text(&base);
93                let text_node = HirText::new(base, text);
94                HirNode::Text(self.arena.alloc(text_node))
95            }
96            HirKind::Internal => {
97                let internal = HirInternal::new(base);
98                HirNode::Internal(self.arena.alloc(internal))
99            }
100            HirKind::Scope => {
101                // Try to extract the name identifier from the scope node
102                let ident = self.extract_scope_ident(&base, node);
103                let scope = HirScope::new(base, ident);
104                HirNode::Scope(self.arena.alloc(scope))
105            }
106            HirKind::Identifier => {
107                let text = self.extract_text(&base);
108                let ident = HirIdent::new(base, text);
109                HirNode::Ident(self.arena.alloc(ident))
110            }
111            other => panic!("unsupported HIR kind for node {:?}", (other, node)),
112        };
113
114        self.hir_map.insert(hir_id, ParentedNode::new(hir_node));
115        hir_id
116    }
117
118    fn collect_children(&mut self, node: Node<'a>, parent_id: HirId) -> Vec<HirId> {
119        let mut cursor = node.walk();
120
121        // In compact mode, skip children for Text nodes to reduce tree size
122        if self.config.compact {
123            let kind = Language::hir_kind(node.kind_id());
124            if kind == HirKind::Text {
125                return Vec::new();
126            }
127        }
128
129        node.children(&mut cursor)
130            .filter_map(|child| {
131                if self.config.compact {
132                    if child.is_error() || child.is_extra() || child.is_missing() || !child.is_named()
133                    {
134                        return None;
135                    }
136                    let child_kind = Language::hir_kind(child.kind_id());
137                    if child_kind == HirKind::Text {
138                        return None;
139                    }
140                    let child_block_kind = Language::block_kind(child.kind_id());
141                    if matches!(child_block_kind, BlockKind::Stmt | BlockKind::Call) {
142                        return None;
143                    }
144                    // Keep Python/Rust block nodes so nested declarations remain visible in compact mode.
145                }
146                Some(self.build_node(child, Some(parent_id)))
147            })
148            .collect()
149    }
150
151    fn should_collect_children(&self, kind: HirKind, block_kind: BlockKind) -> bool {
152        if !self.config.compact {
153            return true;
154        }
155
156        match kind {
157            HirKind::File => true,
158            HirKind::Scope => matches!(
159                block_kind,
160                BlockKind::Root
161                    | BlockKind::Scope
162                    | BlockKind::Class
163                    | BlockKind::Enum
164                    | BlockKind::Impl
165                    | BlockKind::Func
166                    | BlockKind::Const
167            ),
168            HirKind::Internal => matches!(
169                block_kind,
170                BlockKind::Scope | BlockKind::Field | BlockKind::Const | BlockKind::Undefined
171            ),
172            _ => false,
173        }
174    }
175
176    fn make_base(
177        &self,
178        hir_id: HirId,
179        parent: Option<HirId>,
180        node: Node<'a>,
181        kind: HirKind,
182        children: Vec<HirId>,
183    ) -> HirBase<'a> {
184        let field_id = Self::field_id_of(node).unwrap_or(u16::MAX);
185        HirBase {
186            hir_id,
187            parent,
188            node,
189            kind,
190            field_id,
191            children,
192        }
193    }
194
195    fn extract_text(&self, base: &HirBase<'a>) -> String {
196        let start = base.node.start_byte();
197        let end = base.node.end_byte();
198        if end > start && end <= self.file_content.len() {
199            self.file_content[start..end].to_string()
200        } else {
201            String::new()
202        }
203    }
204
205    fn extract_scope_ident(&self, base: &HirBase<'a>, node: Node<'a>) -> Option<&'a HirIdent<'a>> {
206        // Try to get the name field from the tree-sitter node
207        // For Rust, the name field is typically "name"
208        let name_node = node.child_by_field_name("name")?;
209
210        // Create an identifier for the name node
211        let hir_id = self.reserve_hir_id();
212        let ident_base = HirBase {
213            hir_id,
214            parent: Some(base.hir_id),
215            node: name_node,
216            kind: HirKind::Identifier,
217            field_id: u16::MAX,
218            children: Vec::new(),
219        };
220
221        let text = self.extract_text(&ident_base);
222        let ident = HirIdent::new(ident_base, text);
223        Some(self.arena.alloc(ident))
224    }
225
226    fn field_id_of(node: Node<'_>) -> Option<u16> {
227        let parent = node.parent()?;
228        let mut cursor = parent.walk();
229
230        if !cursor.goto_first_child() {
231            return None;
232        }
233
234        loop {
235            if cursor.node().id() == node.id() {
236                return cursor.field_id().map(|id| id.get());
237            }
238            if !cursor.goto_next_sibling() {
239                break;
240            }
241        }
242
243        None
244    }
245}
246
247pub fn build_llmcc_ir_inner<'a, L: LanguageTrait>(
248    arena: &'a Arena<'a>,
249    file_path: Option<String>,
250    file_content: String,
251    tree: &'a tree_sitter::Tree,
252    config: IrBuildConfig,
253) -> Result<(HirId, HashMap<HirId, ParentedNode<'a>>), Box<dyn std::error::Error>> {
254    let builder = HirBuilder::<L>::new(arena, file_path, file_content, config);
255    let root = tree.root_node();
256    let result = builder.build(root);
257    Ok(result)
258}
259
260/// Build IR for all units in the context
261/// TODO: make this run in parallel
262pub fn build_llmcc_ir<'a, L: LanguageTrait>(
263    cc: &'a CompileCtxt<'a>,
264) -> Result<(), Box<dyn std::error::Error>> {
265    build_llmcc_ir_with_config::<L>(cc, IrBuildConfig::default())
266}
267
268/// Build IR for all units in the context with custom config
269pub fn build_llmcc_ir_with_config<'a, L: LanguageTrait>(
270    cc: &'a CompileCtxt<'a>,
271    config: IrBuildConfig,
272) -> Result<(), Box<dyn std::error::Error>> {
273    for index in 0..cc.files.len() {
274        let unit = cc.compile_unit(index);
275        let file_path = unit.file_path().map(|p| p.to_string());
276        let file_content = String::from_utf8_lossy(&unit.file().content()).to_string();
277        let tree = unit.tree();
278
279        let (_file_start_id, hir_map) =
280            build_llmcc_ir_inner::<L>(&cc.arena, file_path, file_content, tree, config)?;
281
282        // Insert all nodes into the compile context
283        for (hir_id, parented_node) in hir_map {
284            cc.hir_map.borrow_mut().insert(hir_id, parented_node);
285        }
286        cc.set_file_start(index, _file_start_id);
287    }
288    Ok(())
289}
llmcc_core/ir_builder.rs

llmcc_core/
ir_builder.rs
