leo_passes/
global_vars_collection.rs1use crate::{CompilerState, Pass, VariableSymbol, VariableType};
44
45use leo_ast::{
46 AleoProgram,
47 AstVisitor,
48 ConstDeclaration,
49 Library,
50 Location,
51 Mapping,
52 Module,
53 ProgramScope,
54 ProgramVisitor,
55 StorageVariable,
56 Stub,
57};
58use leo_errors::Result;
59use leo_span::Symbol;
60
61use indexmap::IndexSet;
62
63pub struct GlobalVarsCollection;
64
65impl Pass for GlobalVarsCollection {
66 type Input = ();
67 type Output = ();
68
69 const NAME: &'static str = "GlobalVarsCollection";
70
71 fn do_pass(_input: Self::Input, state: &mut CompilerState) -> Result<Self::Output> {
72 let ast = std::mem::take(&mut state.ast);
73 let mut visitor = GlobalVarsCollectionVisitor {
74 state,
75 program_name: Symbol::intern(""),
76 module: vec![],
77 parents: IndexSet::new(),
78 };
79
80 ast.visit(
81 |program| visitor.visit_program(program),
82 |_library| {
83 },
85 );
86
87 visitor.state.handler.last_err()?;
88 visitor.state.ast = ast;
89 Ok(())
90 }
91}
92
93struct GlobalVarsCollectionVisitor<'a> {
94 state: &'a mut CompilerState,
96 program_name: Symbol,
98 module: Vec<Symbol>,
100 parents: IndexSet<Symbol>,
102}
103
104impl GlobalVarsCollectionVisitor<'_> {
105 pub fn in_module_scope<T>(&mut self, module: &[Symbol], func: impl FnOnce(&mut Self) -> T) -> T {
107 let parent_module = self.module.clone();
108 self.module = module.to_vec();
109 let result = func(self);
110 self.module = parent_module;
111 result
112 }
113}
114
115impl AstVisitor for GlobalVarsCollectionVisitor<'_> {
116 type AdditionalInput = ();
117 type Output = ();
118
119 fn visit_const(&mut self, input: &ConstDeclaration) {
120 let const_path: Vec<Symbol> = self.module.iter().cloned().chain(std::iter::once(input.place.name)).collect();
123 if let Err(err) = self.state.symbol_table.insert_variable(self.program_name, &const_path, VariableSymbol {
124 type_: None,
125 span: input.place.span,
126 declaration: VariableType::Const,
127 }) {
128 self.state.handler.emit_err(err);
129 }
130 }
131}
132
133impl ProgramVisitor for GlobalVarsCollectionVisitor<'_> {
134 fn visit_program_scope(&mut self, input: &ProgramScope) {
135 self.program_name = input.program_id.as_symbol();
137
138 self.state.symbol_table.add_imported_by(self.program_name, &self.parents);
140
141 input.parents.iter().for_each(|(_, c)| self.visit_type(c));
143 input.consts.iter().for_each(|(_, c)| self.visit_const(c));
144 input.mappings.iter().for_each(|(_, c)| self.visit_mapping(c));
145 input.storage_variables.iter().for_each(|(_, c)| self.visit_storage_variable(c));
146 }
147
148 fn visit_module(&mut self, input: &Module) {
149 self.program_name = input.program_name;
150 self.in_module_scope(&input.path.clone(), |slf| {
151 input.consts.iter().for_each(|(_, c)| slf.visit_const(c));
152 })
153 }
154
155 fn visit_mapping(&mut self, input: &Mapping) {
156 if let Err(err) = self.state.symbol_table.insert_global(
157 Location::new(self.program_name, vec![input.identifier.name]),
158 VariableSymbol { type_: None, span: input.span, declaration: VariableType::Storage },
159 ) {
160 self.state.handler.emit_err(err);
161 }
162 }
163
164 fn visit_storage_variable(&mut self, input: &StorageVariable) {
165 if let Err(err) = self.state.symbol_table.insert_global(
166 Location::new(self.program_name, vec![input.identifier.name]),
167 VariableSymbol { type_: None, span: input.span, declaration: VariableType::Storage },
168 ) {
169 self.state.handler.emit_err(err);
170 }
171 }
172
173 fn visit_stub(&mut self, input: &Stub) {
174 match input {
175 Stub::FromLeo { program, parents } => {
176 self.parents = parents.clone();
177 self.visit_program(program);
178 }
179 Stub::FromAleo { program, parents } => {
180 self.parents = parents.clone();
181 self.visit_aleo_program(program);
182 }
183 Stub::FromLibrary { library, parents } => {
184 self.parents = parents.clone();
185 self.visit_library(library);
186 }
187 }
188 }
189
190 fn visit_library(&mut self, input: &Library) {
191 self.program_name = input.name;
192
193 self.state.symbol_table.add_imported_by(self.program_name, &self.parents);
195
196 self.state.symbol_table.add_as_library(self.program_name);
197
198 input.consts.iter().for_each(|(_, c)| self.visit_const(c));
199 input.modules.values().for_each(|m| self.visit_module(m));
200 }
201
202 fn visit_aleo_program(&mut self, input: &AleoProgram) {
203 self.program_name = input.stub_id.as_symbol();
204
205 self.state.symbol_table.add_imported_by(self.program_name, &self.parents);
207
208 input.mappings.iter().for_each(|(_, c)| self.visit_mapping(c));
209 }
210}