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 std::{cell::RefCell, collections::HashMap, rc::Rc};
23
24mod symbols;
25pub use symbols::*;
26
27/// Maps global and local symbols to information about them.
28///
29/// Scopes are indexed by the NodeID of the function, block, or iteration.
30#[derive(Debug, Default)]
31pub struct SymbolTable {
32    /// Functions indexed by location.
33    functions: IndexMap<Location, FunctionSymbol>,
34
35    /// Records indexed by location.
36    records: IndexMap<Location, Composite>,
37
38    /// Structs indexed by location.
39    structs: IndexMap<Symbol, Composite>,
40
41    /// Consts that have been successfully evaluated.
42    global_consts: IndexMap<Location, Expression>,
43
44    /// Global variables indexed by location.
45    globals: IndexMap<Location, VariableSymbol>,
46
47    /// Local tables index by the NodeID of the function, iteration, or block they're contained in.
48    all_locals: HashMap<NodeID, LocalTable>,
49
50    /// The current LocalTable we're looking at.
51    local: Option<LocalTable>,
52}
53
54#[derive(Clone, Default, Debug)]
55struct LocalTable {
56    inner: Rc<RefCell<LocalTableInner>>,
57}
58
59#[derive(Clone, Default, Debug)]
60struct LocalTableInner {
61    /// The `NodeID` of the function, iteration, or block this table indexes.
62    id: NodeID,
63
64    /// The parent `NodeID` of this scope, if it exists.
65    parent: Option<NodeID>,
66
67    /// The consts we've evaluated in this scope.
68    consts: HashMap<Symbol, Expression>,
69
70    /// Variables in this scope, indexed by name.
71    variables: HashMap<Symbol, VariableSymbol>,
72}
73
74impl LocalTable {
75    fn new(id: NodeID, parent: Option<NodeID>) -> Self {
76        LocalTable {
77            inner: Rc::new(RefCell::new(LocalTableInner {
78                id,
79                parent,
80                consts: HashMap::new(),
81                variables: HashMap::new(),
82            })),
83        }
84    }
85
86    fn dup(&self, new_id: NodeID) -> Self {
87        let mut inner = self.inner.borrow().clone();
88        inner.id = new_id;
89        LocalTable { inner: Rc::new(RefCell::new(inner)) }
90    }
91}
92
93impl SymbolTable {
94    /// Iterator over all the structs (not records) in this program.
95    pub fn iter_structs(&self) -> impl Iterator<Item = (Symbol, &Composite)> {
96        self.structs.iter().map(|(name, comp)| (*name, comp))
97    }
98
99    /// Iterator over all the records in this program.
100    pub fn iter_records(&self) -> impl Iterator<Item = (Location, &Composite)> {
101        self.records.iter().map(|(loc, comp)| (*loc, comp))
102    }
103
104    /// Iterator over all the functions in this program.
105    pub fn iter_functions(&self) -> impl Iterator<Item = (Location, &FunctionSymbol)> {
106        self.functions.iter().map(|(loc, func_symbol)| (*loc, func_symbol))
107    }
108
109    /// Access the struct by this name if it exists.
110    pub fn lookup_struct(&self, name: Symbol) -> Option<&Composite> {
111        self.structs.get(&name)
112    }
113
114    /// Access the record at this location if it exists.
115    pub fn lookup_record(&self, location: Location) -> Option<&Composite> {
116        self.records.get(&location)
117    }
118
119    /// Access the function at this location if it exists.
120    pub fn lookup_function(&self, location: Location) -> Option<&FunctionSymbol> {
121        self.functions.get(&location)
122    }
123
124    /// Access the variable accessible by this name in the current scope.
125    pub fn lookup_variable(&self, program: Symbol, name: Symbol) -> Option<VariableSymbol> {
126        let mut current = self.local.as_ref();
127
128        while let Some(table) = current {
129            let borrowed = table.inner.borrow();
130            let value = borrowed.variables.get(&name);
131            if value.is_some() {
132                return value.cloned();
133            }
134
135            current = borrowed.parent.and_then(|id| self.all_locals.get(&id));
136        }
137
138        self.globals.get(&Location::new(program, name)).cloned()
139    }
140
141    /// Enter the scope of this `NodeID`, creating a table if it doesn't exist yet.
142    ///
143    /// Passing `None` means to enter the global scope.
144    pub fn enter_scope(&mut self, id: Option<NodeID>) {
145        self.local = id.map(|id| {
146            let parent = self.local.as_ref().map(|table| table.inner.borrow().id);
147            let new_local_table = self.all_locals.entry(id).or_insert_with(|| LocalTable::new(id, parent));
148            assert_eq!(parent, new_local_table.inner.borrow().parent, "Entered scopes out of order.");
149            new_local_table.clone()
150        });
151    }
152
153    /// Enter the new scope with id `new_id`, duplicating its local symbol table from the scope at `old_id`.
154    ///
155    /// This is useful for a pass like loop unrolling, in which the loop body must be duplicated multiple times.
156    pub fn enter_scope_duped(&mut self, new_id: NodeID, old_id: NodeID) {
157        let old_local_table = self.all_locals.get(&old_id).expect("Must have an old scope to dup from.");
158        let new_local_table = old_local_table.dup(new_id);
159        let parent = self.local.as_ref().map(|table| table.inner.borrow().id);
160        new_local_table.inner.borrow_mut().parent = parent;
161        self.all_locals.insert(new_id, new_local_table.clone());
162        self.local = Some(new_local_table);
163    }
164
165    /// Enter the parent scope of the current scope (or the global scope if there is no local parent scope).
166    pub fn enter_parent(&mut self) {
167        let parent: Option<NodeID> = self.local.as_ref().and_then(|table| table.inner.borrow().parent);
168        self.local = parent.map(|id| self.all_locals.get(&id).expect("Parent should exist.")).cloned();
169    }
170
171    /// Insert an evaluated const into the current scope.
172    pub fn insert_const(&mut self, program: Symbol, name: Symbol, value: Expression) {
173        if let Some(table) = self.local.as_mut() {
174            table.inner.borrow_mut().consts.insert(name, value);
175        } else {
176            self.global_consts.insert(Location::new(program, name), value);
177        }
178    }
179
180    /// Find the evaluated const accessible by the given name in the current scope.
181    pub fn lookup_const(&self, program: Symbol, name: Symbol) -> Option<Expression> {
182        let mut current = self.local.as_ref();
183
184        while let Some(table) = current {
185            let borrowed = table.inner.borrow();
186            let value = borrowed.consts.get(&name);
187            if value.is_some() {
188                return value.cloned();
189            }
190
191            current = borrowed.parent.and_then(|id| self.all_locals.get(&id));
192        }
193
194        self.global_consts.get(&Location::new(program, name)).cloned()
195    }
196
197    /// Insert a struct at this name.
198    ///
199    /// Since structs are indexed only by name, the program is used only to check shadowing.
200    pub fn insert_struct(&mut self, program: Symbol, name: Symbol, composite: Composite) -> Result<()> {
201        if let Some(old_composite) = self.structs.get(&name) {
202            if eq_struct(&composite, old_composite) {
203                Ok(())
204            } else {
205                Err(AstError::redefining_external_struct(name, composite.span).into())
206            }
207        } else {
208            let location = Location::new(program, name);
209            self.check_shadow_global(location, composite.span)?;
210            self.structs.insert(name, composite);
211            Ok(())
212        }
213    }
214
215    /// Insert a record at this location.
216    pub fn insert_record(&mut self, location: Location, composite: Composite) -> Result<()> {
217        self.check_shadow_global(location, composite.span)?;
218        self.records.insert(location, composite);
219        Ok(())
220    }
221
222    /// Insert a function at this location.
223    pub fn insert_function(&mut self, location: Location, function: Function) -> Result<()> {
224        self.check_shadow_global(location, function.span)?;
225        self.functions.insert(location, FunctionSymbol { function, finalizer: None });
226        Ok(())
227    }
228
229    /// Insert a global at this location.
230    pub fn insert_global(&mut self, location: Location, var: VariableSymbol) -> Result<()> {
231        self.check_shadow_global(location, var.span)?;
232        self.globals.insert(location, var);
233        Ok(())
234    }
235
236    /// Access the global at this location if it exists.
237    pub fn lookup_global(&self, location: Location) -> Option<&VariableSymbol> {
238        self.globals.get(&location)
239    }
240
241    fn check_shadow_global(&self, location: Location, span: Span) -> Result<()> {
242        if self.functions.contains_key(&location) {
243            Err(AstError::shadowed_function(location.name, span).into())
244        } else if self.records.contains_key(&location) {
245            Err(AstError::shadowed_record(location.name, span).into())
246        } else if self.structs.contains_key(&location.name) {
247            Err(AstError::shadowed_struct(location.name, span).into())
248        } else if self.globals.contains_key(&location) {
249            Err(AstError::shadowed_variable(location.name, span).into())
250        } else {
251            Ok(())
252        }
253    }
254
255    fn check_shadow_variable(&self, program: Symbol, name: Symbol, span: Span) -> Result<()> {
256        let mut current = self.local.as_ref();
257
258        while let Some(table) = current {
259            if table.inner.borrow().variables.contains_key(&name) {
260                return Err(AstError::shadowed_variable(name, span).into());
261            }
262            current = table.inner.borrow().parent.map(|id| self.all_locals.get(&id).expect("Parent should exist."));
263        }
264
265        self.check_shadow_global(Location::new(program, name), span)?;
266
267        Ok(())
268    }
269
270    /// Insert a variable into the current scope.
271    pub fn insert_variable(&mut self, program: Symbol, name: Symbol, var: VariableSymbol) -> Result<()> {
272        self.check_shadow_variable(program, name, var.span)?;
273
274        if let Some(table) = self.local.as_mut() {
275            table.inner.borrow_mut().variables.insert(name, var);
276        } else {
277            self.globals.insert(Location::new(program, name), var);
278        }
279
280        Ok(())
281    }
282
283    /// Attach a finalizer to a function.
284    pub fn attach_finalizer(
285        &mut self,
286        caller: Location,
287        callee: Location,
288        future_inputs: Vec<Location>,
289        inferred_inputs: Vec<Type>,
290    ) -> Result<()> {
291        let callee_location = Location::new(callee.program, callee.name);
292
293        if let Some(func) = self.functions.get_mut(&caller) {
294            func.finalizer = Some(Finalizer { location: callee_location, future_inputs, inferred_inputs });
295            Ok(())
296        } else {
297            Err(AstError::function_not_found(caller.name).into())
298        }
299    }
300}
301
302fn eq_struct(new: &Composite, old: &Composite) -> bool {
303    if new.members.len() != old.members.len() {
304        return false;
305    }
306
307    new.members
308        .iter()
309        .zip(old.members.iter())
310        .all(|(member1, member2)| member1.name() == member2.name() && member1.type_.eq_flat_relaxed(&member2.type_))
311}