Implement bug fix

compiler/passes/src/static_single_assignment/rename_program.rs

@@ -15,10 +15,11 @@
 // along with the Leo library. If not, see <https://www.gnu.org/licenses/>.
 
 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 {

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.