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

llmcc_core/
symbol.rs
