Allow ConditionalStatements in finalize in Destructuring and Flattening

compiler/passes/src/destructuring/destructure_program.rs

@@ -16,6 +16,37 @@
 
 use crate::Destructurer;
 
-use leo_ast::ProgramReconstructor;
+use leo_ast::{Finalize, Function, ProgramReconstructor, StatementReconstructor};
 
-impl ProgramReconstructor for Destructurer<'_> {}
+impl ProgramReconstructor for Destructurer<'_> {
+    fn reconstruct_function(&mut self, input: Function) -> Function {
+        Function {
+            annotations: input.annotations,
+            variant: input.variant,
+            identifier: input.identifier,
+            input: input.input,
+            output: input.output,
+            output_type: input.output_type,
+            block: self.reconstruct_block(input.block).0,
+            finalize: input.finalize.map(|finalize| {
+                // Set the `is_finalize` flag before reconstructing the finalize block.
+                self.is_finalize = true;
+                // Reconstruct the finalize block.
+                let finalize = Finalize {
+                    identifier: finalize.identifier,
+                    input: finalize.input,
+                    output: finalize.output,
+                    output_type: finalize.output_type,
+                    block: self.reconstruct_block(finalize.block).0,
+                    span: finalize.span,
+                    id: finalize.id,
+                };
+                // Reset the `is_finalize` flag.
+                self.is_finalize = false;
+                finalize
+            }),
+            span: input.span,
+            id: input.id,
+        }
+    }
+}

compiler/passes/src/destructuring/destructure_statement.rs

@@ -220,8 +220,22 @@ impl StatementReconstructor for Destructurer<'_> {
         (Block { span: block.span, statements, id: self.node_builder.next_id() }, Default::default())
     }
 
-    fn reconstruct_conditional(&mut self, _: ConditionalStatement) -> (Statement, Self::AdditionalOutput) {
+    fn reconstruct_conditional(&mut self, input: ConditionalStatement) -> (Statement, Self::AdditionalOutput) {
-        unreachable!("`ConditionalStatement`s should not be in the AST at this phase of compilation.")
+        // Conditional statements can only exist in finalize blocks.
+        if !self.is_finalize {
+            unreachable!("`ConditionalStatement`s should not be in the AST at this phase of compilation.")
+        } else {
+            (
+                Statement::Conditional(ConditionalStatement {
+                    condition: self.reconstruct_expression(input.condition).0,
+                    then: self.reconstruct_block(input.then).0,
+                    otherwise: input.otherwise.map(|n| Box::new(self.reconstruct_statement(*n).0)),
+                    span: input.span,
+                    id: input.id,
+                }),
+                Default::default(),
+            )
+        }
     }
 
     fn reconstruct_console(&mut self, _: ConsoleStatement) -> (Statement, Self::AdditionalOutput) {

compiler/passes/src/destructuring/destructurer.rs

@@ -30,11 +30,13 @@ pub struct Destructurer<'a> {
     pub(crate) assigner: &'a Assigner,
     /// A mapping between variables and flattened tuple expressions.
     pub(crate) tuples: IndexMap<Symbol, TupleExpression>,
+    /// Whether or not we are currently traversing a finalize block.
+    pub(crate) is_finalize: bool,
 }
 
 impl<'a> Destructurer<'a> {
     pub(crate) fn new(type_table: &'a TypeTable, node_builder: &'a NodeBuilder, assigner: &'a Assigner) -> Self {
-        Self { type_table, node_builder, assigner, tuples: IndexMap::new() }
+        Self { type_table, node_builder, assigner, tuples: IndexMap::new(), is_finalize: false }
     }
 
     /// A wrapper around `assigner.simple_assign_statement` that tracks the type of the lhs.

compiler/passes/src/flattening/flatten_program.rs

@@ -16,7 +16,7 @@
 
 use crate::Flattener;
 
-use leo_ast::{Finalize, Function, ProgramReconstructor, ProgramScope, Statement, StatementReconstructor};
+use leo_ast::{Function, ProgramReconstructor, StatementReconstructor};
 
 impl ProgramReconstructor for Flattener<'_> {
     /// Flattens a program scope.
@@ -40,30 +40,8 @@ impl ProgramReconstructor for Flattener<'_> {
     }
 
     /// Flattens a function's body and finalize block, if it exists.
+    /// Note that the finalize block is not flattened since it uses `branch` instructions to produce correct code in for conditional execution.
     fn reconstruct_function(&mut self, function: Function) -> Function {
-        // First, flatten the finalize block. This allows us to initialize self.finalizes correctly.
-        // Note that this is safe since the finalize block is independent of the function body.
-        let finalize = function.finalize.map(|finalize| {
-            // Flatten the finalize block.
-            let mut block = self.reconstruct_block(finalize.block).0;
-
-            // Get all of the guards and return expression.
-            let returns = self.clear_early_returns();
-
-            // Fold the return statements into the block.
-            self.fold_returns(&mut block, returns);
-
-            Finalize {
-                identifier: finalize.identifier,
-                input: finalize.input,
-                output: finalize.output,
-                output_type: finalize.output_type,
-                block,
-                span: finalize.span,
-                id: finalize.id,
-            }
-        });
-
         // Flatten the function body.
         let mut block = self.reconstruct_block(function.block).0;
 
@@ -81,7 +59,7 @@ impl ProgramReconstructor for Flattener<'_> {
             output: function.output,
             output_type: function.output_type,
             block,
-            finalize,
+            finalize: function.finalize,
             span: function.span,
             id: function.id,
         }

compiler/passes/src/flattening/mod.rs

@@ -18,6 +18,7 @@
 //! The pass flattens `ConditionalStatement`s into a sequence of `AssignStatement`s.
 //! The pass rewrites `ReturnStatement`s into `AssignStatement`s and consolidates the returned values as a single `ReturnStatement` at the end of the function.
 //! The pass rewrites ternary expressions over composite data types, into ternary expressions over the individual fields of the composite data type, followed by an expression constructing the composite data type.
+//! Note that this transformation is only applied to non-finalize code.
 //!
 //! Consider the following Leo code, output by the SSA pass.
 //! ```leo

compiler/passes/src/function_inlining/function_inliner.rs

@@ -32,6 +32,8 @@ pub struct FunctionInliner<'a> {
     pub(crate) reconstructed_functions: Vec<(Symbol, Function)>,
     /// The main program.
     pub(crate) program: Option<Symbol>,
+    /// Whether or not we are currently traversing a finalize block.
+    pub(crate) is_finalize: bool,
 }
 
 impl<'a> FunctionInliner<'a> {
@@ -49,6 +51,7 @@ impl<'a> FunctionInliner<'a> {
             reconstructed_functions: Default::default(),
             type_table,
             program: None,
+            is_finalize: false,
         }
     }
 }

