Merge pull request #27903 from AleoHQ/feat/branch-in-finalize

[Feature] Conditional execution in finalize.

compiler/passes/src/code_generation/generator.rs

@@ -53,6 +53,10 @@ pub struct CodeGenerator<'a> {
     pub(crate) program: &'a Program,
     // The program ID of the current program.
     pub(crate) program_id: Option<ProgramId>,
+    /// A counter to track the next available label.
+    pub(crate) next_label: u64,
+    /// The depth of the current conditional block.
+    pub(crate) conditional_depth: u64,
 }
 
 impl<'a> CodeGenerator<'a> {
@@ -80,6 +84,8 @@ impl<'a> CodeGenerator<'a> {
             futures: Vec::new(),
             program,
             program_id: None,
+            next_label: 0u64,
+            conditional_depth: 0u64,
         }
     }
 }

compiler/passes/src/code_generation/visit_statements.rs

@@ -220,12 +220,59 @@ impl<'a> CodeGenerator<'a> {
     }
 
     fn visit_conditional(&mut self, _input: &'a ConditionalStatement) -> String {
-        // TODO: Once SSA is made optional, create a Leo error informing the user to enable the SSA pass.
-        unreachable!("`ConditionalStatement`s should not be in the AST at this phase of compilation.")
+        if !self.in_finalize {
+            unreachable!("`ConditionalStatement`s should not be in the AST at this phase of compilation.")
+        } else {
+            // Construct a label for the end of the `then` block.
+            let end_then_label = format!("end_then_{}_{}", self.conditional_depth, self.next_label);
+            self.next_label += 1;
+            // Construct a label for the end of the `otherwise` block if it exists.
+            let (has_otherwise, end_otherwise_label) = {
+                match _input.otherwise.is_some() {
+                    true => {
+                        // Construct a label for the end of the `otherwise` block.
+                        let end_otherwise_label =
+                            { format!("end_otherwise_{}_{}", self.conditional_depth, self.next_label) };
+                        self.next_label += 1;
+                        (true, end_otherwise_label)
+                    }
+                    false => (false, String::new()),
+                }
+            };
+
+            // Increment the conditional depth.
+            self.conditional_depth += 1;
+
+            // Create a `branch` instruction.
+            let (condition, mut instructions) = self.visit_expression(&_input.condition);
+            instructions.push_str(&format!("    branch.eq {condition} false to {end_then_label};\n"));
+
+            // Visit the `then` block.
+            instructions.push_str(&self.visit_block(&_input.then));
+            // If the `otherwise` block is present, add a branch instruction to jump to the end of the `otherwise` block.
+            if has_otherwise {
+                instructions.push_str(&format!("    branch.eq true true to {end_otherwise_label};\n"));
+            }
+
+            // Add a label for the end of the `then` block.
+            instructions.push_str(&format!("    position {};\n", end_then_label));
+
+            // Visit the `otherwise` block.
+            if let Some(else_block) = &_input.otherwise {
+                // Visit the `otherwise` block.
+                instructions.push_str(&self.visit_statement(else_block));
+                // Add a label for the end of the `otherwise` block.
+                instructions.push_str(&format!("    position {end_otherwise_label};\n"));
+            }
+
+            // Decrement the conditional depth.
+            self.conditional_depth -= 1;
+
+            instructions
+        }
     }
 
     fn visit_iteration(&mut self, _input: &'a IterationStatement) -> String {
-        // TODO: Once loop unrolling is made optional, create a Leo error informing the user to enable the loop unrolling pass..
         unreachable!("`IterationStatement`s should not be in the AST at this phase of compilation.");
     }
 

compiler/passes/src/common/symbol_table/mod.rs

@@ -181,6 +181,11 @@ impl SymbolTable {
         }
     }
 
+    /// Attempts to lookup a variable in the current scope.
+    pub fn lookup_variable_in_current_scope(&self, location: Location) -> Option<&VariableSymbol> {
+        self.variables.get(&location)
+    }
+
     /// Returns the scope associated with `index`, if it exists in the symbol table.
     pub fn lookup_scope_by_index(&self, index: usize) -> Option<&RefCell<Self>> {
         self.scopes.get(index)

compiler/passes/src/dead_code_elimination/dead_code_eliminator.rs

@@ -26,11 +26,13 @@ pub struct DeadCodeEliminator<'a> {
     pub(crate) used_variables: IndexSet<Symbol>,
     /// Whether or not the variables are necessary.
     pub(crate) is_necessary: bool,
+    /// Whether or not we are currently traversing a finalize block.
+    pub(crate) is_finalize: bool,
 }
 
 impl<'a> DeadCodeEliminator<'a> {
     /// Initializes a new `DeadCodeEliminator`.
     pub fn new(node_builder: &'a NodeBuilder) -> Self {
-        Self { node_builder, used_variables: Default::default(), is_necessary: false }
+        Self { node_builder, used_variables: Default::default(), is_necessary: false, is_finalize: false }
     }
 }

compiler/passes/src/dead_code_elimination/eliminate_program.rs

@@ -33,9 +33,15 @@ impl ProgramReconstructor for DeadCodeEliminator<'_> {
             self.used_variables.clear();
             self.is_necessary = false;
 
+            // Set the `is_finalize` flag.
+            self.is_finalize = true;
+
             // Traverse the finalize block.
             let block = self.reconstruct_block(finalize.block).0;
 
+            // Reset the `is_finalize` flag.
+            self.is_finalize = false;
+
             Finalize {
                 identifier: finalize.identifier,
                 input: finalize.input,

compiler/passes/src/dead_code_elimination/eliminate_statement.rs

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

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) {
-        unreachable!("`ConditionalStatement`s should not be in the AST at this phase of compilation.")
+    fn reconstruct_conditional(&mut self, input: ConditionalStatement) -> (Statement, Self::AdditionalOutput) {
+        // 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, ProgramScope, Statement, 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) {
-        unreachable!("`ConditionalStatement`s should not be in the AST at this phase of compilation.")
+    fn reconstruct_conditional(&mut self, input: ConditionalStatement) -> (Statement, Self::AdditionalOutput) {
+        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.