leo_passes/common/symbol_table/
mod.rs

1// Copyright (C) 2019-2025 Provable Inc.
2// This file is part of the Leo library.
3
4// The Leo library is free software: you can redistribute it and/or modify
5// it under the terms of the GNU General Public License as published by
6// the Free Software Foundation, either version 3 of the License, or
7// (at your option) any later version.
8
9// The Leo library is distributed in the hope that it will be useful,
10// but WITHOUT ANY WARRANTY; without even the implied warranty of
11// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
12// GNU General Public License for more details.
13
14// You should have received a copy of the GNU General Public License
15// along with the Leo library. If not, see <https://www.gnu.org/licenses/>.
16
17use leo_ast::{Composite, Expression, Function, Location, NodeID, Type};
18use leo_errors::{AstError, Result};
19use leo_span::{Span, Symbol};
20
21use indexmap::IndexMap;
22use itertools::Itertools;
23use std::{cell::RefCell, collections::HashMap, rc::Rc};
24
25mod symbols;
26pub use symbols::*;
27
28/// Maps global and local symbols to information about them.
29///
30/// Scopes are indexed by the NodeID of the function, block, or iteration.
31#[derive(Debug, Default)]
32pub struct SymbolTable {
33    /// Functions indexed by location.
34    functions: IndexMap<Location, FunctionSymbol>,
35
36    /// Records indexed by location.
37    records: IndexMap<Location, Composite>,
38
39    /// Structs indexed by a path.
40    structs: IndexMap<Vec<Symbol>, Composite>,
41
42    /// Consts that have been successfully evaluated.
43    global_consts: IndexMap<Location, Expression>,
44
45    /// Global variables indexed by location.
46    globals: IndexMap<Location, VariableSymbol>,
47
48    /// Local tables index by the NodeID of the function, iteration, or block they're contained in.
49    all_locals: HashMap<NodeID, LocalTable>,
50
51    /// The current LocalTable we're looking at.
52    local: Option<LocalTable>,
53}
54
55#[derive(Clone, Default, Debug)]
56struct LocalTable {
57    inner: Rc<RefCell<LocalTableInner>>,
58}
59
60#[derive(Clone, Default, Debug)]
61struct LocalTableInner {
62    /// The `NodeID` of the function, iteration, or block this table indexes.
63    id: NodeID,
64
65    /// The parent `NodeID` of this scope, if it exists.
66    parent: Option<NodeID>,
67
68    /// The children of `NodeID` of this scope
69    children: Vec<NodeID>,
70
71    /// The consts we've evaluated in this scope.
72    consts: HashMap<Symbol, Expression>,
73
74    /// Variables in this scope, indexed by name.
75    variables: HashMap<Symbol, VariableSymbol>,
76}
77
78impl LocalTable {
79    fn new(symbol_table: &mut SymbolTable, id: NodeID, parent: Option<NodeID>) -> Self {
80        // If parent exists, register this scope as its child
81        if let Some(parent_id) = parent {
82            if let Some(parent_table) = symbol_table.all_locals.get_mut(&parent_id) {
83                parent_table.inner.borrow_mut().children.push(id);
84            }
85        }
86
87        LocalTable {
88            inner: Rc::new(RefCell::new(LocalTableInner {
89                id,
90                parent,
91                consts: HashMap::new(),
92                variables: HashMap::new(),
93                children: vec![], // Must still initialize our own children list
94            })),
95        }
96    }
97
98    /// Creates a duplicate of this local scope with a new `NodeID`.
99    ///
100    /// TODO: This currently clones the `children` list, which is incorrect. The new scope may incorrectly
101    /// appear to have descendants that still belong to the original scope. This breaks the structure of
102    /// the scope tree and may cause symbol resolution to behave incorrectly.
103    fn dup(&self, new_id: NodeID) -> Self {
104        let mut inner = self.inner.borrow().clone();
105        inner.id = new_id;
106        LocalTable { inner: Rc::new(RefCell::new(inner)) }
107    }
108}
109
110impl SymbolTable {
111    /// Reset everything except leave global consts that have been evaluated.
112    pub fn reset_but_consts(&mut self) {
113        self.functions.clear();
114        self.records.clear();
115        self.structs.clear();
116        self.globals.clear();
117        self.all_locals.clear();
118        self.local = None;
119    }
120
121    /// Are we currently in the global scope?
122    pub fn global_scope(&self) -> bool {
123        self.local.is_none()
124    }
125
126    /// Iterator over all the structs (not records) in this program.
127    pub fn iter_structs(&self) -> impl Iterator<Item = (&Vec<Symbol>, &Composite)> {
128        self.structs.iter()
129    }
130
131    /// Iterator over all the records in this program.
132    pub fn iter_records(&self) -> impl Iterator<Item = (&Location, &Composite)> {
133        self.records.iter()
134    }
135
136    /// Iterator over all the functions in this program.
137    pub fn iter_functions(&self) -> impl Iterator<Item = (&Location, &FunctionSymbol)> {
138        self.functions.iter()
139    }
140
141    /// Access the struct by this name if it exists.
142    pub fn lookup_struct(&self, path: &[Symbol]) -> Option<&Composite> {
143        self.structs.get(path)
144    }
145
146    /// Access the record at this location if it exists.
147    pub fn lookup_record(&self, location: &Location) -> Option<&Composite> {
148        self.records.get(location)
149    }
150
151    /// Access the function at this location if it exists.
152    pub fn lookup_function(&self, location: &Location) -> Option<&FunctionSymbol> {
153        self.functions.get(location)
154    }
155
156    /// Attempts to look up a variable by a path.
157    ///
158    /// First, it tries to resolve the symbol as a global using the full path under the given program.
159    /// If that fails and the path is non-empty, it falls back to resolving the last component
160    /// of the path as a local symbol.
161    ///
162    /// # Arguments
163    ///
164    /// * `program` - The root symbol representing the program or module context.
165    /// * `path` - A slice of symbols representing the absolute path to the variable.
166    ///
167    /// # Returns
168    ///
169    /// An `Option<VariableSymbol>` containing the resolved symbol if found, otherwise `None`.
170    pub fn lookup_path(&self, program: Symbol, path: &[Symbol]) -> Option<VariableSymbol> {
171        self.lookup_global(&Location::new(program, path.to_vec()))
172            .cloned()
173            .or_else(|| path.last().copied().and_then(|name| self.lookup_local(name)))
174    }
175
176    /// Access the variable accessible by this name in the current scope.
177    pub fn lookup_local(&self, name: Symbol) -> Option<VariableSymbol> {
178        let mut current = self.local.as_ref();
179
180        while let Some(table) = current {
181            let borrowed = table.inner.borrow();
182            let value = borrowed.variables.get(&name);
183            if value.is_some() {
184                return value.cloned();
185            }
186
187            current = borrowed.parent.and_then(|id| self.all_locals.get(&id));
188        }
189
190        None
191    }
192
193    /// Enter the scope of this `NodeID`, creating a table if it doesn't exist yet.
194    ///
195    /// Passing `None` means to enter the global scope.
196    pub fn enter_scope(&mut self, id: Option<NodeID>) {
197        self.local = id.map(|id| {
198            let parent = self.local.as_ref().map(|table| table.inner.borrow().id);
199            let new_local_table = if let Some(existing) = self.all_locals.get(&id) {
200                existing.clone()
201            } else {
202                let new_table = LocalTable::new(self, id, parent);
203                self.all_locals.insert(id, new_table.clone());
204                new_table
205            };
206
207            assert_eq!(parent, new_local_table.inner.borrow().parent, "Entered scopes out of order.");
208            new_local_table.clone()
209        });
210    }
211
212    /// Enter the new scope with id `new_id`, duplicating its local symbol table from the scope at `old_id`.
213    ///
214    /// This is useful for a pass like loop unrolling, in which the loop body must be duplicated multiple times.
215    pub fn enter_scope_duped(&mut self, new_id: NodeID, old_id: NodeID) {
216        let old_local_table = self.all_locals.get(&old_id).expect("Must have an old scope to dup from.");
217        let new_local_table = old_local_table.dup(new_id);
218        let parent = self.local.as_ref().map(|table| table.inner.borrow().id);
219        new_local_table.inner.borrow_mut().parent = parent;
220        self.all_locals.insert(new_id, new_local_table.clone());
221        self.local = Some(new_local_table);
222    }
223
224    /// Enter the parent scope of the current scope (or the global scope if there is no local parent scope).
225    pub fn enter_parent(&mut self) {
226        let parent: Option<NodeID> = self.local.as_ref().and_then(|table| table.inner.borrow().parent);
227        self.local = parent.map(|id| self.all_locals.get(&id).expect("Parent should exist.")).cloned();
228    }
229
230    /// Checks if a `symbol` is local to `scope`.
231    pub fn is_local_to(&self, scope: NodeID, symbol: Symbol) -> bool {
232        self.all_locals.get(&scope).map(|locals| locals.inner.borrow().variables.contains_key(&symbol)).unwrap_or(false)
233    }
234
235    /// Checks whether `symbol` is defined in the current scope (self.local) or any of its
236    /// ancestor scopes, up to and including `scope`.
237    ///
238    /// Returns `false` if the current scope is not a descendant of `scope`.
239    pub fn is_defined_in_scope_or_ancestor_until(&self, scope: NodeID, symbol: Symbol) -> bool {
240        let mut current = self.local.as_ref();
241
242        while let Some(table) = current {
243            let inner = table.inner.borrow();
244
245            // Check if symbol is defined in this scope
246            if inner.variables.contains_key(&symbol) {
247                return true;
248            }
249
250            // Stop when we reach the given upper-bound scope
251            if inner.id == scope {
252                break;
253            }
254
255            // Move to parent
256            current = inner.parent.and_then(|parent_id| self.all_locals.get(&parent_id));
257        }
258
259        false
260    }
261
262    /// Checks if a `symbol` is local to `scope` or any of its child scopes.
263    pub fn is_local_to_or_in_child_scope(&self, scope: NodeID, symbol: Symbol) -> bool {
264        let mut stack = vec![scope];
265
266        while let Some(current_id) = stack.pop() {
267            if let Some(table) = self.all_locals.get(&current_id) {
268                let inner = table.inner.borrow();
269
270                if inner.variables.contains_key(&symbol) {
271                    return true;
272                }
273
274                stack.extend(&inner.children);
275            }
276        }
277
278        false
279    }
280
281    /// Insert an evaluated const into the current scope.
282    pub fn insert_const(&mut self, program: Symbol, path: &[Symbol], value: Expression) {
283        if let Some(table) = self.local.as_mut() {
284            let [const_name] = &path else { panic!("Local consts cannot have paths with more than 1 segment.") };
285            table.inner.borrow_mut().consts.insert(*const_name, value);
286        } else {
287            self.global_consts.insert(Location::new(program, path.to_vec()), value);
288        }
289    }
290
291    /// Find the evaluated const accessible by the given name in the current scope.
292    pub fn lookup_const(&self, program: Symbol, path: &[Symbol]) -> Option<Expression> {
293        let mut current = self.local.as_ref();
294
295        while let Some(table) = current {
296            let borrowed = table.inner.borrow();
297            let value = borrowed.consts.get(path.last().expect("all paths must have at least 1 segment"));
298            if value.is_some() {
299                return value.cloned();
300            }
301
302            current = borrowed.parent.and_then(|id| self.all_locals.get(&id));
303        }
304
305        self.global_consts.get(&Location::new(program, path.to_vec())).cloned()
306    }
307
308    /// Insert a struct at this name.
309    ///
310    /// Since structs are indexed only by name, the program is used only to check shadowing.
311    pub fn insert_struct(&mut self, program: Symbol, path: &[Symbol], composite: Composite) -> Result<()> {
312        if let Some(old_composite) = self.structs.get(path) {
313            if eq_struct(&composite, old_composite) {
314                Ok(())
315            } else {
316                Err(AstError::redefining_external_struct(path.iter().format("::"), composite.span).into())
317            }
318        } else {
319            let location = Location::new(program, path.to_vec());
320            self.check_shadow_global(&location, composite.span)?;
321            self.structs.insert(path.to_vec(), composite);
322            Ok(())
323        }
324    }
325
326    /// Insert a record at this location.
327    pub fn insert_record(&mut self, location: Location, composite: Composite) -> Result<()> {
328        self.check_shadow_global(&location, composite.span)?;
329        self.records.insert(location, composite);
330        Ok(())
331    }
332
333    /// Insert a function at this location.
334    pub fn insert_function(&mut self, location: Location, function: Function) -> Result<()> {
335        self.check_shadow_global(&location, function.span)?;
336        self.functions.insert(location, FunctionSymbol { function, finalizer: None });
337        Ok(())
338    }
339
340    /// Insert a global at this location.
341    pub fn insert_global(&mut self, location: Location, var: VariableSymbol) -> Result<()> {
342        self.check_shadow_global(&location, var.span)?;
343        self.globals.insert(location, var);
344        Ok(())
345    }
346
347    /// Access the global at this location if it exists.
348    pub fn lookup_global(&self, location: &Location) -> Option<&VariableSymbol> {
349        self.globals.get(location)
350    }
351
352    fn check_shadow_global(&self, location: &Location, span: Span) -> Result<()> {
353        let display_name = location.path.iter().format("::");
354        if self.functions.contains_key(location) {
355            Err(AstError::shadowed_function(display_name, span).into())
356        } else if self.records.contains_key(location) {
357            Err(AstError::shadowed_record(display_name, span).into())
358        } else if self.structs.contains_key(&location.path) {
359            Err(AstError::shadowed_struct(display_name, span).into())
360        } else if self.globals.contains_key(location) {
361            Err(AstError::shadowed_variable(display_name, span).into())
362        } else {
363            Ok(())
364        }
365    }
366
367    fn check_shadow_variable(&self, program: Symbol, path: &[Symbol], span: Span) -> Result<()> {
368        let mut current = self.local.as_ref();
369
370        while let Some(table) = current {
371            if let [name] = &path {
372                if table.inner.borrow().variables.contains_key(name) {
373                    return Err(AstError::shadowed_variable(name, span).into());
374                }
375            }
376            current = table.inner.borrow().parent.map(|id| self.all_locals.get(&id).expect("Parent should exist."));
377        }
378
379        self.check_shadow_global(&Location::new(program, path.to_vec()), span)?;
380
381        Ok(())
382    }
383
384    /// Insert a variable into the current scope.
385    pub fn insert_variable(&mut self, program: Symbol, path: &[Symbol], var: VariableSymbol) -> Result<()> {
386        self.check_shadow_variable(program, path, var.span)?;
387
388        if let Some(table) = self.local.as_mut() {
389            let [name] = &path else { panic!("Local variables cannot have paths with more than 1 segment.") };
390            table.inner.borrow_mut().variables.insert(*name, var);
391        } else {
392            self.globals.insert(Location::new(program, path.to_vec()), var);
393        }
394
395        Ok(())
396    }
397
398    /// Attach a finalizer to a function.
399    pub fn attach_finalizer(
400        &mut self,
401        caller: Location,
402        callee: Location,
403        future_inputs: Vec<Location>,
404        inferred_inputs: Vec<Type>,
405    ) -> Result<()> {
406        let callee_location = Location::new(callee.program, callee.path);
407
408        if let Some(func) = self.functions.get_mut(&caller) {
409            func.finalizer = Some(Finalizer { location: callee_location, future_inputs, inferred_inputs });
410            Ok(())
411        } else {
412            Err(AstError::function_not_found(caller.path.iter().format("::")).into())
413        }
414    }
415}
416
417fn eq_struct(new: &Composite, old: &Composite) -> bool {
418    if new.members.len() != old.members.len() {
419        return false;
420    }
421
422    new.members
423        .iter()
424        .zip(old.members.iter())
425        .all(|(member1, member2)| member1.name() == member2.name() && member1.type_.eq_flat_relaxed(&member2.type_))
426}