diff --git a/compiler/passes/src/flattening/flatten_program.rs b/compiler/passes/src/flattening/flatten_program.rs index ca7aa035d9..990977b6c8 100644 --- a/compiler/passes/src/flattening/flatten_program.rs +++ b/compiler/passes/src/flattening/flatten_program.rs @@ -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, diff --git a/compiler/passes/src/flattening/flatten_statement.rs b/compiler/passes/src/flattening/flatten_statement.rs index cdcf5ad142..cc46cd1d23 100644 --- a/compiler/passes/src/flattening/flatten_statement.rs +++ b/compiler/passes/src/flattening/flatten_statement.rs @@ -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()) } } diff --git a/compiler/passes/src/flattening/flattener.rs b/compiler/passes/src/flattening/flattener.rs index af008b045a..9421351b56 100644 --- a/compiler/passes/src/flattening/flattener.rs +++ b/compiler/passes/src/flattening/flattener.rs @@ -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)>, - /// 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, 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);