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.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 Self::different_kind(existing.kind(), kind) => {
101                self.insert_into_scope(owner, ident, global, &fqn, kind)
102            }
103            Some(existing) => {
104                // self.warn_duplicate_symbol(&ident.name, existing.kind(), kind);
105                existing
106            }
107            None => 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.scopes.find_global_suffix(&keys) {
257                    symbol.set_owner(node.hir_id());
258                    symbol.set_unit_index(self.unit.index);
259                    Some(symbol)
260                } else {
261                    let fqn = segments.join("::");
262                    let owner = node.hir_id();
263
264                    // For impl blocks with qualified types like `impl From<T> for module::Type`,
265                    // use the last segment (the actual type name) as the symbol name.
266                    // This ensures proper trie indexing.
267                    let impl_name = segments
268                        .last()
269                        .cloned()
270                        .unwrap_or_else(|| "impl".to_string());
271
272                    // Create a synthetic identifier using the impl target type name
273                    let synthetic_ident = self.unit.cc.arena.alloc(HirIdent::new(
274                        // Create a minimal HirBase using the impl node's information
275                        HirBase {
276                            hir_id: owner,
277                            parent: node.parent(),
278                            node: node.inner_ts_node(),
279                            kind: HirKind::Identifier,
280                            field_id: u16::MAX,
281                            children: Vec::new(),
282                        },
283                        impl_name,
284                    ));
285
286                    let global = false;
287                    let kind = SymbolKind::Impl;
288                    let symbol =
289                        self.scopes
290                            .insert_with(owner, synthetic_ident, global, |symbol| {
291                                symbol.set_owner(owner);
292                                symbol.set_fqn(fqn, self.unit.interner());
293                                symbol.set_kind(kind);
294                                symbol.set_unit_index(self.unit.index);
295                            });
296                    Some(symbol)
297                }
298            });
299        self.visit_children_new_scope(&node, symbol);
300    }
301
302    fn visit_trait_item(&mut self, node: HirNode<'tcx>) {
303        self.visit_children_new_scope(&node, None);
304    }
305
306    fn visit_call_expression(&mut self, node: HirNode<'tcx>) {
307        let enclosing = self
308            .parent_symbol()
309            .map(|symbol| symbol.fqn_name.borrow().clone());
310        let desc = CallDescriptor::from_call(self.unit, &node, enclosing);
311        self.calls.push(desc);
312        self.visit_children(&node);
313    }
314
315    fn visit_const_item(&mut self, node: HirNode<'tcx>) {
316        let kind = node.inner_ts_node().kind();
317        if let Some((symbol, ident, fqn)) =
318            self.create_new_symbol(&node, LangRust::field_name, true, SymbolKind::Const)
319        {
320            let variable = match kind {
321                "const_item" => VariableDescriptor::from_const_item(
322                    self.unit,
323                    &node,
324                    ident.name.clone(),
325                    fqn.clone(),
326                ),
327                "static_item" => VariableDescriptor::from_static_item(
328                    self.unit,
329                    &node,
330                    ident.name.clone(),
331                    fqn.clone(),
332                ),
333                _ => return,
334            };
335            self.variables.push(variable);
336            self.visit_children_new_scope(&node, Some(symbol));
337        } else {
338            self.visit_children(&node);
339        }
340    }
341
342    fn visit_static_item(&mut self, node: HirNode<'tcx>) {
343        self.visit_const_item(node);
344    }
345
346    fn visit_struct_item(&mut self, node: HirNode<'tcx>) {
347        let register_globally = self.should_register_globally(&node);
348        if let Some((symbol, _ident, fqn)) = self.create_new_symbol(
349            &node,
350            LangRust::field_name,
351            register_globally,
352            SymbolKind::Struct,
353        ) {
354            if let Some(desc) = StructDescriptor::from_struct(self.unit, &node, fqn.clone()) {
355                self.structs.push(desc);
356            }
357            self.visit_children_new_scope(&node, Some(symbol));
358        } else {
359            eprintln!("Failed to create struct descriptor for: {:?}", node);
360            self.visit_children(&node);
361        }
362    }
363
364    fn visit_enum_item(&mut self, node: HirNode<'tcx>) {
365        let register_globally = self.should_register_globally(&node);
366        if let Some((symbol, _ident, fqn)) = self.create_new_symbol(
367            &node,
368            LangRust::field_name,
369            register_globally,
370            SymbolKind::Enum,
371        ) {
372            if let Some(desc) = EnumDescriptor::from_enum(self.unit, &node, fqn.clone()) {
373                self.enums.push(desc);
374            }
375            self.visit_children_new_scope(&node, Some(symbol));
376        } else {
377            self.visit_children(&node);
378        }
379    }
380
381    fn visit_enum_variant(&mut self, node: HirNode<'tcx>) {
382        let register_globally = self.enum_variant_should_register_globally(&node);
383        let _ = self.create_new_symbol(
384            &node,
385            LangRust::field_name,
386            register_globally,
387            SymbolKind::EnumVariant,
388        );
389        self.visit_children(&node);
390    }
391
392    fn visit_unknown(&mut self, node: HirNode<'tcx>) {
393        self.visit_children(&node);
394    }
395}
396
397fn impl_type_segments<'tcx>(
398    unit: CompileUnit<'tcx>,
399    type_node: &HirNode<'tcx>,
400) -> Option<Vec<String>> {
401    let ts_node = type_node.inner_ts_node();
402    let expr = parse_type_expr(unit, ts_node);
403    extract_path_segments(&expr)
404}
405
406fn extract_path_segments(expr: &TypeExpr) -> Option<Vec<String>> {
407    match expr {
408        TypeExpr::Path { segments, .. } => Some(segments.clone()),
409        TypeExpr::Reference { inner, .. } => extract_path_segments(inner),
410        TypeExpr::Tuple(items) if items.len() == 1 => extract_path_segments(&items[0]),
411        _ => None,
412    }
413}
414
415pub struct CollectionResult {
416    pub functions: Vec<FunctionDescriptor>,
417    pub variables: Vec<VariableDescriptor>,
418    pub calls: Vec<CallDescriptor>,
419    pub structs: Vec<StructDescriptor>,
420    pub enums: Vec<EnumDescriptor>,
421}
422
423pub fn collect_symbols<'tcx>(
424    unit: CompileUnit<'tcx>,
425    globals: &'tcx Scope<'tcx>,
426) -> CollectionResult {
427    let root = unit.file_start_hir_id().unwrap();
428    let node = unit.hir_node(root);
429    let mut decl_finder = DeclCollector::new(unit, globals);
430    decl_finder.visit_node(node);
431    CollectionResult {
432        functions: decl_finder.take_functions(),
433        variables: decl_finder.take_variables(),
434        calls: decl_finder.take_calls(),
435        structs: decl_finder.take_structs(),
436        enums: decl_finder.take_enums(),
437    }
438}
llmcc_rust/collect.rs

llmcc_rust/
collect.rs
