use crate::{CompilerState, Pass, VariableSymbol, VariableType};
use leo_ast::{
AleoProgram,
AstVisitor,
ConstDeclaration,
Location,
Mapping,
Module,
ProgramScope,
ProgramVisitor,
StorageVariable,
Stub,
};
use leo_errors::Result;
use leo_span::Symbol;
use indexmap::IndexSet;
pub struct GlobalVarsCollection;
impl Pass for GlobalVarsCollection {
type Input = ();
type Output = ();
const NAME: &'static str = "GlobalVarsCollection";
fn do_pass(_input: Self::Input, state: &mut CompilerState) -> Result<Self::Output> {
let ast = std::mem::take(&mut state.ast);
let mut visitor = GlobalVarsCollectionVisitor {
state,
program_name: Symbol::intern(""),
module: vec![],
parents: IndexSet::new(),
};
visitor.visit_program(ast.as_repr());
visitor.state.handler.last_err()?;
visitor.state.ast = ast;
Ok(())
}
}
struct GlobalVarsCollectionVisitor<'a> {
state: &'a mut CompilerState,
program_name: Symbol,
module: Vec<Symbol>,
parents: IndexSet<Symbol>,
}
impl GlobalVarsCollectionVisitor<'_> {
pub fn in_module_scope<T>(&mut self, module: &[Symbol], func: impl FnOnce(&mut Self) -> T) -> T {
let parent_module = self.module.clone();
self.module = module.to_vec();
let result = func(self);
self.module = parent_module;
result
}
}
impl AstVisitor for GlobalVarsCollectionVisitor<'_> {
type AdditionalInput = ();
type Output = ();
fn visit_const(&mut self, input: &ConstDeclaration) {
let const_path: Vec<Symbol> = self.module.iter().cloned().chain(std::iter::once(input.place.name)).collect();
if let Err(err) = self.state.symbol_table.insert_variable(self.program_name, &const_path, VariableSymbol {
type_: None,
span: input.place.span,
declaration: VariableType::Const,
}) {
self.state.handler.emit_err(err);
}
}
}
impl ProgramVisitor for GlobalVarsCollectionVisitor<'_> {
fn visit_program_scope(&mut self, input: &ProgramScope) {
self.program_name = input.program_id.name.name;
self.state.symbol_table.add_imported_by(self.program_name, &self.parents);
input.consts.iter().for_each(|(_, c)| self.visit_const(c));
input.mappings.iter().for_each(|(_, c)| self.visit_mapping(c));
input.storage_variables.iter().for_each(|(_, c)| self.visit_storage_variable(c));
}
fn visit_module(&mut self, input: &Module) {
self.program_name = input.program_name;
self.in_module_scope(&input.path.clone(), |slf| {
input.consts.iter().for_each(|(_, c)| slf.visit_const(c));
})
}
fn visit_mapping(&mut self, input: &Mapping) {
if let Err(err) = self.state.symbol_table.insert_global(
Location::new(self.program_name, vec![input.identifier.name]),
VariableSymbol { type_: None, span: input.span, declaration: VariableType::Storage },
) {
self.state.handler.emit_err(err);
}
}
fn visit_storage_variable(&mut self, input: &StorageVariable) {
if let Err(err) = self.state.symbol_table.insert_global(
Location::new(self.program_name, vec![input.identifier.name]),
VariableSymbol { type_: None, span: input.span, declaration: VariableType::Storage },
) {
self.state.handler.emit_err(err);
}
}
fn visit_stub(&mut self, input: &Stub) {
match input {
Stub::FromLeo { program, parents } => {
self.parents = parents.clone();
self.visit_program(program);
}
Stub::FromAleo { program, parents } => {
self.parents = parents.clone();
self.visit_aleo_program(program);
}
}
}
fn visit_aleo_program(&mut self, input: &AleoProgram) {
self.program_name = input.stub_id.name.name;
self.state.symbol_table.add_imported_by(self.program_name, &self.parents);
input.mappings.iter().for_each(|(_, c)| self.visit_mapping(c));
}
}