use super::{FlatteningVisitor, ReturnGuard};
use leo_ast::{AstReconstructor, 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,
structs: input.structs.into_iter().map(|(i, c)| (i, self.reconstruct_struct(c))).collect(),
mappings: input.mappings.into_iter().map(|(id, mapping)| (id, self.reconstruct_mapping(mapping))).collect(),
functions: input.functions.into_iter().map(|(i, f)| (i, self.reconstruct_function(f))).collect(),
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(),
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(plain.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,
}
}
}