From b7ec5105c4f292abda217c3206e706663ace3af0 Mon Sep 17 00:00:00 2001
From: Patrick LaFontaine <32135464+Pat-Lafon@users.noreply.github.com>
Date: Fri, 15 Sep 2023 21:39:29 -0400
Subject: [PATCH] Address #295

---
 brilirs/Makefile                              |  2 +-
 brilirs/src/basic_block.rs                    | 44 +++++++++----------
 test/interp/core/non_linear_control_flow.bril | 26 +++++++++++
 test/interp/core/non_linear_control_flow.out  |  2 +
 4 files changed, 51 insertions(+), 23 deletions(-)
 create mode 100644 test/interp/core/non_linear_control_flow.bril
 create mode 100644 test/interp/core/non_linear_control_flow.out

diff --git a/brilirs/Makefile b/brilirs/Makefile
index f142c5f13..1d3f3168d 100644
--- a/brilirs/Makefile
+++ b/brilirs/Makefile
@@ -65,4 +65,4 @@ plot: plot.svg
 # This is primarily used for running examples and debuging a bril program
 .PHONY: example
 example:
-	bril2json < ../benchmarks/sqrt.bril | cargo run
+	bril2json < ../benchmarks/float/sqrt.bril | cargo run
diff --git a/brilirs/src/basic_block.rs b/brilirs/src/basic_block.rs
index bcd20c00a..5f78df0a3 100644
--- a/brilirs/src/basic_block.rs
+++ b/brilirs/src/basic_block.rs
@@ -289,28 +289,28 @@ impl BBFunction {
     }
     let last_idx = self.blocks.len() - 1;
     for (i, block) in self.blocks.iter_mut().enumerate() {
-      // If we're before the last block
-      if i < last_idx {
-        // Get the last instruction
-        let last_instr = block.instrs.last().cloned();
-        if let Some(bril_rs::Instruction::Effect {
-          op: bril_rs::EffectOps::Jump | bril_rs::EffectOps::Branch,
-          labels,
-          ..
-        }) = last_instr
-        {
-          for l in labels {
-            block
-              .exit
-              .push(*label_map.get(&l).ok_or(InterpError::MissingLabel(l))?);
-          }
-        } else if let Some(bril_rs::Instruction::Effect {
-          op: bril_rs::EffectOps::Return,
-          ..
-        }) = last_instr
-        {
-          // We are done, there is no exit from this block
-        } else {
+      // Get the last instruction
+      let last_instr = block.instrs.last().cloned();
+      if let Some(bril_rs::Instruction::Effect {
+        op: bril_rs::EffectOps::Jump | bril_rs::EffectOps::Branch,
+        labels,
+        ..
+      }) = last_instr
+      {
+        for l in labels {
+          block
+            .exit
+            .push(*label_map.get(&l).ok_or(InterpError::MissingLabel(l))?);
+        }
+      } else if let Some(bril_rs::Instruction::Effect {
+        op: bril_rs::EffectOps::Return,
+        ..
+      }) = last_instr
+      {
+        // We are done, there is no exit from this block
+      } else {
+        // If we're before the last block
+        if i < last_idx {
           block.exit.push(i + 1);
         }
       }
diff --git a/test/interp/core/non_linear_control_flow.bril b/test/interp/core/non_linear_control_flow.bril
new file mode 100644
index 000000000..9211d8f54
--- /dev/null
+++ b/test/interp/core/non_linear_control_flow.bril
@@ -0,0 +1,26 @@
+# ARGS: false false
+@main(b0: bool, b1: bool) {
+    jmp .start;
+  .end:
+    print x_0_2;
+    print x_1_2;
+    ret;
+  .l_1_3:
+    jmp .end;
+  .l_1_2:
+    x_1_2 : int = const 0;
+    jmp .l_1_3;
+  .l_1_1:
+    x_1_1 : int = const 1;
+    jmp .l_1_3;
+  .l_0_3:
+    br b1 .l_1_1 .l_1_2;
+  .l_0_2:
+    x_0_2 : int = const 2;
+    jmp .l_0_3;
+  .l_0_1:
+    x_0_1 : int = const 3;
+    jmp .l_0_3;
+  .start:
+    br b0 .l_0_1 .l_0_2;
+}
\ No newline at end of file
diff --git a/test/interp/core/non_linear_control_flow.out b/test/interp/core/non_linear_control_flow.out
new file mode 100644
index 000000000..389e26214
--- /dev/null
+++ b/test/interp/core/non_linear_control_flow.out
@@ -0,0 +1,2 @@
+2
+0