Fix flattening pass

This commit is contained in:
d0cd 2022-10-14 22:30:08 -04:00
parent 1f977e5c45
commit c6fd32c032
3 changed files with 20 additions and 57 deletions

View File

@ -16,7 +16,7 @@
use crate::Flattener;
use leo_ast::{Finalize, FinalizeStatement, Function, ProgramReconstructor, Statement, StatementReconstructor, Type};
use leo_ast::{Finalize, Function, ProgramReconstructor, ReturnStatement, Statement, StatementReconstructor, Type};
impl ProgramReconstructor for Flattener<'_> {
/// Flattens a function's body and finalize block, if it exists.
@ -65,43 +65,13 @@ impl ProgramReconstructor for Flattener<'_> {
let mut block = self.reconstruct_block(function.block).0;
// Get all of the guards and return expression.
// TODO: Verify that there is always at least one
let returns = self.clear_early_returns();
// If the function contains return statements, then we fold them into a single return statement.
self.fold_returns(&mut block, returns);
// If the function has a finalize block, then type checking guarantees that it has at least one finalize statement.
if finalize.is_some() {
// Get all of the guards and finalize expression.
let finalize_arguments = self.clear_early_finalizes();
let arguments = match finalize_arguments.iter().all(|component| component.is_empty()) {
// If the finalize statement takes no arguments, then output an empty vector.
true => vec![],
// If the function contains finalize statements with at least one argument, then we fold them into a vector of arguments.
// Note that `finalizes` is always initialized to the appropriate number of vectors.
false => {
// Construct an expression for each argument to the finalize statement.
finalize_arguments
.into_iter()
.enumerate()
.map(|(i, component)| {
let (expression, stmts) = self.fold_guards(format!("fin${i}$").as_str(), component);
// Add all of the accumulated statements to the end of the block.
block.statements.extend(stmts);
expression
})
.collect()
}
};
// Add the `FinalizeStatement` to the end of the block.
block.statements.push(Statement::Finalize(FinalizeStatement {
arguments,
span: Default::default(),
}));
}
Function {
annotations: function.annotations,

View File

@ -19,9 +19,9 @@ use itertools::Itertools;
use std::borrow::Borrow;
use leo_ast::{
AssignStatement, BinaryExpression, BinaryOperation, Block, ConditionalStatement, ConsoleFunction, ConsoleStatement,
DefinitionStatement, Expression, ExpressionReconstructor, FinalizeStatement, Identifier, IterationStatement, Node,
ReturnStatement, Statement, StatementReconstructor, TupleExpression, Type, UnaryExpression, UnaryOperation,
AssignStatement, BinaryExpression, BinaryOperation, Block, ConditionalStatement, DefinitionStatement, Expression,
ExpressionReconstructor, IterationStatement, Node, ReturnStatement, Statement, StatementReconstructor,
UnaryExpression, UnaryOperation,
};
impl StatementReconstructor for Flattener<'_> {
@ -367,23 +367,6 @@ impl StatementReconstructor for Flattener<'_> {
unreachable!("`DefinitionStatement`s should not exist in the AST at this phase of compilation.")
}
/// Replaces a finalize statement with an empty block statement.
/// Stores the arguments to the finalize statement, which are later folded into a single finalize statement at the end of the function.
fn reconstruct_finalize(&mut self, input: FinalizeStatement) -> (Statement, Self::AdditionalOutput) {
// Construct the associated guard.
let guard = self.construct_guard();
// For each finalize argument, add it and its associated guard to the appropriate list of finalize arguments.
// Note that type checking guarantees that the number of arguments in a finalize statement is equal to the number of arguments in to the finalize block.
for (i, argument) in input.arguments.into_iter().enumerate() {
// Note that the argument is not reconstructed.
// Note that this unwrap is safe since we initialize `self.finalizes` with a number of vectors equal to the number of finalize arguments.
self.finalizes.get_mut(i).unwrap().push((guard.clone(), argument));
}
(Statement::dummy(Default::default()), Default::default())
}
// TODO: Error message requesting the user to enable loop-unrolling.
fn reconstruct_iteration(&mut self, _input: IterationStatement) -> (Statement, Self::AdditionalOutput) {
unreachable!("`IterationStatement`s should not be in the AST at this phase of compilation.");
@ -402,12 +385,22 @@ impl StatementReconstructor for Flattener<'_> {
Expression::Identifier(identifier) if self.tuples.contains_key(&identifier.name) => {
// Note that the `unwrap` is safe since the match arm checks that the entry exists in `self.tuples`.
let tuple = self.tuples.get(&identifier.name).unwrap().clone();
self.returns.push((guard, Expression::Tuple(tuple)))
self.returns.push((guard.clone(), Expression::Tuple(tuple)))
}
// Otherwise, add the expression directly.
_ => self.returns.push((guard, input.expression)),
_ => self.returns.push((guard.clone(), input.expression)),
};
// Add each finalize argument to the list of finalize arguments.
if let Some(arguments) = input.finalize_args {
// For each finalize argument, add it and its associated guard to the appropriate list of finalize arguments.
// Note that type checking guarantees that the number of arguments in a finalize statement is equal to the number of arguments in to the finalize block.
for (i, argument) in arguments.into_iter().enumerate() {
// Note that this unwrap is safe since we initialize `self.finalizes` with a number of vectors equal to the number of finalize arguments.
self.finalizes.get_mut(i).unwrap().push((guard.clone(), argument));
}
}
(Statement::dummy(Default::default()), Default::default())
}
}

View File

@ -38,8 +38,8 @@ pub struct Flattener<'a> {
/// Note that returns are inserted in the order they are encountered during a pre-order traversal of the AST.
/// Note that type checking guarantees that there is at most one return in a basic block.
pub(crate) returns: Vec<(Option<Expression>, Expression)>,
/// A list containing tuples of guards and expressions associated with `FinalizeStatement`s.
/// A guard is an expression that evaluates to true on the execution path of the `FinalizeStatement`.
/// A list containing tuples of guards and expressions associated with finalize arguments.
/// A guard is an expression that evaluates to true on the execution path of the finalize argument.
/// Note that finalizes are inserted in the order they are encountered during a pre-order traversal of the AST.
/// Note that type checking guarantees that there is at most one finalize in a basic block.
pub(crate) finalizes: Vec<Vec<(Option<Expression>, Expression)>>,
@ -201,7 +201,7 @@ impl<'a> Flattener<'a> {
if !returns.is_empty() {
let (expression, stmts) = self.fold_guards("ret$", returns);
// TODO: Flatten tuples in the return statements.
// TODO: Flatten tuples in the return statements once they are allowed.
// Add all of the accumulated statements to the end of the block.
block.statements.extend(stmts);