Skip to main content

ty_python_core/
symbol.rs

1use bitflags::bitflags;
2use hashbrown::hash_table::Entry;
3use ruff_index::{IndexVec, newtype_index};
4use ruff_python_ast::name::Name;
5use rustc_hash::FxHasher;
6use std::hash::{Hash as _, Hasher as _};
7use std::ops::{Deref, DerefMut};
8
9// Selected using performance and memory profiling across the 162-project ecosystem corpus.
10// Symbol-name equality is cheap enough that raising the cutoff from 8 to 16 reduced retained
11// memory without a measurable performance regression.
12const LINEAR_SEARCH_THRESHOLD: usize = 16;
13
14/// Uniquely identifies a symbol in a given scope.
15#[newtype_index]
16#[derive(Ord, PartialOrd, get_size2::GetSize)]
17pub struct ScopedSymbolId;
18
19/// A symbol in a given scope.
20#[derive(Debug, Clone, PartialEq, Eq, get_size2::GetSize)]
21pub struct Symbol {
22    name: Name,
23    flags: SymbolFlags,
24}
25
26impl std::fmt::Display for Symbol {
27    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
28        self.name.fmt(f)
29    }
30}
31
32bitflags! {
33    /// Flags that can be queried to obtain information about a symbol in a given scope.
34    ///
35    /// See the doc-comment at the top of [`super::use_def`] for explanations of what it
36    /// means for a symbol to be *bound* as opposed to *declared*.
37    #[derive(Copy, Clone, Debug, Eq, PartialEq)]
38    struct SymbolFlags: u8 {
39        const IS_USED               = 1 << 0;
40        const IS_BOUND              = 1 << 1;
41        const IS_DECLARED           = 1 << 2;
42        const MARKED_GLOBAL         = 1 << 3;
43        const MARKED_NONLOCAL       = 1 << 4;
44        /// true if the symbol is assigned more than once, or if it is assigned even though it is already in use
45        const IS_REASSIGNED         = 1 << 5;
46        const IS_PARAMETER          = 1 << 6;
47    }
48}
49
50impl get_size2::GetSize for SymbolFlags {}
51
52impl Symbol {
53    pub(crate) const fn new(name: Name) -> Self {
54        Self {
55            name,
56            flags: SymbolFlags::empty(),
57        }
58    }
59
60    pub fn name(&self) -> &Name {
61        &self.name
62    }
63
64    /// Is the symbol used in its containing scope?
65    pub fn is_used(&self) -> bool {
66        self.flags.contains(SymbolFlags::IS_USED)
67    }
68
69    /// Is the symbol given a value in its containing scope?
70    pub const fn is_bound(&self) -> bool {
71        self.flags.contains(SymbolFlags::IS_BOUND)
72    }
73
74    /// Is the symbol declared in its containing scope?
75    pub fn is_declared(&self) -> bool {
76        self.flags.contains(SymbolFlags::IS_DECLARED)
77    }
78
79    /// Is the symbol `global` its containing scope?
80    pub fn is_global(&self) -> bool {
81        self.flags.contains(SymbolFlags::MARKED_GLOBAL)
82    }
83
84    /// Is the symbol `nonlocal` its containing scope?
85    pub fn is_nonlocal(&self) -> bool {
86        self.flags.contains(SymbolFlags::MARKED_NONLOCAL)
87    }
88
89    /// Is the symbol defined in this scope, vs referring to some enclosing scope?
90    ///
91    /// There are three common cases where a name refers to an enclosing scope:
92    ///
93    /// 1. explicit `global` variables
94    /// 2. explicit `nonlocal` variables
95    /// 3. "free" variables, which are used in a scope where they're neither bound nor declared
96    ///
97    /// Note that even if `is_local` is false, that doesn't necessarily mean there's an enclosing
98    /// scope that resolves the reference. The symbol could be a built-in like `print`, or a name
99    /// error at runtime, or a global variable added dynamically with e.g. `globals()`.
100    ///
101    /// XXX: There's a fourth case that we don't (can't) handle here. A variable that's bound or
102    /// declared (anywhere) in a class body, but used before it's bound (at runtime), resolves
103    /// (unbelievably) to the global scope. For example:
104    /// ```py
105    /// x = 42
106    /// def f():
107    ///     x = 43
108    ///     class Foo:
109    ///         print(x)  # 42 (never 43)
110    ///         if secrets.randbelow(2):
111    ///             x = 44
112    ///         print(x)  # 42 or 44
113    /// ```
114    /// In cases like this, the resolution isn't known until runtime, and in fact it varies from
115    /// one use to the next. The semantic index alone can't resolve this, and instead it's a
116    /// special case in type inference (see `infer_place_load`).
117    pub fn is_local(&self) -> bool {
118        !self.is_global() && !self.is_nonlocal() && (self.is_bound() || self.is_declared())
119    }
120
121    pub const fn is_reassigned(&self) -> bool {
122        self.flags.contains(SymbolFlags::IS_REASSIGNED)
123    }
124
125    pub(crate) fn is_parameter(&self) -> bool {
126        self.flags.contains(SymbolFlags::IS_PARAMETER)
127    }
128
129    pub(super) fn mark_global(&mut self) {
130        self.insert_flags(SymbolFlags::MARKED_GLOBAL);
131    }
132
133    pub(super) fn mark_nonlocal(&mut self) {
134        self.insert_flags(SymbolFlags::MARKED_NONLOCAL);
135    }
136
137    pub(super) fn mark_bound(&mut self) {
138        if self.is_bound() || self.is_used() {
139            self.insert_flags(SymbolFlags::IS_REASSIGNED);
140        }
141
142        self.insert_flags(SymbolFlags::IS_BOUND);
143    }
144
145    pub(super) fn mark_used(&mut self) {
146        self.insert_flags(SymbolFlags::IS_USED);
147    }
148
149    pub(super) fn mark_declared(&mut self) {
150        self.insert_flags(SymbolFlags::IS_DECLARED);
151    }
152
153    pub(super) fn mark_parameter(&mut self) {
154        self.insert_flags(SymbolFlags::IS_PARAMETER);
155    }
156
157    fn insert_flags(&mut self, flags: SymbolFlags) {
158        self.flags.insert(flags);
159    }
160}
161
162/// Map from symbol name to its ID.
163///
164/// Uses a hash table to avoid storing the name twice.
165#[derive(Debug, Default, get_size2::GetSize)]
166struct SymbolReverseTable(hashbrown::HashTable<ScopedSymbolId>);
167
168impl SymbolReverseTable {
169    fn symbol_id(
170        &self,
171        symbols: &IndexVec<ScopedSymbolId, Symbol>,
172        name: &str,
173    ) -> Option<ScopedSymbolId> {
174        self.0
175            .find(Self::hash_name(name), |id| symbols[*id].name == name)
176            .copied()
177    }
178
179    fn entry<'a>(
180        &'a mut self,
181        symbols: &IndexVec<ScopedSymbolId, Symbol>,
182        symbol: &Symbol,
183    ) -> Entry<'a, ScopedSymbolId> {
184        self.0.entry(
185            Self::hash_name(symbol.name()),
186            |id| &symbols[*id].name == symbol.name(),
187            |id| Self::hash_name(&symbols[*id].name),
188        )
189    }
190
191    fn shrink_to_fit(&mut self, symbols: &IndexVec<ScopedSymbolId, Symbol>) {
192        self.0
193            .shrink_to_fit(|id| Self::hash_name(&symbols[*id].name));
194    }
195
196    fn hash_name(name: &str) -> u64 {
197        let mut h = FxHasher::default();
198        name.hash(&mut h);
199        h.finish()
200    }
201}
202
203/// The symbols of a given scope.
204///
205/// Allows lookup by name and a symbol's ID.
206#[derive(Default, get_size2::GetSize)]
207pub(super) struct SymbolTable {
208    symbols: IndexVec<ScopedSymbolId, Symbol>,
209    /// Reverse lookup retained only when linear search would be expensive.
210    reverse: Option<Box<SymbolReverseTable>>,
211}
212
213impl SymbolTable {
214    /// Look up a symbol by its ID.
215    ///
216    /// ## Panics
217    /// If the ID is not valid for this symbol table.
218    #[track_caller]
219    pub(crate) fn symbol(&self, id: ScopedSymbolId) -> &Symbol {
220        &self.symbols[id]
221    }
222
223    /// Look up a symbol by its ID, mutably.
224    ///
225    /// ## Panics
226    /// If the ID is not valid for this symbol table.
227    #[track_caller]
228    pub(crate) fn symbol_mut(&mut self, id: ScopedSymbolId) -> &mut Symbol {
229        &mut self.symbols[id]
230    }
231
232    /// Look up the ID of a symbol by its name.
233    pub(crate) fn symbol_id(&self, name: &str) -> Option<ScopedSymbolId> {
234        if let Some(reverse) = self.reverse.as_deref() {
235            return reverse.symbol_id(&self.symbols, name);
236        }
237
238        self.symbols
239            .iter_enumerated()
240            .find_map(|(id, symbol)| (symbol.name == name).then_some(id))
241    }
242
243    /// Iterate over the symbols in this symbol table.
244    pub(crate) fn iter(&self) -> std::slice::Iter<'_, Symbol> {
245        self.symbols.iter()
246    }
247}
248
249impl PartialEq for SymbolTable {
250    fn eq(&self, other: &Self) -> bool {
251        // It's sufficient to compare the symbols as the map is only a reverse lookup.
252        self.symbols == other.symbols
253    }
254}
255
256impl Eq for SymbolTable {}
257
258impl std::fmt::Debug for SymbolTable {
259    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
260        f.debug_tuple("SymbolTable").field(&self.symbols).finish()
261    }
262}
263
264#[derive(Debug, Default)]
265pub(super) struct SymbolTableBuilder {
266    table: SymbolTable,
267    reverse: SymbolReverseTable,
268}
269
270impl SymbolTableBuilder {
271    pub(super) fn symbol_id(&self, name: &str) -> Option<ScopedSymbolId> {
272        self.reverse.symbol_id(&self.table.symbols, name)
273    }
274
275    /// Add a new symbol to this scope or update the flags if a symbol with the same name already exists.
276    pub(super) fn add(&mut self, symbol: Symbol) -> (ScopedSymbolId, bool) {
277        let entry = self.reverse.entry(&self.table.symbols, &symbol);
278
279        match entry {
280            Entry::Occupied(entry) => {
281                let id = *entry.get();
282
283                if !symbol.flags.is_empty() {
284                    self.symbols[id].flags.insert(symbol.flags);
285                }
286
287                (id, false)
288            }
289            Entry::Vacant(entry) => {
290                let id = self.table.symbols.push(symbol);
291                entry.insert(id);
292                (id, true)
293            }
294        }
295    }
296
297    pub(super) fn build(self) -> SymbolTable {
298        let Self {
299            mut table,
300            mut reverse,
301        } = self;
302        table.symbols.shrink_to_fit();
303
304        if table.symbols.len() > LINEAR_SEARCH_THRESHOLD {
305            reverse.shrink_to_fit(&table.symbols);
306            table.reverse = Some(Box::new(reverse));
307        }
308
309        table
310    }
311}
312
313impl Deref for SymbolTableBuilder {
314    type Target = SymbolTable;
315
316    fn deref(&self) -> &Self::Target {
317        &self.table
318    }
319}
320
321impl DerefMut for SymbolTableBuilder {
322    fn deref_mut(&mut self) -> &mut Self::Target {
323        &mut self.table
324    }
325}