From 962526dab5e58a1432085b5077ff56f0f62b9063 Mon Sep 17 00:00:00 2001
From: Pranav Gaddamadugu <23022326+d0cd@users.noreply.github.com>
Date: Tue, 14 May 2024 15:20:28 -0700
Subject: [PATCH] Fix flattening

---
 .../eliminate_program.rs                      |  1 +
 .../passes/src/flattening/flatten_program.rs  | 40 +++++++++++--------
 compiler/passes/src/flattening/mod.rs         |  2 +-
 examples/basic_bank/README.md                 |  2 +-
 .../execution/cond_exec_in_finalize.out       | 12 +++---
 5 files changed, 33 insertions(+), 24 deletions(-)

diff --git a/compiler/passes/src/dead_code_elimination/eliminate_program.rs b/compiler/passes/src/dead_code_elimination/eliminate_program.rs
index 2a760912a5..1c58e28d95 100644
--- a/compiler/passes/src/dead_code_elimination/eliminate_program.rs
+++ b/compiler/passes/src/dead_code_elimination/eliminate_program.rs
@@ -23,6 +23,7 @@ impl ProgramReconstructor for DeadCodeEliminator<'_> {
         // Reset the state of the dead code eliminator.
         self.used_variables.clear();
         self.is_necessary = false;
+        self.is_async = input.variant.is_async_function();
 
         // Traverse the function body.
         let block = self.reconstruct_block(input.block).0;
diff --git a/compiler/passes/src/flattening/flatten_program.rs b/compiler/passes/src/flattening/flatten_program.rs
index 6a837f5215..6c09fe4959 100644
--- a/compiler/passes/src/flattening/flatten_program.rs
+++ b/compiler/passes/src/flattening/flatten_program.rs
@@ -41,25 +41,33 @@ impl ProgramReconstructor for Flattener<'_> {
 
     /// Flattens a function's body
     fn reconstruct_function(&mut self, function: Function) -> Function {
-        // Flatten the function body.
-        let mut block = self.reconstruct_block(function.block).0;
+        // 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;
 
-        // Get all of the guards and return expression.
-        let returns = self.clear_early_returns();
+            // 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);
+            // 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,
+            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,
+            }
         }
     }
 }
diff --git a/compiler/passes/src/flattening/mod.rs b/compiler/passes/src/flattening/mod.rs
index ac35cdb12b..366e254864 100644
--- a/compiler/passes/src/flattening/mod.rs
+++ b/compiler/passes/src/flattening/mod.rs
@@ -18,7 +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.
+//! Note that this transformation is not applied to async functioins.
 //!
 //! Consider the following Leo code, output by the SSA pass.
 //! ```leo
diff --git a/examples/basic_bank/README.md b/examples/basic_bank/README.md
index 67cab4c84d..07d07b033c 100644
--- a/examples/basic_bank/README.md
+++ b/examples/basic_bank/README.md
@@ -32,7 +32,7 @@ Can you find any others?
 - record ownership
 - loops and bounded iteration
 - mappings
-- finalize
+- `async`/`await`
 
 ## Running the Program
 
diff --git a/tests/expectations/execution/cond_exec_in_finalize.out b/tests/expectations/execution/cond_exec_in_finalize.out
index 9c1468f42d..818cb31259 100644
--- a/tests/expectations/execution/cond_exec_in_finalize.out
+++ b/tests/expectations/execution/cond_exec_in_finalize.out
@@ -9,11 +9,11 @@ outputs:
           initial_ast: 27c7103cf4aef0e2bb8a1abb044f3aa3a02b6d58c08a0f2cca78d13f7b061665
           unrolled_ast: 27c7103cf4aef0e2bb8a1abb044f3aa3a02b6d58c08a0f2cca78d13f7b061665
           ssa_ast: 5d2d844bb95e6a45dffd4786aa55c2483f7cda49617cb74671e9687c1cb75a74
