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, NodeBuilder, NodeID, Type};
18use leo_errors::{AstError, Color, Label, LeoError, 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            && let Some(parent_table) = symbol_table.all_locals.get_mut(&parent_id)
83        {
84            parent_table.inner.borrow_mut().children.push(id);
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    fn clear_but_consts(&mut self) {
99        self.inner.borrow_mut().variables.clear();
100    }
101
102    /// Recursively duplicates this table and all children.
103    /// `new_parent` is the NodeID of the parent in the new tree (None for root).
104    pub fn dup(
105        &self,
106        node_builder: &NodeBuilder,
107        all_locals: &mut HashMap<NodeID, LocalTable>,
108        new_parent: Option<NodeID>,
109    ) -> LocalTable {
110        let inner = self.inner.borrow();
111
112        // Generate a new ID for this table
113        let new_id = node_builder.next_id();
114
115        // Recursively duplicate children with new_id as their parent
116        let new_children: Vec<NodeID> = inner
117            .children
118            .iter()
119            .map(|child_id| {
120                let child_table = all_locals.get(child_id).expect("Child must exist").clone();
121                let duped_child = child_table.dup(node_builder, all_locals, Some(new_id));
122                let child_new_id = duped_child.inner.borrow().id;
123                all_locals.insert(child_new_id, duped_child);
124                child_new_id
125            })
126            .collect();
127
128        // Duplicate this table with correct parent
129        let new_table = LocalTable {
130            inner: Rc::new(RefCell::new(LocalTableInner {
131                id: new_id,
132                parent: new_parent,
133                children: new_children,
134                consts: inner.consts.clone(),
135                variables: inner.variables.clone(),
136            })),
137        };
138
139        // Register in all_locals
140        all_locals.insert(new_id, new_table.clone());
141
142        new_table
143    }
144}
145
146impl SymbolTable {
147    /// Reset everything except leave consts that have been evaluated.
148    pub fn reset_but_consts(&mut self) {
149        self.functions.clear();
150        self.records.clear();
151        self.structs.clear();
152        self.globals.clear();
153
154        // clear all non-const locals
155        for local_table in self.all_locals.values_mut() {
156            local_table.clear_but_consts();
157        }
158
159        self.local = None;
160    }
161
162    /// Are we currently in the global scope?
163    pub fn global_scope(&self) -> bool {
164        self.local.is_none()
165    }
166
167    /// Iterator over all the structs (not records) in this program.
168    pub fn iter_structs(&self) -> impl Iterator<Item = (&Vec<Symbol>, &Composite)> {
169        self.structs.iter()
170    }
171
172    /// Iterator over all the records in this program.
173    pub fn iter_records(&self) -> impl Iterator<Item = (&Location, &Composite)> {
174        self.records.iter()
175    }
176
177    /// Iterator over all the functions in this program.
178    pub fn iter_functions(&self) -> impl Iterator<Item = (&Location, &FunctionSymbol)> {
179        self.functions.iter()
180    }
181
182    /// Access the struct by this name if it exists.
183    pub fn lookup_struct(&self, path: &[Symbol]) -> Option<&Composite> {
184        self.structs.get(path)
185    }
186
187    /// Access the record at this location if it exists.
188    pub fn lookup_record(&self, location: &Location) -> Option<&Composite> {
189        self.records.get(location)
190    }
191
192    /// Access the function at this location if it exists.
193    pub fn lookup_function(&self, location: &Location) -> Option<&FunctionSymbol> {
194        self.functions.get(location)
195    }
196
197    /// Attempts to look up a variable by a path.
198    ///
199    /// First, it tries to resolve the symbol as a global using the full path under the given program.
200    /// If that fails and the path is non-empty, it falls back to resolving the last component
201    /// of the path as a local symbol.
202    ///
203    /// # Arguments
204    ///
205    /// * `program` - The root symbol representing the program or module context.
206    /// * `path` - A slice of symbols representing the absolute path to the variable.
207    ///
208    /// # Returns
209    ///
210    /// An `Option<VariableSymbol>` containing the resolved symbol if found, otherwise `None`.
211    pub fn lookup_path(&self, program: Symbol, path: &[Symbol]) -> Option<VariableSymbol> {
212        self.lookup_global(&Location::new(program, path.to_vec()))
213            .cloned()
214            .or_else(|| path.last().copied().and_then(|name| self.lookup_local(name)))
215    }
216
217    /// Access the variable accessible by this name in the current scope.
218    pub fn lookup_local(&self, name: Symbol) -> Option<VariableSymbol> {
219        let mut current = self.local.as_ref();
220
221        while let Some(table) = current {
222            let borrowed = table.inner.borrow();
223            let value = borrowed.variables.get(&name);
224            if value.is_some() {
225                return value.cloned();
226            }
227
228            current = borrowed.parent.and_then(|id| self.all_locals.get(&id));
229        }
230
231        None
232    }
233
234    /// Enter the scope of this `NodeID`, creating a table if it doesn't exist yet.
235    ///
236    /// Passing `None` means to enter the global scope.
237    pub fn enter_scope(&mut self, id: Option<NodeID>) {
238        self.local = id.map(|id| {
239            let parent = self.local.as_ref().map(|table| table.inner.borrow().id);
240            let new_local_table = if let Some(existing) = self.all_locals.get(&id) {
241                existing.clone()
242            } else {
243                let new_table = LocalTable::new(self, id, parent);
244                self.all_locals.insert(id, new_table.clone());
245                new_table
246            };
247
248            assert_eq!(parent, new_local_table.inner.borrow().parent, "Entered scopes out of order.");
249            new_local_table.clone()
250        });
251    }
252
253    /// Enter a new scope by duplicating the local table at `old_id` recursively.
254    ///
255    /// Each scope in the subtree receives a fresh NodeID from `node_builder`.
256    /// The new root scope's parent is set to the current scope (if any).
257    /// Returns the NodeID of the new duplicated root scope.
258    pub fn enter_scope_duped(&mut self, old_id: NodeID, node_builder: &NodeBuilder) -> usize {
259        let old_local_table = self.all_locals.get(&old_id).expect("Must have an old scope to dup from.").clone();
260
261        // Recursively duplicate the table and all its children
262        let new_local_table =
263            old_local_table.dup(node_builder, &mut self.all_locals, self.local.as_ref().map(|t| t.inner.borrow().id));
264
265        let new_id = new_local_table.inner.borrow().id;
266
267        // Update current scope
268        self.local = Some(new_local_table);
269
270        new_id
271    }
272
273    /// Enter the parent scope of the current scope (or the global scope if there is no local parent scope).
274    pub fn enter_parent(&mut self) {
275        let parent: Option<NodeID> = self.local.as_ref().and_then(|table| table.inner.borrow().parent);
276        self.local = parent.map(|id| self.all_locals.get(&id).expect("Parent should exist.")).cloned();
277    }
278
279    /// Checks if a `symbol` is local to `scope`.
280    pub fn is_local_to(&self, scope: NodeID, symbol: Symbol) -> bool {
281        self.all_locals.get(&scope).map(|locals| locals.inner.borrow().variables.contains_key(&symbol)).unwrap_or(false)
282    }
283
284    /// Checks whether `symbol` is defined in the current scope (self.local) or any of its
285    /// ancestor scopes, up to and including `scope`.
286    ///
287    /// Returns `false` if the current scope is not a descendant of `scope`.
288    pub fn is_defined_in_scope_or_ancestor_until(&self, scope: NodeID, symbol: Symbol) -> bool {
289        let mut current = self.local.as_ref();
290
291        while let Some(table) = current {
292            let inner = table.inner.borrow();
293
294            // Check if symbol is defined in this scope
295            if inner.variables.contains_key(&symbol) {
296                return true;
297            }
298
299            // Stop when we reach the given upper-bound scope
300            if inner.id == scope {
301                break;
302            }
303
304            // Move to parent
305            current = inner.parent.and_then(|parent_id| self.all_locals.get(&parent_id));
306        }
307
308        false
309    }
310
311    /// Checks if a `symbol` is local to `scope` or any of its child scopes.
312    pub fn is_local_to_or_in_child_scope(&self, scope: NodeID, symbol: Symbol) -> bool {
313        let mut stack = vec![scope];
314
315        while let Some(current_id) = stack.pop() {
316            if let Some(table) = self.all_locals.get(&current_id) {
317                let inner = table.inner.borrow();
318
319                if inner.variables.contains_key(&symbol) {
320                    return true;
321                }
322
323                stack.extend(&inner.children);
324            }
325        }
326
327        false
328    }
329
330    /// Insert an evaluated const into the current scope.
331    pub fn insert_const(&mut self, program: Symbol, path: &[Symbol], value: Expression) {
332        if let Some(table) = self.local.as_mut() {
333            let [const_name] = &path else { panic!("Local consts cannot have paths with more than 1 segment.") };
334            table.inner.borrow_mut().consts.insert(*const_name, value);
335        } else {
336            self.global_consts.insert(Location::new(program, path.to_vec()), value);
337        }
338    }
339
340    /// Find the evaluated const accessible by the given name in the current scope.
341    pub fn lookup_const(&self, program: Symbol, path: &[Symbol]) -> Option<Expression> {
342        let mut current = self.local.as_ref();
343
344        while let Some(table) = current {
345            let borrowed = table.inner.borrow();
346            let value = borrowed.consts.get(path.last().expect("all paths must have at least 1 segment"));
347            if value.is_some() {
348                return value.cloned();
349            }
350
351            current = borrowed.parent.and_then(|id| self.all_locals.get(&id));
352        }
353
354        self.global_consts.get(&Location::new(program, path.to_vec())).cloned()
355    }
356
357    /// Insert a struct at this name.
358    ///
359    /// Since structs are indexed only by name, the program is used only to check shadowing.
360    pub fn insert_struct(&mut self, program: Symbol, path: &[Symbol], composite: Composite) -> Result<()> {
361        if let Some(old_composite) = self.structs.get(path) {
362            if eq_struct(&composite, old_composite) {
363                Ok(())
364            } else {
365                Err(AstError::redefining_external_struct(path.iter().format("::"), old_composite.span).into())
366            }
367        } else {
368            let location = Location::new(program, path.to_vec());
369            self.check_shadow_global(&location, composite.identifier.span)?;
370            self.structs.insert(path.to_vec(), composite);
371            Ok(())
372        }
373    }
374
375    /// Insert a record at this location.
376    pub fn insert_record(&mut self, location: Location, composite: Composite) -> Result<()> {
377        self.check_shadow_global(&location, composite.identifier.span)?;
378        self.records.insert(location, composite);
379        Ok(())
380    }
381
382    /// Insert a function at this location.
383    pub fn insert_function(&mut self, location: Location, function: Function) -> Result<()> {
384        self.check_shadow_global(&location, function.identifier.span)?;
385        self.functions.insert(location, FunctionSymbol { function, finalizer: None });
386        Ok(())
387    }
388
389    /// Insert a global at this location.
390    pub fn insert_global(&mut self, location: Location, var: VariableSymbol) -> Result<()> {
391        self.check_shadow_global(&location, var.span)?;
392        self.globals.insert(location, var);
393        Ok(())
394    }
395
396    /// Access the global at this location if it exists.
397    pub fn lookup_global(&self, location: &Location) -> Option<&VariableSymbol> {
398        self.globals.get(location)
399    }
400
401    pub fn emit_shadow_error(name: Symbol, span: Span, prev_span: Span) -> LeoError {
402        AstError::name_defined_multiple_times(name, span)
403            .with_labels(vec![
404                Label::new(format!("previous definition of `{name}` here"), prev_span).with_color(Color::Blue),
405                Label::new(format!("`{name}` redefined here"), span),
406            ])
407            .into()
408    }
409
410    fn check_shadow_global(&self, location: &Location, span: Span) -> Result<()> {
411        let name = location.path.last().expect("location path must have at least one segment.");
412
413        self.functions
414            .get(location)
415            .map(|f| f.function.identifier.span)
416            .or_else(|| self.records.get(location).map(|r| r.identifier.span))
417            .or_else(|| self.structs.get(&location.path).map(|s| s.identifier.span))
418            .or_else(|| self.globals.get(location).map(|g| g.span))
419            .map_or_else(|| Ok(()), |prev_span| Err(Self::emit_shadow_error(*name, span, prev_span)))
420    }
421
422    fn check_shadow_variable(&self, program: Symbol, path: &[Symbol], span: Span) -> Result<()> {
423        let mut current = self.local.as_ref();
424
425        while let Some(table) = current {
426            if let [name] = &path
427                && let Some(var_symbol) = table.inner.borrow().variables.get(name)
428            {
429                return Err(Self::emit_shadow_error(*name, span, var_symbol.span));
430            }
431            current = table.inner.borrow().parent.map(|id| self.all_locals.get(&id).expect("Parent should exist."));
432        }
433
434        self.check_shadow_global(&Location::new(program, path.to_vec()), span)?;
435
436        Ok(())
437    }
438
439    /// Insert a variable into the current scope.
440    pub fn insert_variable(&mut self, program: Symbol, path: &[Symbol], var: VariableSymbol) -> Result<()> {
441        self.check_shadow_variable(program, path, var.span)?;
442
443        if let Some(table) = self.local.as_mut() {
444            let [name] = &path else { panic!("Local variables cannot have paths with more than 1 segment.") };
445            table.inner.borrow_mut().variables.insert(*name, var);
446        } else {
447            self.globals.insert(Location::new(program, path.to_vec()), var);
448        }
449
450        Ok(())
451    }
452
453    /// Attach a finalizer to a function.
454    pub fn attach_finalizer(
455        &mut self,
456        caller: Location,
457        callee: Location,
458        future_inputs: Vec<Location>,
459        inferred_inputs: Vec<Type>,
460    ) -> Result<()> {
461        let callee_location = Location::new(callee.program, callee.path);
462
463        if let Some(func) = self.functions.get_mut(&caller) {
464            func.finalizer = Some(Finalizer { location: callee_location, future_inputs, inferred_inputs });
465            Ok(())
466        } else {
467            Err(AstError::function_not_found(caller.path.iter().format("::")).into())
468        }
469    }
470}
471
472fn eq_struct(new: &Composite, old: &Composite) -> bool {
473    if new.members.len() != old.members.len() {
474        return false;
475    }
476
477    new.members
478        .iter()
479        .zip(old.members.iter())
480        .all(|(member1, member2)| member1.name() == member2.name() && member1.type_.eq_flat_relaxed(&member2.type_))
481}
leo_passes/common/symbol_table/mod.rs

leo_passes/common/symbol_table/
mod.rs
