llmcc_rust/
collect.rs

1use std::mem;
2
3use llmcc_core::context::CompileUnit;
4use llmcc_core::ir::{HirBase, HirId, HirIdent, HirKind, HirNode};
5use llmcc_core::symbol::{Scope, ScopeStack, Symbol, SymbolKind};
6
7use crate::descriptor::function::parse_type_expr;
8use crate::descriptor::{
9    CallDescriptor, EnumDescriptor, FnVisibility, FunctionDescriptor, StructDescriptor, TypeExpr,
10    VariableDescriptor,
11};
12use crate::token::{AstVisitorRust, LangRust};
13
14/// DeclCollector collects all symbol declarations (functions, structs, enums, variables, etc.)
15/// from the AST during the first pass of semantic analysis.
16///
17/// For each symbol encountered:
18/// 1. Insert it into the local scope stack for local name resolution
19/// 2. If it's public/global, also register it in the global symbol table for cross-file resolution
20///
21/// This enables later phases to resolve both local names (via scope stack) and qualified names
22/// (via global symbol table).
23#[derive(Debug)]
24struct DeclCollector<'tcx> {
25    unit: CompileUnit<'tcx>,
26    scopes: ScopeStack<'tcx>,
27    functions: Vec<FunctionDescriptor>,
28    variables: Vec<VariableDescriptor>,
29    calls: Vec<CallDescriptor>,
30    structs: Vec<StructDescriptor>,
31    enums: Vec<EnumDescriptor>,
32}
33
34impl<'tcx> DeclCollector<'tcx> {
35    pub fn new(unit: CompileUnit<'tcx>, globals: &'tcx Scope<'tcx>) -> Self {
36        let mut scopes = ScopeStack::new(&unit.cc.arena, &unit.cc.interner, &unit.cc.symbol_map);
37        // TODO: make create a new symbol assoicate unit file name
38        scopes.push_with_symbol(globals, None);
39        Self {
40            unit,
41            scopes,
42            functions: Vec::new(),
43            variables: Vec::new(),
44            calls: Vec::new(),
45            structs: Vec::new(),
46            enums: Vec::new(),
47        }
48    }
49
50    fn parent_symbol(&self) -> Option<&'tcx Symbol> {
51        self.scopes.scoped_symbol()
52    }
53
54    fn scoped_fqn(&self, _node: &HirNode<'tcx>, name: &str) -> String {
55        if let Some(parent) = self.parent_symbol() {
56            let parent_fqn = parent.fqn_name.borrow();
57            if parent_fqn.is_empty() {
58                name.to_string()
59            } else {
60                format!("{}::{}", parent_fqn.as_str(), name)
61            }
62        } else {
63            name.to_string()
64        }
65    }
66
67    fn take_functions(&mut self) -> Vec<FunctionDescriptor> {
68        mem::take(&mut self.functions)
69    }
70
71    fn take_variables(&mut self) -> Vec<VariableDescriptor> {
72        mem::take(&mut self.variables)
73    }
74
75    fn take_calls(&mut self) -> Vec<CallDescriptor> {
76        mem::take(&mut self.calls)
77    }
78
79    fn take_structs(&mut self) -> Vec<StructDescriptor> {
80        mem::take(&mut self.structs)
81    }
82
83    fn take_enums(&mut self) -> Vec<EnumDescriptor> {
84        mem::take(&mut self.enums)
85    }
86
87    fn create_new_symbol(
88        &mut self,
89        node: &HirNode<'tcx>,
90        field_id: u16,
91        global: bool,
92        kind: SymbolKind,
93    ) -> Option<(&'tcx Symbol, &'tcx HirIdent<'tcx>, String)> {
94        let ident_node = node.opt_child_by_field(self.unit, field_id)?;
95        let ident = ident_node.as_ident()?;
96        let fqn = self.scoped_fqn(node, &ident.name);
97        let owner = node.hir_id();
98
99        let symbol = match self.scopes.find_symbol_local(&ident.name) {
100            Some(existing) if existing.unit_index() == Some(self.unit.index) => {
101                if Self::different_kind(existing.kind(), kind) {
102                    self.insert_into_scope(owner, ident, global, &fqn, kind)
103                } else {
104                    existing
105                }
106            }
107            _ => self.insert_into_scope(owner, ident, global, &fqn, kind),
108        };
109
110        Some((symbol, ident, fqn))
111    }
112
113    fn different_kind(existing_kind: SymbolKind, new_kind: SymbolKind) -> bool {
114        existing_kind != SymbolKind::Unknown && existing_kind != new_kind
115    }
116
117    // fn warn_duplicate_symbol(&self, name: &str, existing_kind: SymbolKind, new_kind: SymbolKind) {
118    //     eprintln!(
119    //         "warning: duplicate symbol '{}' found in the same scope. existing kind: {:?}, new kind: {:?}. skipping insertion.",
120    //         name, existing_kind, new_kind
121    //     );
122    // }
123
124    fn insert_into_scope(
125        &mut self,
126        owner: HirId,
127        ident: &'tcx HirIdent<'tcx>,
128        global: bool,
129        fqn: &str,
130        kind: SymbolKind,
131    ) -> &'tcx Symbol {
132        let interner = self.unit.interner();
133        let unit_index = self.unit.index;
134
135        self.scopes.insert_with(owner, ident, global, |symbol| {
136            symbol.set_owner(owner);
137            symbol.set_fqn(fqn.to_string(), interner);
138            symbol.set_kind(kind);
139            symbol.set_unit_index(unit_index);
140        })
141    }
142
143    fn has_public_visibility(&self, node: &HirNode<'tcx>) -> bool {
144        let ts_node = node.inner_ts_node();
145        let Some(name_node) = ts_node.child_by_field_name("name") else {
146            return false;
147        };
148        let header = self
149            .unit
150            .file()
151            .opt_get_text(ts_node.start_byte(), name_node.start_byte())
152            .unwrap_or_default();
153
154        !matches!(FnVisibility::from_header(&header), FnVisibility::Private)
155    }
156
157    fn should_register_globally(&self, node: &HirNode<'tcx>) -> bool {
158        self.parent_symbol().is_none() || self.has_public_visibility(node)
159    }
160
161    fn enum_variant_should_register_globally(&self, _node: &HirNode<'tcx>) -> bool {
162        let Some(enum_symbol) = self.scopes.scoped_symbol() else {
163            return false;
164        };
165
166        if enum_symbol.kind() != SymbolKind::Enum {
167            return false;
168        }
169
170        let parent_node = self.unit.hir_node(enum_symbol.owner());
171        self.has_public_visibility(&parent_node)
172    }
173
174    fn visit_children_new_scope(
175        &mut self,
176        node: &HirNode<'tcx>,
177        scoped_symbol: Option<&'tcx Symbol>,
178    ) {
179        let depth = self.scopes.depth();
180        let scope = self.unit.alloc_scope(node.hir_id());
181
182        let symbol = scoped_symbol.or_else(|| self.scopes.scoped_symbol());
183        self.scopes.push_with_symbol(scope, symbol);
184        self.visit_children(node);
185        self.scopes.pop_until(depth);
186    }
187}
188
189impl<'tcx> AstVisitorRust<'tcx> for DeclCollector<'tcx> {
190    fn unit(&self) -> CompileUnit<'tcx> {
191        self.unit
192    }
193
194    fn visit_source_file(&mut self, node: HirNode<'tcx>) {
195        self.visit_children_new_scope(&node, None);
196    }
197
198    fn visit_function_item(&mut self, node: HirNode<'tcx>) {
199        let register_globally = self.should_register_globally(&node);
200        if let Some((symbol, _ident, fqn)) = self.create_new_symbol(
201            &node,
202            LangRust::field_name,
203            register_globally,
204            SymbolKind::Function,
205        ) {
206            if let Some(desc) = FunctionDescriptor::from_hir(self.unit, &node, fqn.clone()) {
207                self.functions.push(desc);
208            }
209            self.visit_children_new_scope(&node, Some(symbol));
210        } else {
211            self.visit_children(&node);
212        }
213    }
214
215    fn visit_let_declaration(&mut self, node: HirNode<'tcx>) {
216        if let Some((_symbol, ident, fqn)) =
217            self.create_new_symbol(&node, LangRust::field_pattern, false, SymbolKind::Variable)
218        {
219            let var =
220                VariableDescriptor::from_let(self.unit, &node, ident.name.clone(), fqn.clone());
221            self.variables.push(var);
222        }
223        self.visit_children(&node);
224    }
225
226    fn visit_block(&mut self, node: HirNode<'tcx>) {
227        self.visit_children_new_scope(&node, None);
228    }
229
230    fn visit_parameter(&mut self, node: HirNode<'tcx>) {
231        let _ = self.create_new_symbol(&node, LangRust::field_pattern, false, SymbolKind::Variable);
232        self.visit_children(&node);
233    }
234
235    fn visit_mod_item(&mut self, node: HirNode<'tcx>) {
236        let symbol = self
237            .create_new_symbol(&node, LangRust::field_name, true, SymbolKind::Module)
238            .map(|(symbol, _ident, _)| symbol);
239        self.visit_children_new_scope(&node, symbol);
240    }
241
242    fn visit_impl_item(&mut self, node: HirNode<'tcx>) {
243        let symbol = node
244            .opt_child_by_field(self.unit, LangRust::field_type)
245            .and_then(|type_node| {
246                let segments = impl_type_segments(self.unit, &type_node)?;
247                if segments.is_empty() {
248                    return None;
249                }
250
251                let keys: Vec<_> = segments
252                    .iter()
253                    .map(|segment| self.unit.interner().intern(segment))
254                    .collect();
255
256                if let Some(symbol) = self
257                    .scopes
258                    .find_global_suffix_in_unit(&keys, self.unit.index)
259                    .or_else(|| self.scopes.find_global_suffix(&keys))
260                {
261                    symbol.set_owner(node.hir_id());
262                    symbol.set_unit_index(self.unit.index);
263                    Some(symbol)
264                } else {
265                    let fqn = segments.join("::");
266                    let owner = node.hir_id();
267
268                    // For impl blocks with qualified types like `impl From<T> for module::Type`,
269                    // use the last segment (the actual type name) as the symbol name.
270                    // This ensures proper trie indexing.
271                    let impl_name = segments
272                        .last()
273                        .cloned()
274                        .unwrap_or_else(|| "impl".to_string());
275
276                    // Create a synthetic identifier using the impl target type name
277                    let synthetic_ident = self.unit.cc.arena.alloc(HirIdent::new(
278                        // Create a minimal HirBase using the impl node's information
279                        HirBase {
280                            hir_id: owner,
281                            parent: node.parent(),
282                            node: node.inner_ts_node(),
283                            kind: HirKind::Identifier,
284                            field_id: u16::MAX,
285                            children: Vec::new(),
286                        },
287                        impl_name,
288                    ));
289
290                    let global = false;
291                    let kind = SymbolKind::Impl;
292                    let symbol =
293                        self.scopes
294                            .insert_with(owner, synthetic_ident, global, |symbol| {
295                                symbol.set_owner(owner);
296                                symbol.set_fqn(fqn, self.unit.interner());
297                                symbol.set_kind(kind);
298                                symbol.set_unit_index(self.unit.index);
299                            });
300                    Some(symbol)
301                }
302            });
303        self.visit_children_new_scope(&node, symbol);
304    }
305
306    fn visit_trait_item(&mut self, node: HirNode<'tcx>) {
307        self.visit_struct_item(node);
308    }
309
310    fn visit_function_signature_item(&mut self, node: HirNode<'tcx>) {
311        self.visit_function_item(node);
312    }
313
314    fn visit_call_expression(&mut self, node: HirNode<'tcx>) {
315        let enclosing = self
316            .parent_symbol()
317            .map(|symbol| symbol.fqn_name.borrow().clone());
318        let desc = CallDescriptor::from_call(self.unit, &node, enclosing);
319        self.calls.push(desc);
320        self.visit_children(&node);
321    }
322
323    fn visit_const_item(&mut self, node: HirNode<'tcx>) {
324        let ts_kind = node.inner_ts_node().kind();
325        let symbol_kind = if ts_kind == "static_item" {
326            SymbolKind::Static
327        } else {
328            SymbolKind::Const
329        };
330        if let Some((symbol, ident, fqn)) =
331            self.create_new_symbol(&node, LangRust::field_name, true, symbol_kind)
332        {
333            let variable = match ts_kind {
334                "const_item" => VariableDescriptor::from_const_item(
335                    self.unit,
336                    &node,
337                    ident.name.clone(),
338                    fqn.clone(),
339                ),
340                "static_item" => VariableDescriptor::from_static_item(
341                    self.unit,
342                    &node,
343                    ident.name.clone(),
344                    fqn.clone(),
345                ),
346                _ => return,
347            };
348            self.variables.push(variable);
349            self.visit_children_new_scope(&node, Some(symbol));
350        } else {
351            self.visit_children(&node);
352        }
353    }
354
355    fn visit_static_item(&mut self, node: HirNode<'tcx>) {
356        self.visit_const_item(node);
357    }
358
359    fn visit_struct_item(&mut self, node: HirNode<'tcx>) {
360        let register_globally = self.should_register_globally(&node);
361        if let Some((symbol, _ident, fqn)) = self.create_new_symbol(
362            &node,
363            LangRust::field_name,
364            register_globally,
365            SymbolKind::Struct,
366        ) {
367            if let Some(desc) = StructDescriptor::from_struct(self.unit, &node, fqn.clone()) {
368                self.structs.push(desc);
369            }
370            self.visit_children_new_scope(&node, Some(symbol));
371        } else {
372            eprintln!("Failed to create struct descriptor for: {:?}", node);
373            self.visit_children(&node);
374        }
375    }
376
377    fn visit_enum_item(&mut self, node: HirNode<'tcx>) {
378        let register_globally = self.should_register_globally(&node);
379        if let Some((symbol, _ident, fqn)) = self.create_new_symbol(
380            &node,
381            LangRust::field_name,
382            register_globally,
383            SymbolKind::Enum,
384        ) {
385            if let Some(desc) = EnumDescriptor::from_enum(self.unit, &node, fqn.clone()) {
386                self.enums.push(desc);
387            }
388            self.visit_children_new_scope(&node, Some(symbol));
389        } else {
390            self.visit_children(&node);
391        }
392    }
393
394    fn visit_enum_variant(&mut self, node: HirNode<'tcx>) {
395        let register_globally = self.enum_variant_should_register_globally(&node);
396        let _ = self.create_new_symbol(
397            &node,
398            LangRust::field_name,
399            register_globally,
400            SymbolKind::EnumVariant,
401        );
402        self.visit_children(&node);
403    }
404
405    fn visit_unknown(&mut self, node: HirNode<'tcx>) {
406        self.visit_children(&node);
407    }
408}
409
410fn impl_type_segments<'tcx>(
411    unit: CompileUnit<'tcx>,
412    type_node: &HirNode<'tcx>,
413) -> Option<Vec<String>> {
414    let ts_node = type_node.inner_ts_node();
415    let expr = parse_type_expr(unit, ts_node);
416    extract_path_segments(&expr)
417}
418
419fn extract_path_segments(expr: &TypeExpr) -> Option<Vec<String>> {
420    match expr {
421        TypeExpr::Path { segments, .. } => Some(segments.clone()),
422        TypeExpr::Reference { inner, .. } => extract_path_segments(inner),
423        TypeExpr::Tuple(items) if items.len() == 1 => extract_path_segments(&items[0]),
424        _ => None,
425    }
426}
427
428pub struct CollectionResult {
429    pub functions: Vec<FunctionDescriptor>,
430    pub variables: Vec<VariableDescriptor>,
431    pub calls: Vec<CallDescriptor>,
432    pub structs: Vec<StructDescriptor>,
433    pub enums: Vec<EnumDescriptor>,
434}
435
436pub fn collect_symbols<'tcx>(
437    unit: CompileUnit<'tcx>,
438    globals: &'tcx Scope<'tcx>,
439) -> CollectionResult {
440    let root = unit.file_start_hir_id().unwrap();
441    let node = unit.hir_node(root);
442    let mut decl_finder = DeclCollector::new(unit, globals);
443    decl_finder.visit_node(node);
444    CollectionResult {
445        functions: decl_finder.take_functions(),
446        variables: decl_finder.take_variables(),
447        calls: decl_finder.take_calls(),
448        structs: decl_finder.take_structs(),
449        enums: decl_finder.take_enums(),
450    }
451}
llmcc_rust/collect.rs

llmcc_rust/
collect.rs