-          flattened_ast: df415243c990cb8a30c4e863f60fcfebbcbd3ddd8d698db902fe44c3726b9da5
-          destructured_ast: 011a8cb6f4a57de5947f443b7e7641f1151fb80f0ec1e973795d06162a046283
-          inlined_ast: 9a9ca324838b653871e9309e62cb3d0d3258479f732296d798d21017d4cb3391
-          dce_ast: 9a9ca324838b653871e9309e62cb3d0d3258479f732296d798d21017d4cb3391
-          bytecode: c874534d8f0d3fbd6314e9c33907a596cf23f1f6d070776423b0316f40251d44
+          flattened_ast: 8cb0fae560bad8dba170c24763ed47e1d9a211a03b3f63a7a8c8777c8b83b931
+          destructured_ast: 55c8f0065d1d91bf3d99c0ac656fdd00ddd6e196767dd083cbd2e4c188afa737
+          inlined_ast: b5b1c94a03bd528b6a5dd3273304476903401287a79e90b3edb7cbae5114fa89
+          dce_ast: b5b1c94a03bd528b6a5dd3273304476903401287a79e90b3edb7cbae5114fa89
+          bytecode: 4c9b370477c0b08af171a80b375e88adaa53684d7419d52effb5504bbc5e9b59
           errors: ""
           warnings: ""
       execute:
@@ -39,7 +39,7 @@ outputs:
             global_state_root: sr1qqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqq6gk0xu
             proof: proof1qyqsqqqqqqqqqqqpqqqqqqqqqqq0rxwjqlxmc8acm0xaxeewz905709n3wjz45gmdpv6q0yy2pqhzj0dj7zaq2x32veqxg2aqveccuyqqys9vejps3830l50n6znqk47m7tf00xv4qnpg9vpz4l45rqs7n2hdyxutqjjd6w3j3le0p64u5dpjq2jhfant7td2w0wrr0r5v762k8lujm2r3rq283jfkl0aqrnfuhhklhxpe5zrhvcndj95p8drpw7fsqggqfk9e40l7nw600kel4p36ggaapgjpzpqlanwuuwlj3mas52p6qefk7xkacgcldxka9saqf4wwgqu9e87lv9jy68fzr65nujqdlxkds23l30nqejxfh0d6ycvvmgfzutps9ggafhelp0vc7zude43yjsrnx3348l4wzgu4nncayjlc88sw60h02lkn79ak8hszv4gqejvdlam4v5ylx5wsde7y5af0jrwkn2qp6st8j0ns5hafgh6mex29dt7ah26clzzmgm6ack4wvq64905ja9xuzsex6rf3uvahk89wzhlzx5eqwq7894l2l0f7vkmglk8d85987qttja0ug87594vxrkaadd008ary7fw94xyzga40cea6rlqqksq6qk52l549ze48pwenug4pg8g9lrfe3xqsl6dezyxa2a7vtnk83rhhw4uv9h4q4cas9v9r2meqkxmgspwe2r4gxyhn8zu2rraz9xspvj2skkarq30rv58zdqa6a6fftucu8l6788afecfl5r22dta56xfm6ze432hcxupq54sqpsug3dh2a6jqdkpw75pn7kdf60yfj0nru6w920cp808jxch6v3rj04rej26rvyqk5uz90k5kdlrg4c6hdxje2qtueffyw4mjr93yyuq8qjccstrq4qaad77dv42yl3wf5d0pqgceeurf4dtrzey0dt6jtz0tmynmvyzhcqycrcks09wslyy699lufkjcwknxju3t88vh9vffj0xmr5fzsh6cr72mjhwsrelh8lt023xeg0x9zvwn63pgz2es37p5dzlc3kut8sqp2tfy5vpazc58c5u25wg69aqa9vgjxkk94uwugx2n80jhwjqpugpxx9huu297v629hdjchld26c394rydn9t4rckf8ka63q3reelczq9uqdrgjh4azp353rx5x3matwx550cu9kgtwd7jnf76gl67nppkqfqvqqqqqqqqqqqwt6l5ch7kf7lrjfkpz83v2yw4mek6f6www8lwjucqpev25xf2x63tztw2x2mmqpp022yy8faa5csqqz5chxf3fde47ngzntjzhche0lv9ya3ywfagh8qxs490phu0n8pu0pz966d627zse28nap87cn65cpqy370fjf7z03rtrwshkz3fdu3hmknkdcx2sw72yv7elklj780n4qrq3lx7dmmuez3pj84zhhdhvuk9j9cpzxu74hfjkym23dkqya4cqkrtf5xhl06kpnqs6lfahpxjtysyqqwxuqrc
           verified: true
-          status: rejected
+          status: accepted
           errors: ""
           warnings: ""
         - execution: