From e3bcda669c1da5152deed07579d3a658e9ff9946 Mon Sep 17 00:00:00 2001
From: Philippe Laferriere <plafer@protonmail.com>
Date: Wed, 14 Aug 2024 16:54:15 -0400
Subject: [PATCH] chore: fix tests that didn't assume merging

---
 assembly/src/assembler/mod.rs   |   3 +-
 assembly/src/assembler/tests.rs |  30 +++-------
 assembly/src/tests.rs           | 102 +++++++++++++-------------------
 3 files changed, 52 insertions(+), 83 deletions(-)

diff --git a/assembly/src/assembler/mod.rs b/assembly/src/assembler/mod.rs
index d2ba89f89..8070c4d2f 100644
--- a/assembly/src/assembler/mod.rs
+++ b/assembly/src/assembler/mod.rs
@@ -746,7 +746,8 @@ fn merge_basic_blocks(
     let mut decorators = DecoratorList::new();
 
     for &basic_block_node_id in contiguous_basic_block_ids {
-        // TODO(plafer): document unwraps
+        // It is safe to unwrap here, since we already checked that all IDs in
+        // `contiguous_basic_block_ids` are `BasicBlockNode`s
         let basic_block_node = mast_forest_builder
             .get_mast_node(basic_block_node_id)
             .unwrap()
diff --git a/assembly/src/assembler/tests.rs b/assembly/src/assembler/tests.rs
index a111b65b4..1d838bcb9 100644
--- a/assembly/src/assembler/tests.rs
+++ b/assembly/src/assembler/tests.rs
@@ -92,22 +92,6 @@ fn nested_blocks() -> Result<(), Report> {
 
     let program = assembler.assemble_program(program).unwrap();
 
-    let exec_bar_node_id = {
-        // bar procedure
-        let basic_block_1_id = expected_mast_forest_builder
-            .ensure_block(vec![Operation::Push(17_u32.into())], None)
-            .unwrap();
-
-        // Basic block representing the `foo` procedure
-        let basic_block_2_id = expected_mast_forest_builder
-            .ensure_block(vec![Operation::Push(19_u32.into())], None)
-            .unwrap();
-
-        expected_mast_forest_builder
-            .ensure_join(basic_block_1_id, basic_block_2_id)
-            .unwrap()
-    };
-
     // basic block representing foo::bar.baz procedure
     let exec_foo_bar_baz_node_id = expected_mast_forest_builder
         .ensure_block(vec![Operation::Push(29_u32.into())], None)
@@ -134,13 +118,15 @@ fn nested_blocks() -> Result<(), Report> {
     let r#true2 = expected_mast_forest_builder.ensure_split(r#true3, r#false3).unwrap();
 
     let r#while = {
-        let push_basic_block_id = {
-            expected_mast_forest_builder
-                .ensure_block(vec![Operation::Push(23u32.into())], None)
-                .unwrap()
-        };
         let body_node_id = expected_mast_forest_builder
-            .ensure_join(exec_bar_node_id, push_basic_block_id)
+            .ensure_block(
+                vec![
+                    Operation::Push(17u32.into()),
+                    Operation::Push(19u32.into()),
+                    Operation::Push(23u32.into()),
+                ],
+                None,
+            )
             .unwrap();
 
         expected_mast_forest_builder.ensure_loop(body_node_id).unwrap()
diff --git a/assembly/src/tests.rs b/assembly/src/tests.rs
index d84374dee..0660b6c6f 100644
--- a/assembly/src/tests.rs
+++ b/assembly/src/tests.rs
@@ -1285,17 +1285,11 @@ begin
                 end
             else
                 join
-                    join
+                    basic_block mul push(8) push(8) end
+                    if.true
                         basic_block mul end
-                        basic_block push(8) end
-                    end
-                    join
-                        basic_block push(8) end
-                        if.true
-                            basic_block mul end
-                        else
-                            basic_block noop end
-                        end
+                    else
+                        basic_block noop end
                     end
                 end
             end
@@ -1320,10 +1314,7 @@ fn program_with_one_procedure() -> TestResult {
     let program = context.assemble(source)?;
     let expected = "\
 begin
-    join
-        basic_block push(2) push(3) add end
-        basic_block push(3) push(7) mul end
-    end
+    basic_block push(2) push(3) add push(3) push(7) mul end
 end";
     assert_str_eq!(format!("{program}"), expected);
     Ok(())
@@ -1342,24 +1333,21 @@ fn program_with_nested_procedure() -> TestResult {
     let program = context.assemble(source)?;
     let expected = "\
 begin
-    join
-        join
-            join
-                basic_block push(2) push(4) add end
-                basic_block push(3) push(7) mul end
-            end
-            join
-                basic_block push(11) end
-                join
-                    join
-                        basic_block push(5) end
-                        basic_block push(3) push(7) mul end
-                    end
-                    basic_block add end
-                end
-            end
-        end
-        basic_block neg add end
+    basic_block
+        push(2)
+        push(4)
+        add
+        push(3)
+        push(7)
+        mul
+        push(11)
+        push(5)
+        push(3)
+        push(7)
+        mul
+        add
+        neg
+        add
     end
 end";
     assert_str_eq!(format!("{program}"), expected);
@@ -1386,23 +1374,23 @@ fn program_with_proc_locals() -> TestResult {
     let program = context.assemble(source)?;
     let expected = "\
 begin
-    join
-        basic_block push(4) push(3) push(2) end
-        basic_block
-            push(1)
-            fmpupdate
-            pad
-            fmpadd
-            mstore
-            drop
-            add
-            pad
-            fmpadd
-            mload
-            mul
-            push(18446744069414584320)
-            fmpupdate
-        end
+    basic_block
+        push(4)
+        push(3)
+        push(2)
+        push(1)
+        fmpupdate
+        pad
+        fmpadd
+        mstore
+        drop
+        add
+        pad
+        fmpadd
+        mload
+        mul
+        push(18446744069414584320)
+        fmpupdate
     end
 end";
     assert_str_eq!(format!("{program}"), expected);
@@ -1603,7 +1591,7 @@ begin
     join
         join
             basic_block push(4) push(3) end
-            external.0x20234ee941e53a15886e733cc8e041198c6e90d2a16ea18ce1030e8c3596dd38
+            external.0xc2545da99d3a1f3f38d957c7893c44d78998d8ea8b11aba7e22c8c2b2a213dae
         end
         call.0x20234ee941e53a15886e733cc8e041198c6e90d2a16ea18ce1030e8c3596dd38
     end
@@ -2060,17 +2048,11 @@ begin
                 end
             else
                 join
-                    join
+                    basic_block mul push(8) push(8) end
+                    if.true
                         basic_block mul end
-                        basic_block push(8) end
-                    end
-                    join
-                        basic_block push(8) end
-                        if.true
-                            basic_block mul end
-                        else
-                            basic_block noop end
-                        end
+                    else
+                        basic_block noop end
                     end
                 end
             end