use crate::{CompilerState, Pass};
use leo_ast::{
AleoProgram,
AstVisitor,
Composite,
ConstDeclaration,
Function,
FunctionStub,
Location,
Mapping,
MappingType,
Module,
OptionalType,
ProgramScope,
ProgramVisitor,
StorageVariable,
Type,
Variant,
};
use leo_errors::Result;
use leo_span::Symbol;
pub struct GlobalItemsCollection;
impl Pass for GlobalItemsCollection {
type Input = ();
type Output = ();
const NAME: &'static str = "GlobalItemsCollection";
fn do_pass(_input: Self::Input, state: &mut CompilerState) -> Result<Self::Output> {
let ast = std::mem::take(&mut state.ast);
let mut visitor = GlobalItemsCollectionVisitor { state, program_name: Symbol::intern(""), module: vec![] };
visitor.visit_program(ast.as_repr());
visitor.state.handler.last_err()?;
visitor.state.ast = ast;
Ok(())
}
}
struct GlobalItemsCollectionVisitor<'a> {
state: &'a mut CompilerState,
program_name: Symbol,
module: Vec<Symbol>,
}
impl GlobalItemsCollectionVisitor<'_> {
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 GlobalItemsCollectionVisitor<'_> {
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();
self.state.symbol_table.set_global_type(&Location::new(self.program_name, const_path), input.type_.clone());
}
}
impl ProgramVisitor for GlobalItemsCollectionVisitor<'_> {
fn visit_program_scope(&mut self, input: &ProgramScope) {
self.program_name = input.program_id.name.name;
input.consts.iter().for_each(|(_, c)| self.visit_const(c));
input.composites.iter().for_each(|(_, c)| self.visit_composite(c));
input.mappings.iter().for_each(|(_, c)| self.visit_mapping(c));
input.storage_variables.iter().for_each(|(_, c)| self.visit_storage_variable(c));
input.functions.iter().for_each(|(_, c)| self.visit_function(c));
if let Some(c) = input.constructor.as_ref() {
self.visit_constructor(c);
}
}
fn visit_module(&mut self, input: &Module) {
self.program_name = input.program_name;
self.in_module_scope(&input.path.clone(), |slf| {
input.composites.iter().for_each(|(_, c)| slf.visit_composite(c));
input.functions.iter().for_each(|(_, c)| slf.visit_function(c));
input.consts.iter().for_each(|(_, c)| slf.visit_const(c));
})
}
fn visit_composite(&mut self, input: &Composite) {
let full_name = self.module.iter().cloned().chain(std::iter::once(input.name())).collect::<Vec<Symbol>>();
if input.is_record {
if let Err(err) =
self.state.symbol_table.insert_record(Location::new(self.program_name, full_name), input.clone())
{
self.state.handler.emit_err(err);
}
} else if let Err(err) =
self.state.symbol_table.insert_struct(Location::new(self.program_name, full_name.clone()), input.clone())
{
self.state.handler.emit_err(err);
}
}
fn visit_mapping(&mut self, input: &Mapping) {
self.state.symbol_table.set_global_type(
&Location::new(self.program_name, vec![input.identifier.name]),
Type::Mapping(MappingType {
key: Box::new(input.key_type.clone()),
value: Box::new(input.value_type.clone()),
}),
);
}
fn visit_storage_variable(&mut self, input: &StorageVariable) {
let type_ = match input.type_ {
Type::Vector(_) => input.type_.clone(),
_ => Type::Optional(OptionalType { inner: Box::new(input.type_.clone()) }),
};
self.state.symbol_table.set_global_type(&Location::new(self.program_name, vec![input.identifier.name]), type_);
}
fn visit_function(&mut self, input: &Function) {
let full_name = self.module.iter().cloned().chain(std::iter::once(input.name())).collect::<Vec<Symbol>>();
if let Err(err) =
self.state.symbol_table.insert_function(Location::new(self.program_name, full_name), input.clone())
{
self.state.handler.emit_err(err);
}
}
fn visit_aleo_program(&mut self, input: &AleoProgram) {
self.program_name = input.stub_id.name.name;
input.functions.iter().for_each(|(_, c)| self.visit_function_stub(c));
input.composites.iter().for_each(|(_, c)| self.visit_composite_stub(c));
input.mappings.iter().for_each(|(_, c)| self.visit_mapping(c));
}
fn visit_function_stub(&mut self, input: &FunctionStub) {
let location = Location::new(self.program_name, vec![input.name()]);
if let Err(err) = self.state.symbol_table.insert_function(location.clone(), Function::from(input.clone())) {
self.state.handler.emit_err(err);
}
if matches!(input.variant, Variant::AsyncTransition) {
let name = Symbol::intern(&format!("finalize/{}", input.name()));
if let Err(err) = self.state.symbol_table.attach_finalizer(
location,
Location::new(self.program_name, vec![name]),
Vec::new(),
Vec::new(),
) {
self.state.handler.emit_err(err);
}
}
}
fn visit_composite_stub(&mut self, input: &Composite) {
if input.is_record {
if let Err(err) = self
.state
.symbol_table
.insert_record(Location::new(self.program_name, vec![input.name()]), input.clone())
{
self.state.handler.emit_err(err);
}
} else if let Err(err) =
self.state.symbol_table.insert_struct(Location::new(self.program_name, vec![input.name()]), input.clone())
{
self.state.handler.emit_err(err);
}
}
}