use crate::{FunctionInliner, Replacer};
use leo_ast::{
CallExpression,
Expression,
ExpressionReconstructor,
Identifier,
ReturnStatement,
Statement,
StatementReconstructor,
Type,
UnitExpression,
Variant,
};
use indexmap::IndexMap;
use itertools::Itertools;
impl ExpressionReconstructor for FunctionInliner<'_> {
type AdditionalOutput = Vec<Statement>;
fn reconstruct_call(&mut self, input: CallExpression) -> (Expression, Self::AdditionalOutput) {
if input.program.unwrap() != self.program.unwrap() {
return (Expression::Call(input), Default::default());
}
let function_name = match *input.function {
Expression::Identifier(identifier) => identifier.name,
_ => unreachable!("Parser guarantees that `input.function` is always an identifier."),
};
let (_, callee) = self.reconstructed_functions.iter().find(|(symbol, _)| *symbol == function_name).unwrap();
match callee.variant {
Variant::Inline => {
let parameter_to_argument = callee
.input
.iter()
.map(|input| input.identifier().name)
.zip_eq(input.arguments)
.collect::<IndexMap<_, _>>();
let replace = |identifier: &Identifier| {
parameter_to_argument.get(&identifier.name).cloned().unwrap_or(Expression::Identifier(*identifier))
};
let mut inlined_statements =
Replacer::new(replace).reconstruct_block(callee.block.clone()).0.statements;
let result = match inlined_statements.last() {
Some(Statement::Return(_)) => {
match inlined_statements.pop().unwrap() {
Statement::Return(ReturnStatement { expression, .. }) => expression,
_ => panic!("This branch checks that the last statement is a return statement."),
}
}
_ => {
let id = self.node_builder.next_id();
self.type_table.insert(id, Type::Unit);
Expression::Unit(UnitExpression { span: Default::default(), id })
}
};
(result, inlined_statements)
}
Variant::Function | Variant::AsyncFunction | Variant::Transition | Variant::AsyncTransition => {
(Expression::Call(input), Default::default())
}
}
}
}