llmcc_core/
symbol.rs

1use std::cell::{Cell, RefCell};
2use std::collections::HashMap;
3
4use crate::graph_builder::BlockId;
5use crate::interner::{InternPool, InternedStr};
6use crate::ir::{Arena, HirId, HirIdent};
7use crate::trie::SymbolTrie;
8use std::sync::atomic::{AtomicU32, Ordering};
9
10static NEXT_SYMBOL_ID: AtomicU32 = AtomicU32::new(1);
11
12#[derive(Copy, Clone, PartialEq, Eq, Hash, Debug, Default)]
13pub struct SymId(pub u32);
14
15impl std::fmt::Display for SymId {
16    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
17        write!(f, "{}", self.0)
18    }
19}
20
21#[derive(Clone, Copy, Debug, PartialEq, Eq)]
22pub enum SymbolKind {
23    Unknown,
24    Module,
25    Struct,
26    Enum,
27    Function,
28    Variable,
29    Const,
30    Static,
31    Trait,
32    Impl,
33    EnumVariant,
34}
35
36/// Canonical representation of an item bound in a scope (functions, variables, types, etc.).
37#[derive(Debug)]
38pub struct Scope<'tcx> {
39    /// Trie for fast symbol lookup by suffix
40    trie: RefCell<SymbolTrie<'tcx>>,
41    /// The HIR node that owns this scope
42    owner: HirId,
43    /// The symbol that introduced this scope, if any
44    symbol: Cell<Option<&'tcx Symbol>>,
45}
46
47impl<'tcx> Scope<'tcx> {
48    pub fn new(owner: HirId) -> Self {
49        Self {
50            trie: RefCell::new(SymbolTrie::default()),
51            owner,
52            symbol: Cell::new(None),
53        }
54    }
55
56    pub fn owner(&self) -> HirId {
57        self.owner
58    }
59
60    pub fn symbol(&self) -> Option<&'tcx Symbol> {
61        self.symbol.get()
62    }
63
64    pub fn set_symbol(&self, symbol: Option<&'tcx Symbol>) {
65        self.symbol.set(symbol);
66    }
67
68    pub fn insert(&self, symbol: &'tcx Symbol, interner: &InternPool) -> SymId {
69        let sym_id = symbol.id;
70        self.trie.borrow_mut().insert_symbol(symbol, interner);
71        sym_id
72    }
73
74    pub fn get_id(&self, key: InternedStr) -> Option<SymId> {
75        let hits = self.trie.borrow().lookup_symbol_suffix(&[key]);
76        hits.first().map(|symbol| symbol.id)
77    }
78
79    pub fn lookup_suffix_once(&self, suffix: &[InternedStr]) -> Option<&'tcx Symbol> {
80        self.trie
81            .borrow()
82            .lookup_symbol_suffix(suffix)
83            .into_iter()
84            .next()
85    }
86
87    pub fn lookup_suffix_symbols(&self, suffix: &[InternedStr]) -> Vec<&'tcx Symbol> {
88        self.trie.borrow().lookup_symbol_suffix(suffix)
89    }
90
91    pub fn format_compact(&self) -> String {
92        let count = self.trie.borrow().total_symbols();
93        format!("{}/{}", self.owner, count)
94    }
95
96    pub fn all_symbols(&self) -> Vec<&'tcx Symbol> {
97        self.trie.borrow().symbols()
98    }
99}
100
101#[derive(Debug)]
102pub struct ScopeStack<'tcx> {
103    arena: &'tcx Arena<'tcx>,
104    interner: &'tcx InternPool,
105    stack: Vec<&'tcx Scope<'tcx>>,
106    symbol_map: &'tcx RefCell<HashMap<SymId, &'tcx Symbol>>,
107}
108
109impl<'tcx> ScopeStack<'tcx> {
110    pub fn new(
111        arena: &'tcx Arena<'tcx>,
112        interner: &'tcx InternPool,
113        symbol_map: &'tcx RefCell<HashMap<SymId, &'tcx Symbol>>,
114    ) -> Self {
115        Self {
116            arena,
117            interner,
118            stack: Vec::new(),
119            symbol_map,
120        }
121    }
122
123    pub fn depth(&self) -> usize {
124        self.stack.len()
125    }
126
127    pub fn push(&mut self, scope: &'tcx Scope<'tcx>) {
128        self.push_with_symbol(scope, None);
129    }
130
131    pub fn push_with_symbol(&mut self, scope: &'tcx Scope<'tcx>, symbol: Option<&'tcx Symbol>) {
132        scope.set_symbol(symbol);
133        self.stack.push(scope);
134    }
135
136    pub fn pop(&mut self) -> Option<&'tcx Scope<'tcx>> {
137        self.stack.pop()
138    }
139
140    pub fn pop_until(&mut self, depth: usize) {
141        while self.depth() > depth {
142            self.pop();
143        }
144    }
145
146    pub fn top(&self) -> Option<&'tcx Scope<'tcx>> {
147        self.stack.last().copied()
148    }
149
150    pub fn scoped_symbol(&self) -> Option<&'tcx Symbol> {
151        self.stack.iter().rev().find_map(|scope| scope.symbol())
152    }
153
154    pub fn iter(&self) -> impl DoubleEndedIterator<Item = &'tcx Scope<'tcx>> + '_ {
155        self.stack.iter().copied()
156    }
157
158    pub fn lookup_scoped_suffix_once(&self, suffix: &[InternedStr]) -> Option<&'tcx Symbol> {
159        self.find_scoped_suffix_with_filters(suffix, None, None)
160    }
161
162    pub fn find_scoped_suffix_with_filters(
163        &self,
164        suffix: &[InternedStr],
165        kind: Option<SymbolKind>,
166        file: Option<usize>,
167    ) -> Option<&'tcx Symbol> {
168        for scope in self.iter().rev() {
169            let symbols = scope.trie.borrow().lookup_symbol_suffix(suffix);
170            if let Some(symbol) = select_symbol(symbols, kind, file) {
171                return Some(symbol);
172            }
173        }
174        None
175    }
176
177    fn scope_for_insertion(&mut self, global: bool) -> Result<&'tcx Scope<'tcx>, &'static str> {
178        if global {
179            self.stack.first().copied().ok_or("no global scope exists")
180        } else {
181            self.stack
182                .last()
183                .copied()
184                .ok_or("no active scope available")
185        }
186    }
187
188    pub fn insert_symbol(
189        &mut self,
190        symbol: &'tcx Symbol,
191        global: bool,
192    ) -> Result<SymId, &'static str> {
193        let scope = self.scope_for_insertion(global)?;
194        Ok(scope.insert(symbol, self.interner))
195    }
196
197    pub fn find_symbol_id(&self, name: &str) -> Option<SymId> {
198        let key = self.interner.intern(name);
199        self.iter().rev().find_map(|scope| scope.get_id(key))
200    }
201
202    fn find_symbol_local_by_key(&self, key: InternedStr) -> Option<&'tcx Symbol> {
203        let scopes = if self.stack.len() > 1 {
204            &self.stack[1..]
205        } else {
206            &self.stack[..]
207        };
208
209        scopes.iter().rev().find_map(|scope| {
210            scope
211                .trie
212                .borrow()
213                .lookup_symbol_suffix(&[key])
214                .into_iter()
215                .next()
216        })
217    }
218
219    pub fn find_symbol_local(&self, name: &str) -> Option<&'tcx Symbol> {
220        let key = self.interner.intern(name);
221        self.find_symbol_local_by_key(key)
222    }
223
224    pub fn find_global_suffix_vec(&self, suffix: &[InternedStr]) -> Vec<&'tcx Symbol> {
225        self.stack
226            .first()
227            .map(|scope| scope.trie.borrow().lookup_symbol_suffix(suffix))
228            .unwrap_or_default()
229    }
230
231    pub fn find_global_suffix(&self, suffix: &[InternedStr]) -> Option<&'tcx Symbol> {
232        self.find_global_suffix_with_filters(suffix, None, None)
233    }
234
235    pub fn find_global_suffix_with_filters(
236        &self,
237        suffix: &[InternedStr],
238        kind: Option<SymbolKind>,
239        file: Option<usize>,
240    ) -> Option<&'tcx Symbol> {
241        let symbols = self.find_global_suffix_vec(suffix);
242        select_symbol(symbols, kind, file)
243    }
244
245    /// Find a global symbol that matches the provided suffix but restrict results to a unit index.
246    pub fn find_global_suffix_in_unit(
247        &self,
248        suffix: &[InternedStr],
249        unit_index: usize,
250    ) -> Option<&'tcx Symbol> {
251        self.find_global_suffix_with_filters(suffix, None, Some(unit_index))
252    }
253
254    pub fn insert_with<F>(
255        &mut self,
256        owner: HirId,
257        ident: &HirIdent<'tcx>,
258        global: bool,
259        init: F,
260    ) -> &'tcx Symbol
261    where
262        F: FnOnce(&'tcx Symbol),
263    {
264        let key = self.interner.intern(&ident.name);
265
266        let symbol = self.alloc_symbol(owner, ident, key);
267        init(symbol);
268
269        self.insert_symbol(symbol, false)
270            .expect("failed to insert symbol into scope");
271        if global {
272            self.insert_symbol(symbol, true)
273                .expect("failed to insert symbol into global scope");
274        }
275
276        symbol
277    }
278
279    fn alloc_symbol(&self, owner: HirId, ident: &HirIdent<'tcx>, key: InternedStr) -> &'tcx Symbol {
280        let symbol = Symbol::new(owner, ident.name.clone(), key);
281        let symbol = self.arena.alloc(symbol);
282        self.symbol_map.borrow_mut().insert(symbol.id, symbol);
283        symbol
284    }
285}
286
287/// Canonical representation of an item bound in a scope (functions, variables, types, etc.).
288#[derive(Debug, Clone)]
289pub struct Symbol {
290    /// Monotonic identifier assigned when the symbol is created.
291    pub id: SymId,
292    /// Owning HIR node that introduces the symbol (e.g. function, struct, module).
293    pub owner: Cell<HirId>,
294    /// Unqualified name exactly as written in source.
295    pub name: String,
296    /// Interned key for `name`, used for fast lookup.
297    pub name_key: InternedStr,
298    /// Fully qualified name cached as a string (updated as scopes are resolved).
299    pub fqn_name: RefCell<String>,
300    /// Interned key for the fully qualified name.
301    pub fqn_key: RefCell<InternedStr>,
302    /// All symbols that this symbols depends on, most general relation, could be
303    /// another relation, like field_of, type_of, called_by, calls etc.
304    /// we dont do very clear sepration becase we want llm models to do that, we
305    /// only need to tell models some symbols having depends relations
306    pub depends: RefCell<Vec<SymId>>,
307    pub depended: RefCell<Vec<SymId>>,
308    pub kind: Cell<SymbolKind>,
309    /// Which compile unit this symbol defined
310    pub unit_index: Cell<Option<usize>>,
311    /// Optional block id associated with this symbol (for graph building)
312    pub block_id: Cell<Option<BlockId>>,
313}
314
315impl Symbol {
316    pub fn new(owner: HirId, name: String, name_key: InternedStr) -> Self {
317        let id = NEXT_SYMBOL_ID.fetch_add(1, Ordering::SeqCst);
318        let sym_id = SymId(id);
319
320        let fqn_key = name_key;
321
322        Self {
323            id: sym_id,
324            owner: Cell::new(owner),
325            name: name.clone(),
326            name_key,
327            fqn_name: RefCell::new(name),
328            fqn_key: RefCell::new(fqn_key),
329            depends: RefCell::new(Vec::new()),
330            depended: RefCell::new(Vec::new()),
331            kind: Cell::new(SymbolKind::Unknown),
332            unit_index: Cell::new(None),
333            block_id: Cell::new(None),
334        }
335    }
336
337    pub fn owner(&self) -> HirId {
338        self.owner.get()
339    }
340
341    pub fn set_owner(&self, owner: HirId) {
342        self.owner.set(owner);
343    }
344
345    pub fn format_compact(&self) -> String {
346        format!("{}->{} \"{}\"", self.id, self.owner.get(), self.name)
347    }
348
349    pub fn set_fqn(&self, fqn: String, interner: &InternPool) {
350        let key = interner.intern(&fqn);
351        *self.fqn_name.borrow_mut() = fqn;
352        *self.fqn_key.borrow_mut() = key;
353    }
354
355    pub fn kind(&self) -> SymbolKind {
356        self.kind.get()
357    }
358
359    pub fn set_kind(&self, kind: SymbolKind) {
360        self.kind.set(kind);
361    }
362
363    pub fn unit_index(&self) -> Option<usize> {
364        self.unit_index.get()
365    }
366
367    pub fn set_unit_index(&self, file: usize) {
368        if self.unit_index.get().is_none() {
369            self.unit_index.set(Some(file));
370        }
371    }
372
373    pub fn add_depends_on(&self, sym_id: SymId) {
374        if sym_id == self.id {
375            return;
376        }
377        let mut deps = self.depends.borrow_mut();
378        if deps.contains(&sym_id) {
379            return;
380        }
381        deps.push(sym_id);
382    }
383
384    pub fn add_depended_by(&self, sym_id: SymId) {
385        if sym_id == self.id {
386            return;
387        }
388        let mut deps = self.depended.borrow_mut();
389        if deps.contains(&sym_id) {
390            return;
391        }
392        deps.push(sym_id);
393    }
394
395    pub fn add_dependency(&self, other: &Symbol) {
396        self.add_depends_on(other.id);
397        other.add_depended_by(self.id);
398    }
399
400    pub fn block_id(&self) -> Option<BlockId> {
401        self.block_id.get()
402    }
403
404    pub fn set_block_id(&self, block_id: Option<BlockId>) {
405        self.block_id.set(block_id);
406    }
407}
408
409fn select_symbol(
410    candidates: Vec<&Symbol>,
411    kind: Option<SymbolKind>,
412    file: Option<usize>,
413) -> Option<&Symbol> {
414    if candidates.is_empty() {
415        return None;
416    }
417
418    if let Some(kind) = kind {
419        let matches: Vec<&Symbol> = candidates
420            .iter()
421            .copied()
422            .filter(|symbol| symbol.kind() == kind)
423            .collect();
424
425        if let Some(file) = file {
426            if let Some(symbol) = matches
427                .iter()
428                .copied()
429                .find(|symbol| symbol.unit_index() == Some(file))
430            {
431                return Some(symbol);
432            }
433        }
434
435        if !matches.is_empty() {
436            return Some(matches[0]);
437        }
438    }
439
440    if let Some(file) = file {
441        if let Some(symbol) = candidates
442            .iter()
443            .copied()
444            .find(|candidate| candidate.unit_index() == Some(file))
445        {
446            return Some(symbol);
447        }
448    }
449
450    candidates.into_iter().next()
451}
llmcc_core/symbol.rs

llmcc_core/
symbol.rs
