diff --git a/compiler/passes/src/static_single_assignment/rename_program.rs b/compiler/passes/src/static_single_assignment/rename_program.rs index 21235be7dc..ee9bf0a158 100644 --- a/compiler/passes/src/static_single_assignment/rename_program.rs +++ b/compiler/passes/src/static_single_assignment/rename_program.rs @@ -15,10 +15,11 @@ // along with the Leo library. If not, see . use crate::StaticSingleAssigner; +use itertools::Itertools; use leo_ast::{ Expression, Function, FunctionInput, ProgramReconstructor, ReturnStatement, Statement, StatementReconstructor, - TernaryExpression, + TernaryExpression, TupleExpression, }; impl ProgramReconstructor for StaticSingleAssigner<'_> { @@ -48,20 +49,40 @@ impl ProgramReconstructor for StaticSingleAssigner<'_> { // Type checking guarantees that there exists at least one return statement in the function body. let (_, last_return_expression) = returns.pop().unwrap(); + // TODO: Document handling tuples // Fold all return expressions into a single ternary expression. - let expression = - returns - .into_iter() - .rev() - .fold(last_return_expression, |acc, (guard, expression)| match guard { - None => unreachable!("All return statements except for the last one must have a guard."), - Some(guard) => Expression::Ternary(TernaryExpression { + let expression = returns + .into_iter() + .rev() + .fold(last_return_expression, |acc, (guard, expr)| match guard { + None => unreachable!("All return statements except for the last one must have a guard."), + Some(guard) => match (acc, expr) { + (Expression::Tuple(acc_tuple), Expression::Tuple(expr_tuple)) => { + Expression::Tuple(TupleExpression { + elements: acc_tuple + .elements + .into_iter() + .zip_eq(expr_tuple.elements.into_iter()) + .map(|(if_true, if_false)| { + Expression::Ternary(TernaryExpression { + condition: Box::new(guard.clone()), + if_true: Box::new(if_true), + if_false: Box::new(if_false), + span: Default::default(), + }) + }) + .collect(), + span: Default::default(), + }) + } + (acc, expr) => Expression::Ternary(TernaryExpression { condition: Box::new(guard), - if_true: Box::new(expression), - if_false: Box::new(acc), + if_true: Box::new(acc), + if_false: Box::new(expr), span: Default::default(), }), - }); + }, + }); // Add the `ReturnStatement` to the end of the block. block.statements.push(Statement::Return(ReturnStatement { diff --git a/compiler/passes/src/static_single_assignment/rename_statement.rs b/compiler/passes/src/static_single_assignment/rename_statement.rs index ce29447e5f..05b39259cd 100644 --- a/compiler/passes/src/static_single_assignment/rename_statement.rs +++ b/compiler/passes/src/static_single_assignment/rename_statement.rs @@ -26,19 +26,12 @@ use leo_span::Symbol; use indexmap::IndexSet; impl StatementReconstructor for StaticSingleAssigner<'_> { - /// Transforms a `ReturnStatement` into an `AssignStatement`, - /// storing the variable and the associated guard in `self.early_returns`. + /// Transforms a `ReturnStatement` into an empty `BlockStatement`, + /// storing the expression and the associated guard in `self.early_returns`. /// - /// Note that this pass assumes that there is at most one `ReturnStatement` in a block. + /// Note that type checking guarantees that there is at most one `ReturnStatement` in a block. fn reconstruct_return(&mut self, input: ReturnStatement) -> Statement { - // Create a fresh name for the expression in the return statement. - let symbol = self.unique_symbol("$return"); - self.rename_table.update(symbol, symbol); - - // Initialize a new `AssignStatement` for the return expression. - let place = Expression::Identifier(Identifier::new(symbol)); - - // Add the variable and associated guard. + // Construct the associated guard. let guard = match self.condition_stack.is_empty() { true => None, false => { @@ -53,9 +46,13 @@ impl StatementReconstructor for StaticSingleAssigner<'_> { })) } }; - self.early_returns.push((guard, place.clone())); - Self::simple_assign_statement(place, self.reconstruct_expression(input.expression).0) + // Reconstruct the expression and add it to the early returns. + let expression = self.reconstruct_expression(input.expression).0; + self.early_returns.push((guard, expression)); + + // Return an empty block. + Statement::dummy(input.span) } /// Reconstructs the `DefinitionStatement` into an `AssignStatement`, renaming the left-hand-side as appropriate.