compiler/passes/src/function_inlining/inline_program.rs

@@ -16,7 +16,7 @@
 
 use crate::FunctionInliner;
 
-use leo_ast::{Function, ProgramReconstructor, ProgramScope};
+use leo_ast::{Finalize, Function, ProgramReconstructor, ProgramScope, StatementReconstructor};
 use leo_span::Symbol;
 
 use indexmap::IndexMap;
@@ -60,4 +60,35 @@ impl ProgramReconstructor for FunctionInliner<'_> {
             span: input.span,
         }
     }
+
+    fn reconstruct_function(&mut self, input: Function) -> Function {
+        Function {
+            annotations: input.annotations,
+            variant: input.variant,
+            identifier: input.identifier,
+            input: input.input,
+            output: input.output,
+            output_type: input.output_type,
+            block: self.reconstruct_block(input.block).0,
+            finalize: input.finalize.map(|finalize| {
+                // Set the `is_finalize` flag before reconstructing the finalize block.
+                self.is_finalize = true;
+                // Reconstruct the finalize block.
+                let finalize = Finalize {
+                    identifier: finalize.identifier,
+                    input: finalize.input,
+                    output: finalize.output,
+                    output_type: finalize.output_type,
+                    block: self.reconstruct_block(finalize.block).0,
+                    span: finalize.span,
+                    id: finalize.id,
+                };
+                // Reset the `is_finalize` flag.
+                self.is_finalize = false;
+                finalize
+            }),
+            span: input.span,
+            id: input.id,
+        }
+    }
 }

compiler/passes/src/function_inlining/inline_statement.rs

@@ -70,8 +70,21 @@ impl StatementReconstructor for FunctionInliner<'_> {
     }
 
     /// Flattening removes conditional statements from the program.
-    fn reconstruct_conditional(&mut self, _: ConditionalStatement) -> (Statement, Self::AdditionalOutput) {
+    fn reconstruct_conditional(&mut self, input: ConditionalStatement) -> (Statement, Self::AdditionalOutput) {
-        unreachable!("`ConditionalStatement`s should not be in the AST at this phase of compilation.")
+        if !self.is_finalize {
+            unreachable!("`ConditionalStatement`s should not be in the AST at this phase of compilation.")
+        } else {
+            (
+                Statement::Conditional(ConditionalStatement {
+                    condition: self.reconstruct_expression(input.condition).0,
+                    then: self.reconstruct_block(input.then).0,
+                    otherwise: input.otherwise.map(|n| Box::new(self.reconstruct_statement(*n).0)),
+                    span: input.span,
+                    id: input.id,
+                }),
+                Default::default(),
+            )
+        }
     }
 
     /// Parsing guarantees that console statements are not present in the program.