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                    if child_block_kind == BlockKind::Scope && child.kind() == "block" {
145                        return None;
146                    }
147                }
148                Some(self.build_node(child, Some(parent_id)))
149            })
150            .collect()
151    }
152
153    fn should_collect_children(&self, kind: HirKind, block_kind: BlockKind) -> bool {
154        if !self.config.compact {
155            return true;
156        }
157
158        match kind {
159            HirKind::File => true,
160            HirKind::Scope => matches!(
161                block_kind,
162                BlockKind::Root
163                    | BlockKind::Scope
164                    | BlockKind::Class
165                    | BlockKind::Enum
166                    | BlockKind::Impl
167                    | BlockKind::Func
168                    | BlockKind::Const
169            ),
170            HirKind::Internal => matches!(
171                block_kind,
172                BlockKind::Scope | BlockKind::Field | BlockKind::Const | BlockKind::Undefined
173            ),
174            _ => false,
175        }
176    }
177
178    fn make_base(
179        &self,
180        hir_id: HirId,
181        parent: Option<HirId>,
182        node: Node<'a>,
183        kind: HirKind,
184        children: Vec<HirId>,
185    ) -> HirBase<'a> {
186        let field_id = Self::field_id_of(node).unwrap_or(u16::MAX);
187        HirBase {
188            hir_id,
189            parent,
190            node,
191            kind,
192            field_id,
193            children,
194        }
195    }
196
197    fn extract_text(&self, base: &HirBase<'a>) -> String {
198        let start = base.node.start_byte();
199        let end = base.node.end_byte();
200        if end > start && end <= self.file_content.len() {
201            self.file_content[start..end].to_string()
202        } else {
203            String::new()
204        }
205    }
206
207    fn extract_scope_ident(&self, base: &HirBase<'a>, node: Node<'a>) -> Option<&'a HirIdent<'a>> {
208        // Try to get the name field from the tree-sitter node
209        // For Rust, the name field is typically "name"
210        let name_node = node.child_by_field_name("name")?;
211
212        // Create an identifier for the name node
213        let hir_id = self.reserve_hir_id();
214        let ident_base = HirBase {
215            hir_id,
216            parent: Some(base.hir_id),
217            node: name_node,
218            kind: HirKind::Identifier,
219            field_id: u16::MAX,
220            children: Vec::new(),
221        };
222
223        let text = self.extract_text(&ident_base);
224        let ident = HirIdent::new(ident_base, text);
225        Some(self.arena.alloc(ident))
226    }
227
228    fn field_id_of(node: Node<'_>) -> Option<u16> {
229        let parent = node.parent()?;
230        let mut cursor = parent.walk();
231
232        if !cursor.goto_first_child() {
233            return None;
234        }
235
236        loop {
237            if cursor.node().id() == node.id() {
238                return cursor.field_id().map(|id| id.get());
239            }
240            if !cursor.goto_next_sibling() {
241                break;
242            }
243        }
244
245        None
246    }
247}
248
249pub fn build_llmcc_ir_inner<'a, L: LanguageTrait>(
250    arena: &'a Arena<'a>,
251    file_path: Option<String>,
252    file_content: String,
253    tree: &'a tree_sitter::Tree,
254    config: IrBuildConfig,
255) -> Result<(HirId, HashMap<HirId, ParentedNode<'a>>), Box<dyn std::error::Error>> {
256    let builder = HirBuilder::<L>::new(arena, file_path, file_content, config);
257    let root = tree.root_node();
258    let result = builder.build(root);
259    Ok(result)
260}
261
262/// Build IR for all units in the context
263/// TODO: make this run in parallel
264pub fn build_llmcc_ir<'a, L: LanguageTrait>(
265    cc: &'a CompileCtxt<'a>,
266) -> Result<(), Box<dyn std::error::Error>> {
267    build_llmcc_ir_with_config::<L>(cc, IrBuildConfig::default())
268}
269
270/// Build IR for all units in the context with custom config
271pub fn build_llmcc_ir_with_config<'a, L: LanguageTrait>(
272    cc: &'a CompileCtxt<'a>,
273    config: IrBuildConfig,
274) -> Result<(), Box<dyn std::error::Error>> {
275    for index in 0..cc.files.len() {
276        let unit = cc.compile_unit(index);
277        let file_path = unit.file_path().map(|p| p.to_string());
278        let file_content = String::from_utf8_lossy(&unit.file().content()).to_string();
279        let tree = unit.tree();
280
281        let (_file_start_id, hir_map) =
282            build_llmcc_ir_inner::<L>(&cc.arena, file_path, file_content, tree, config)?;
283
284        // Insert all nodes into the compile context
285        for (hir_id, parented_node) in hir_map {
286            cc.hir_map.borrow_mut().insert(hir_id, parented_node);
287        }
288        cc.set_file_start(index, _file_start_id);
289    }
290    Ok(())
291}
llmcc_core/ir_builder.rs

llmcc_core/
ir_builder.rs
