use super::{FlatteningVisitor, ReturnGuard};
use leo_ast::{
AstReconstructor,
Constructor,
Expression,
Function,
ProgramReconstructor,
ProgramScope,
ReturnStatement,
Statement,
};
impl ProgramReconstructor for FlatteningVisitor<'_> {
fn reconstruct_program_scope(&mut self, input: ProgramScope) -> ProgramScope {
self.program = input.program_id.name.name;
ProgramScope {
program_id: input.program_id,
consts: input
.consts
.into_iter()
.map(|(i, c)| match self.reconstruct_const(c) {
(Statement::Const(declaration), _) => (i, declaration),
_ => panic!("`reconstruct_const` can only return `Statement::Const`"),
})
.collect(),
composites: input.composites.into_iter().map(|(i, c)| (i, self.reconstruct_composite(c))).collect(),
mappings: input.mappings.into_iter().map(|(id, mapping)| (id, self.reconstruct_mapping(mapping))).collect(),
storage_variables: input
.storage_variables
.into_iter()
.map(|(id, storage_variable)| (id, self.reconstruct_storage_variable(storage_variable)))
.collect(),
functions: input.functions.into_iter().map(|(i, f)| (i, self.reconstruct_function(f))).collect(),
constructor: input.constructor.map(|c| self.reconstruct_constructor(c)),
span: input.span,
}
}
fn reconstruct_function(&mut self, function: Function) -> Function {
self.is_async = function.variant.is_async_function();
let mut block = self.reconstruct_block(function.block).0;
let returns = std::mem::take(&mut self.returns);
let expression_returns: Vec<(Option<Expression>, ReturnStatement)> = returns
.into_iter()
.map(|(guard, statement)| match guard {
ReturnGuard::None => (None, statement),
ReturnGuard::Unconstructed(plain) | ReturnGuard::Constructed { plain, .. } => {
(Some(leo_ast::Path::from(plain).to_local().into()), statement)
}
})
.collect();
self.fold_returns(&mut block, expression_returns);
Function {
annotations: function.annotations,
variant: function.variant,
identifier: function.identifier,
const_parameters: function.const_parameters,
input: function.input,
output: function.output,
output_type: function.output_type,
block,
span: function.span,
id: function.id,
}
}
fn reconstruct_constructor(&mut self, constructor: Constructor) -> Constructor {
self.is_async = true;
let mut block = self.reconstruct_block(constructor.block).0;
let returns = std::mem::take(&mut self.returns);
let expression_returns: Vec<(Option<Expression>, ReturnStatement)> = returns
.into_iter()
.map(|(guard, statement)| match guard {
ReturnGuard::None => (None, statement),
ReturnGuard::Unconstructed(plain) | ReturnGuard::Constructed { plain, .. } => {
(Some(plain.into()), statement)
}
})
.collect();
self.fold_returns(&mut block, expression_returns);
Constructor { annotations: constructor.annotations, block, span: constructor.span, id: constructor.id }
}
}