1use std::collections::HashMap;
2use std::time::{Duration, Instant};
3
4use llmcc_core::context::CompileUnit;
5use llmcc_core::ir::{HirId, HirIdent, HirNode};
6use llmcc_core::symbol::{Scope, Symbol, SymbolKind};
7
8use crate::descriptor::function::parse_type_expr;
9use crate::descriptor::{
10 CallDescriptor, EnumDescriptor, FnVisibility, FunctionDescriptor, StructDescriptor, TypeExpr,
11 VariableDescriptor,
12};
13use crate::token::{AstVisitorRust, LangRust};
14
15#[derive(Debug)]
16pub struct CollectionResult {
17 pub functions: Vec<FunctionDescriptor>,
18 pub variables: Vec<VariableDescriptor>,
19 pub calls: Vec<CallDescriptor>,
20 pub structs: Vec<StructDescriptor>,
21 pub enums: Vec<EnumDescriptor>,
22}
23
24#[derive(Debug)]
25pub struct SymbolSpec {
26 pub owner: HirId,
27 pub name: String,
28 pub fqn: String,
29 pub kind: SymbolKind,
30 pub unit_index: usize,
31}
32
33#[derive(Debug)]
34pub struct ScopeSpec {
35 pub owner: Option<HirId>,
36 pub symbol_index: Option<usize>,
37 pub symbols: Vec<usize>,
38}
39
40#[derive(Debug)]
41pub struct CollectedSymbols {
42 pub result: CollectionResult,
43 pub symbols: Vec<SymbolSpec>,
44 pub scopes: Vec<ScopeSpec>,
45}
46
47#[derive(Debug)]
48pub struct SymbolBatch {
49 pub collected: CollectedSymbols,
50 pub total_time: Duration,
51 pub visit_time: Duration,
52}
53
54#[derive(Debug)]
55struct ScopeInfo {
56 owner: Option<HirId>,
57 symbol_index: Option<usize>,
58 symbols: Vec<usize>,
59 locals: HashMap<String, usize>,
60}
61
62#[derive(Debug)]
63struct DeclCollector<'tcx> {
64 unit: CompileUnit<'tcx>,
65 scope_infos: Vec<ScopeInfo>,
66 scope_lookup: HashMap<HirId, usize>,
67 scope_stack: Vec<usize>,
68 symbols: Vec<SymbolSpec>,
69 functions: Vec<FunctionDescriptor>,
70 variables: Vec<VariableDescriptor>,
71 calls: Vec<CallDescriptor>,
72 structs: Vec<StructDescriptor>,
73 enums: Vec<EnumDescriptor>,
74}
75
76impl<'tcx> DeclCollector<'tcx> {
77 pub fn new(unit: CompileUnit<'tcx>) -> Self {
78 Self {
79 unit,
80 scope_infos: vec![ScopeInfo {
81 owner: None,
82 symbol_index: None,
83 symbols: Vec::new(),
84 locals: HashMap::new(),
85 }],
86 scope_lookup: HashMap::new(),
87 scope_stack: vec![0],
88 symbols: Vec::new(),
89 functions: Vec::new(),
90 variables: Vec::new(),
91 calls: Vec::new(),
92 structs: Vec::new(),
93 enums: Vec::new(),
94 }
95 }
96
97 fn current_scope_index(&self) -> usize {
98 *self
99 .scope_stack
100 .last()
101 .expect("scope stack should never be empty")
102 }
103
104 fn ensure_scope(&mut self, owner: HirId) -> usize {
105 if let Some(&idx) = self.scope_lookup.get(&owner) {
106 return idx;
107 }
108
109 let idx = self.scope_infos.len();
110 self.scope_infos.push(ScopeInfo {
111 owner: Some(owner),
112 symbol_index: None,
113 symbols: Vec::new(),
114 locals: HashMap::new(),
115 });
116 self.scope_lookup.insert(owner, idx);
117 idx
118 }
119
120 fn parent_symbol(&self) -> Option<&SymbolSpec> {
121 for &scope_idx in self.scope_stack.iter().rev() {
122 if let Some(symbol_idx) = self.scope_infos[scope_idx].symbol_index {
123 if let Some(symbol) = self.symbols.get(symbol_idx) {
124 return Some(symbol);
125 }
126 }
127 }
128 None
129 }
130
131 fn scoped_fqn(&self, _node: &HirNode<'tcx>, name: &str) -> String {
132 if let Some(parent) = self.parent_symbol() {
133 if parent.fqn.is_empty() {
134 name.to_string()
135 } else {
136 format!("{}::{}", parent.fqn, name)
137 }
138 } else {
139 name.to_string()
140 }
141 }
142
143 fn create_new_symbol(
144 &mut self,
145 node: &HirNode<'tcx>,
146 field_id: u16,
147 global: bool,
148 kind: SymbolKind,
149 ) -> Option<(usize, &'tcx HirIdent<'tcx>, String)> {
150 let ident_node = node.opt_child_by_field(self.unit, field_id)?;
151 let ident = ident_node.as_ident()?;
152 let name = ident.name.clone();
153 let owner = node.hir_id();
154
155 if let Some(existing_idx) = self.find_symbol_local(&name) {
156 let existing_kind = self.symbols[existing_idx].kind;
157 if Self::different_kind(existing_kind, kind) {
158 let fqn = self.scoped_fqn(node, &name);
159 let idx = self.insert_symbol(owner, name.clone(), fqn.clone(), kind, global);
160 return Some((idx, ident, fqn));
161 } else {
162 let fqn = self.symbols[existing_idx].fqn.clone();
163 return Some((existing_idx, ident, fqn));
164 }
165 }
166
167 let fqn = self.scoped_fqn(node, &name);
168 let idx = self.insert_symbol(owner, name.clone(), fqn.clone(), kind, global);
169 Some((idx, ident, fqn))
170 }
171
172 fn different_kind(existing_kind: SymbolKind, new_kind: SymbolKind) -> bool {
173 existing_kind != SymbolKind::Unknown && existing_kind != new_kind
174 }
175
176 fn insert_symbol(
177 &mut self,
178 owner: HirId,
179 name: String,
180 fqn: String,
181 kind: SymbolKind,
182 global: bool,
183 ) -> usize {
184 let idx = self.symbols.len();
185 self.symbols.push(SymbolSpec {
186 owner,
187 name: name.clone(),
188 fqn,
189 kind,
190 unit_index: self.unit.index,
191 });
192
193 let current_scope = self.current_scope_index();
194 self.scope_infos[current_scope]
195 .locals
196 .insert(name.clone(), idx);
197 self.scope_infos[current_scope].symbols.push(idx);
198
199 if global {
200 self.scope_infos[0].locals.insert(name, idx);
201 self.scope_infos[0].symbols.push(idx);
202 }
203
204 idx
205 }
206
207 fn find_symbol_local(&self, name: &str) -> Option<usize> {
208 if self.scope_stack.len() <= 1 {
209 return None;
210 }
211
212 for &scope_idx in self.scope_stack[1..].iter().rev() {
213 if let Some(&symbol_idx) = self.scope_infos[scope_idx].locals.get(name) {
214 return Some(symbol_idx);
215 }
216 }
217
218 None
219 }
220
221 fn has_public_visibility(&self, node: &HirNode<'tcx>) -> bool {
222 let ts_node = node.inner_ts_node();
223 let Some(name_node) = ts_node.child_by_field_name("name") else {
224 return false;
225 };
226 let header = self
227 .unit
228 .file()
229 .opt_get_text(ts_node.start_byte(), name_node.start_byte())
230 .unwrap_or_default();
231
232 !matches!(FnVisibility::from_header(&header), FnVisibility::Private)
233 }
234
235 fn should_register_globally(&self, node: &HirNode<'tcx>) -> bool {
236 self.parent_symbol().is_none() || self.has_public_visibility(node)
237 }
238
239 fn enum_variant_should_register_globally(&self, _node: &HirNode<'tcx>) -> bool {
240 let Some(enum_symbol) = self.parent_symbol() else {
241 return false;
242 };
243
244 if enum_symbol.kind != SymbolKind::Enum {
245 return false;
246 }
247
248 let parent_node = self.unit.hir_node(enum_symbol.owner);
249 self.has_public_visibility(&parent_node)
250 }
251
252 fn visit_children_new_scope(&mut self, node: &HirNode<'tcx>, scoped_symbol: Option<usize>) {
253 let owner = node.hir_id();
254 let scope_idx = self.ensure_scope(owner);
255 if let Some(symbol_idx) = scoped_symbol {
256 self.scope_infos[scope_idx].symbol_index = Some(symbol_idx);
257 }
258
259 self.scope_stack.push(scope_idx);
260 self.visit_children(node);
261 self.scope_stack.pop();
262 }
263
264 fn visit_children(&mut self, node: &HirNode<'tcx>) {
265 for id in node.children() {
266 let child = self.unit.hir_node(*id);
267 self.visit_node(child);
268 }
269 }
270
271 fn find_symbol_by_fqn_in_unit(&self, fqn: &str) -> Option<usize> {
272 self.symbols
273 .iter()
274 .enumerate()
275 .rev()
276 .find(|(_, symbol)| symbol.fqn == fqn && symbol.unit_index == self.unit.index)
277 .map(|(idx, _)| idx)
278 }
279
280 fn find_symbol_by_fqn(&self, fqn: &str) -> Option<usize> {
281 self.symbols
282 .iter()
283 .enumerate()
284 .rev()
285 .find(|(_, symbol)| symbol.fqn == fqn)
286 .map(|(idx, _)| idx)
287 }
288
289 fn finish(self) -> CollectedSymbols {
290 let scopes = self
291 .scope_infos
292 .into_iter()
293 .map(|info| ScopeSpec {
294 owner: info.owner,
295 symbol_index: info.symbol_index,
296 symbols: info.symbols,
297 })
298 .collect();
299
300 CollectedSymbols {
301 result: CollectionResult {
302 functions: self.functions,
303 variables: self.variables,
304 calls: self.calls,
305 structs: self.structs,
306 enums: self.enums,
307 },
308 symbols: self.symbols,
309 scopes,
310 }
311 }
312}
313
314impl<'tcx> AstVisitorRust<'tcx> for DeclCollector<'tcx> {
315 fn unit(&self) -> CompileUnit<'tcx> {
316 self.unit
317 }
318
319 fn visit_source_file(&mut self, node: HirNode<'tcx>) {
320 self.visit_children_new_scope(&node, None);
321 }
322
323 fn visit_function_item(&mut self, node: HirNode<'tcx>) {
324 let register_globally = self.should_register_globally(&node);
325 if let Some((symbol_idx, _ident, fqn)) = self.create_new_symbol(
326 &node,
327 LangRust::field_name,
328 register_globally,
329 SymbolKind::Function,
330 ) {
331 if let Some(desc) = FunctionDescriptor::from_hir(self.unit, &node, fqn.clone()) {
332 self.functions.push(desc);
333 }
334 self.visit_children_new_scope(&node, Some(symbol_idx));
335 } else {
336 self.visit_children(&node);
337 }
338 }
339
340 fn visit_let_declaration(&mut self, node: HirNode<'tcx>) {
341 if let Some((_symbol_idx, ident, fqn)) =
342 self.create_new_symbol(&node, LangRust::field_pattern, false, SymbolKind::Variable)
343 {
344 let var =
345 VariableDescriptor::from_let(self.unit, &node, ident.name.clone(), fqn.clone());
346 self.variables.push(var);
347 }
348 self.visit_children(&node);
349 }
350
351 fn visit_block(&mut self, node: HirNode<'tcx>) {
352 self.visit_children_new_scope(&node, None);
353 }
354
355 fn visit_parameter(&mut self, node: HirNode<'tcx>) {
356 let _ = self.create_new_symbol(&node, LangRust::field_pattern, false, SymbolKind::Variable);
357 self.visit_children(&node);
358 }
359
360 fn visit_mod_item(&mut self, node: HirNode<'tcx>) {
361 let symbol_idx = self
362 .create_new_symbol(&node, LangRust::field_name, true, SymbolKind::Module)
363 .map(|(symbol_idx, _ident, _)| symbol_idx);
364 self.visit_children_new_scope(&node, symbol_idx);
365 }
366
367 fn visit_impl_item(&mut self, node: HirNode<'tcx>) {
368 let symbol_idx = node
369 .opt_child_by_field(self.unit, LangRust::field_type)
370 .and_then(|type_node| {
371 let segments = impl_type_segments(self.unit, &type_node)?;
372 if segments.is_empty() {
373 return None;
374 }
375
376 let fqn = segments.join("::");
377 let impl_name = segments
378 .last()
379 .cloned()
380 .unwrap_or_else(|| "impl".to_string());
381
382 if let Some(idx) = self.find_symbol_by_fqn_in_unit(&fqn) {
383 let symbol = &mut self.symbols[idx];
384 symbol.owner = node.hir_id();
385 symbol.unit_index = self.unit.index;
386 return Some(idx);
387 }
388
389 if let Some(idx) = self.find_symbol_by_fqn(&fqn) {
390 let symbol = &mut self.symbols[idx];
391 symbol.owner = node.hir_id();
392 symbol.unit_index = self.unit.index;
393 return Some(idx);
394 }
395
396 if let Some(idx) = self.scope_infos[0].locals.get(&impl_name).copied() {
397 let symbol = &mut self.symbols[idx];
398 symbol.owner = node.hir_id();
399 symbol.unit_index = self.unit.index;
400 return Some(idx);
401 }
402
403 Some(self.insert_symbol(node.hir_id(), impl_name, fqn, SymbolKind::Impl, false))
404 });
405 self.visit_children_new_scope(&node, symbol_idx);
406 }
407
408 fn visit_trait_item(&mut self, node: HirNode<'tcx>) {
409 self.visit_struct_item(node);
410 }
411
412 fn visit_function_signature_item(&mut self, node: HirNode<'tcx>) {
413 self.visit_function_item(node);
414 }
415
416 fn visit_call_expression(&mut self, node: HirNode<'tcx>) {
417 let enclosing = self.parent_symbol().map(|symbol| symbol.fqn.clone());
418 let desc = CallDescriptor::from_call(self.unit, &node, enclosing);
419 self.calls.push(desc);
420 self.visit_children(&node);
421 }
422
423 fn visit_const_item(&mut self, node: HirNode<'tcx>) {
424 let ts_kind = node.inner_ts_node().kind();
425 let symbol_kind = if ts_kind == "static_item" {
426 SymbolKind::Static
427 } else {
428 SymbolKind::Const
429 };
430 if let Some((symbol_idx, ident, fqn)) =
431 self.create_new_symbol(&node, LangRust::field_name, true, symbol_kind)
432 {
433 let variable = match ts_kind {
434 "const_item" => VariableDescriptor::from_const_item(
435 self.unit,
436 &node,
437 ident.name.clone(),
438 fqn.clone(),
439 ),
440 "static_item" => VariableDescriptor::from_static_item(
441 self.unit,
442 &node,
443 ident.name.clone(),
444 fqn.clone(),
445 ),
446 _ => return,
447 };
448 self.variables.push(variable);
449 self.visit_children_new_scope(&node, Some(symbol_idx));
450 } else {
451 self.visit_children(&node);
452 }
453 }
454
455 fn visit_static_item(&mut self, node: HirNode<'tcx>) {
456 self.visit_const_item(node);
457 }
458
459 fn visit_struct_item(&mut self, node: HirNode<'tcx>) {
460 let register_globally = self.should_register_globally(&node);
461 if let Some((symbol_idx, _ident, fqn)) = self.create_new_symbol(
462 &node,
463 LangRust::field_name,
464 register_globally,
465 SymbolKind::Struct,
466 ) {
467 if let Some(desc) = StructDescriptor::from_struct(self.unit, &node, fqn.clone()) {
468 self.structs.push(desc);
469 }
470 self.visit_children_new_scope(&node, Some(symbol_idx));
471 } else {
472 eprintln!("Failed to create struct descriptor for: {:?}", node);
473 self.visit_children(&node);
474 }
475 }
476
477 fn visit_enum_item(&mut self, node: HirNode<'tcx>) {
478 let register_globally = self.should_register_globally(&node);
479 if let Some((symbol_idx, _ident, fqn)) = self.create_new_symbol(
480 &node,
481 LangRust::field_name,
482 register_globally,
483 SymbolKind::Enum,
484 ) {
485 if let Some(desc) = EnumDescriptor::from_enum(self.unit, &node, fqn.clone()) {
486 self.enums.push(desc);
487 }
488 self.visit_children_new_scope(&node, Some(symbol_idx));
489 } else {
490 self.visit_children(&node);
491 }
492 }
493
494 fn visit_enum_variant(&mut self, node: HirNode<'tcx>) {
495 let register_globally = self.enum_variant_should_register_globally(&node);
496 let _ = self.create_new_symbol(
497 &node,
498 LangRust::field_name,
499 register_globally,
500 SymbolKind::EnumVariant,
501 );
502 self.visit_children(&node);
503 }
504
505 fn visit_unknown(&mut self, node: HirNode<'tcx>) {
506 self.visit_children(&node);
507 }
508}
509
510fn impl_type_segments<'tcx>(
511 unit: CompileUnit<'tcx>,
512 type_node: &HirNode<'tcx>,
513) -> Option<Vec<String>> {
514 let ts_node = type_node.inner_ts_node();
515 let expr = parse_type_expr(unit, ts_node);
516 extract_path_segments(&expr)
517}
518
519fn extract_path_segments(expr: &TypeExpr) -> Option<Vec<String>> {
520 match expr {
521 TypeExpr::Path { segments, .. } => Some(segments.clone()),
522 TypeExpr::Reference { inner, .. } => extract_path_segments(inner),
523 TypeExpr::Tuple(items) if items.len() == 1 => extract_path_segments(&items[0]),
524 _ => None,
525 }
526}
527
528fn apply_collected_symbols<'tcx>(
529 unit: CompileUnit<'tcx>,
530 globals: &'tcx Scope<'tcx>,
531 collected: &CollectedSymbols,
532) {
533 let interner = unit.interner();
534 let mut created_symbols = Vec::with_capacity(collected.symbols.len());
535
536 {
537 let mut symbol_map = unit.cc.symbol_map.write().unwrap();
538 for spec in &collected.symbols {
539 let key = interner.intern(&spec.name);
540 let symbol = unit
541 .cc
542 .arena
543 .alloc(Symbol::new(spec.owner, spec.name.clone(), key));
544 symbol.set_kind(spec.kind);
545 symbol.set_unit_index(spec.unit_index);
546 symbol.set_fqn(spec.fqn.clone(), interner);
547 symbol_map.insert(symbol.id, symbol);
548 created_symbols.push(symbol);
549 }
550 }
551
552 for scope_spec in &collected.scopes {
553 let target_scope = match scope_spec.owner {
554 Some(owner) => unit.alloc_scope(owner),
555 None => globals,
556 };
557
558 if let Some(symbol_idx) = scope_spec.symbol_index {
559 if let Some(symbol) = created_symbols.get(symbol_idx) {
560 target_scope.set_symbol(Some(symbol));
561 }
562 }
563
564 for &symbol_idx in &scope_spec.symbols {
565 if let Some(symbol) = created_symbols.get(symbol_idx) {
566 target_scope.insert(symbol, interner);
567 }
568 }
569 }
570}
571
572pub fn collect_symbols_batch<'tcx>(unit: CompileUnit<'tcx>) -> SymbolBatch {
573 let collect_start = Instant::now();
574 let root = unit.file_start_hir_id().unwrap();
575 let node = unit.hir_node(root);
576 let mut decl_finder = DeclCollector::new(unit);
577
578 let visit_start = Instant::now();
579 decl_finder.visit_node(node);
580 let visit_time = visit_start.elapsed();
581
582 let collected = decl_finder.finish();
583 let total_time = collect_start.elapsed();
584
585 SymbolBatch {
586 collected,
587 total_time,
588 visit_time,
589 }
590}
591
592pub fn apply_symbol_batch<'tcx>(
593 unit: CompileUnit<'tcx>,
594 globals: &'tcx Scope<'tcx>,
595 batch: SymbolBatch,
596) -> CollectionResult {
597 let SymbolBatch {
598 collected,
599 total_time,
600 visit_time,
601 } = batch;
602
603 let counts = (
604 collected.result.functions.len(),
605 collected.result.structs.len(),
606 collected.result.variables.len(),
607 collected.result.enums.len(),
608 collected.result.calls.len(),
609 );
610
611 apply_collected_symbols(unit, globals, &collected);
612
613 if total_time.as_millis() > 10 {
614 tracing::trace!(
615 "[COLLECT][rust] File {:?}: total={:.2}ms, visit={:.2}ms, fns={}, structs={}, vars={}, enums={}, calls={}",
616 unit.file_path().unwrap_or("unknown"),
617 total_time.as_secs_f64() * 1000.0,
618 visit_time.as_secs_f64() * 1000.0,
619 counts.0,
620 counts.1,
621 counts.2,
622 counts.3,
623 counts.4,
624 );
625 }
626
627 let CollectedSymbols { result, .. } = collected;
628 result
629}
630
631pub fn collect_symbols<'tcx>(
632 unit: CompileUnit<'tcx>,
633 globals: &'tcx Scope<'tcx>,
634) -> CollectionResult {
635 let batch = collect_symbols_batch(unit);
636 apply_symbol_batch(unit, globals, batch)
637}