Fix flattening

compiler/compiler/tests/utilities/mod.rs

@@ -281,6 +281,8 @@ pub fn compile_and_process<'a>(parsed: &'a mut Compiler<'a, CurrentNetwork>) ->
     // Compile Leo program to bytecode.
     let bytecode = parsed.code_generation_pass(&st, &struct_graph, &call_graph)?;
 
+    println!("Bytecode: {bytecode}");
+
     Ok(bytecode)
 }
 

compiler/passes/src/flattening/flatten_program.rs

@@ -41,33 +41,28 @@ impl ProgramReconstructor for Flattener<'_> {
 
     /// Flattens a function's body
     fn reconstruct_function(&mut self, function: Function) -> Function {
-        // If the function is an async function, then return it as is.
-        // Note that async functions are not flattened since it uses `branch` instructions to produce correct code in for conditional execution.
-        if function.variant.is_async_function() {
-            return function;
-        }
-        // Otherwise, flatten the function body.
-        else {
-            // Flatten the function body.
-            let mut block = self.reconstruct_block(function.block).0;
+        // Set when the function is an async function.
+        self.is_async = function.variant.is_async_function();
 
-            // Get all of the guards and return expression.
-            let returns = self.clear_early_returns();
+        // Flatten the function body.
+        let mut block = self.reconstruct_block(function.block).0;
 
-            // Fold the return statements into the block.
-            self.fold_returns(&mut block, returns);
+        // Get all of the guards and return expression.
+        let returns = self.clear_early_returns();
 
-            Function {
-                annotations: function.annotations,
-                variant: function.variant,
-                identifier: function.identifier,
-                input: function.input,
-                output: function.output,
-                output_type: function.output_type,
-                block,
-                span: function.span,
-                id: function.id,
-            }
+        // Fold the return statements into the block.
+        self.fold_returns(&mut block, returns);
+
+        Function {
+            annotations: function.annotations,
+            variant: function.variant,
+            identifier: function.identifier,
+            input: function.input,
+            output: function.output,
+            output_type: function.output_type,
+            block,
+            span: function.span,
+            id: function.id,
         }
     }
 }

compiler/passes/src/flattening/flatten_statement.rs

@@ -61,6 +61,11 @@ impl StatementReconstructor for Flattener<'_> {
     fn reconstruct_assert(&mut self, input: AssertStatement) -> (Statement, Self::AdditionalOutput) {
         let mut statements = Vec::new();
 
+        // If we are traversing an async function, then we can return the assert as it.
+        if self.is_async {
+            return (Statement::Assert(input), statements);
+        }
+
         // Flatten the arguments of the assert statement.
         let assert = AssertStatement {
             span: input.span,
@@ -222,6 +227,29 @@ impl StatementReconstructor for Flattener<'_> {
     fn reconstruct_conditional(&mut self, conditional: ConditionalStatement) -> (Statement, Self::AdditionalOutput) {
         let mut statements = Vec::with_capacity(conditional.then.statements.len());
 
+        // If we are traversing an async function, reconstruct the if and else blocks, but do not flatten them.
+        if self.is_async {
+            let then_block = self.reconstruct_block(conditional.then).0;
+            let otherwise_block = match conditional.otherwise {
+                Some(statement) => match *statement {
+                    Statement::Block(block) => self.reconstruct_block(block).0,
+                    _ => unreachable!("SSA guarantees that the `otherwise` is always a `Block`"),
+                },
+                None => Block { span: Default::default(), statements: Vec::new(), id: self.node_builder.next_id() },
+            };
+
+            return (
+                Statement::Conditional(ConditionalStatement {
+                    condition: conditional.condition,
+                    then: then_block,
+                    otherwise: Some(Box::new(Statement::Block(otherwise_block))),
+                    span: conditional.span,
+                    id: conditional.id,
+                }),
+                statements,
+            );
+        }
+
         // Add condition to the condition stack.
         self.condition_stack.push(conditional.condition.clone());
 
@@ -269,6 +297,10 @@ impl StatementReconstructor for Flattener<'_> {
     /// Transforms a return statement into an empty block statement.
     /// Stores the arguments to the return statement, which are later folded into a single return statement at the end of the function.
     fn reconstruct_return(&mut self, input: ReturnStatement) -> (Statement, Self::AdditionalOutput) {
+        // If we are traversing an async function, return as is.
+        if self.is_async {
+            return (Statement::Return(input), Default::default());
+        }
         // Construct the associated guard.
         let guard = self.construct_guard();
 

compiler/passes/src/flattening/flattener.rs

@@ -67,6 +67,8 @@ pub struct Flattener<'a> {
     pub(crate) returns: Vec<(Option<Expression>, ReturnStatement)>,
     /// The program name.
     pub(crate) program: Option<Symbol>,
+    /// Whether the function is an async function.
+    pub(crate) is_async: bool,
 }
 
 impl<'a> Flattener<'a> {
@@ -84,6 +86,7 @@ impl<'a> Flattener<'a> {
             condition_stack: Vec::new(),
             returns: Vec::new(),
             program: None,
+            is_async: false,
         }
     }
 

tests/expectations/compiler/finalize/only_finalize_with_flattening.out

@@ -9,10 +9,10 @@ outputs:
           initial_ast: 31d6210343f8a439e9a731aa2b344a0a35bb4828fcbfb5b7f3d3c532aa0d49e0
           unrolled_ast: 31d6210343f8a439e9a731aa2b344a0a35bb4828fcbfb5b7f3d3c532aa0d49e0
           ssa_ast: ed0d528c739439b087da26d083d1a1c6705e5e9b020f6dbb6d1510929df3079f
-          flattened_ast: 6022517c96f9f27f698a3763e2745248829606c98f0089a6ec969c21e46c17e8
-          destructured_ast: d3da411156bdf490343160266cbc638347e8699097e129137eaf2eb328d5227c
-          inlined_ast: ab6e5a75c050c438cf92aeb65992405668b7587c04c088d5b158a9bbecf31c55
-          dce_ast: ab6e5a75c050c438cf92aeb65992405668b7587c04c088d5b158a9bbecf31c55
-          bytecode: d6edcce70bf27b2fad397a62ae0bee08448a0c157d89e49867d843d83a04bfb7
+          flattened_ast: 853653a2db1cb618cf7fe214cc90ff1ee6f952dbc6e945b6b3c4b34ca07e906d
+          destructured_ast: 0fa73a1e802c55fe76758295dbeb0e4a9340fd95c14b05b033fd9aec1a039fed
+          inlined_ast: 17df43530f88543010218898759c3459f380370012717bebaec75f9a27d40130
+          dce_ast: 17df43530f88543010218898759c3459f380370012717bebaec75f9a27d40130
+          bytecode: ed6b317a872c3d7083d96539a0f0601af2dcc1a6d198d897edf7f6d4db26e47a
           errors: ""
           warnings: ""