llmcc_core/
symbol.rs

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

llmcc_core/
symbol.rs
