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
9const LINEAR_SEARCH_THRESHOLD: usize = 16;
13
14#[newtype_index]
16#[derive(Ord, PartialOrd, get_size2::GetSize)]
17pub struct ScopedSymbolId;
18
19#[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 #[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 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 pub fn is_used(&self) -> bool {
66 self.flags.contains(SymbolFlags::IS_USED)
67 }
68
69 pub const fn is_bound(&self) -> bool {
71 self.flags.contains(SymbolFlags::IS_BOUND)
72 }
73
74 pub fn is_declared(&self) -> bool {
76 self.flags.contains(SymbolFlags::IS_DECLARED)
77 }
78
79 pub fn is_global(&self) -> bool {
81 self.flags.contains(SymbolFlags::MARKED_GLOBAL)
82 }
83
84 pub fn is_nonlocal(&self) -> bool {
86 self.flags.contains(SymbolFlags::MARKED_NONLOCAL)
87 }
88
89 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#[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#[derive(Default, get_size2::GetSize)]
207pub(super) struct SymbolTable {
208 symbols: IndexVec<ScopedSymbolId, Symbol>,
209 reverse: Option<Box<SymbolReverseTable>>,
211}
212
213impl SymbolTable {
214 #[track_caller]
219 pub(crate) fn symbol(&self, id: ScopedSymbolId) -> &Symbol {
220 &self.symbols[id]
221 }
222
223 #[track_caller]
228 pub(crate) fn symbol_mut(&mut self, id: ScopedSymbolId) -> &mut Symbol {
229 &mut self.symbols[id]
230 }
231
232 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 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 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 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}