fix: (SSA parser) translate blocks in logical order rather than syntax order

compiler/noirc_evaluator/src/ssa/parser/into_ssa.rs

@@ -1,4 +1,7 @@
-use std::{collections::HashMap, sync::Arc};
+use std::{
+    collections::{HashMap, HashSet, VecDeque},
+    sync::Arc,
+};
 
 use acvm::acir::circuit::ErrorSelector;
 use noirc_errors::call_stack::CallStackId;
@@ -157,17 +160,68 @@ impl Translator {
             } else {
                 self.builder.insert_block()
             };
-            let entry = self.blocks.entry(self.current_function_id()).or_default();
-            entry.insert(block.name.clone(), block_id);
+            let blocks = self.blocks.entry(self.current_function_id()).or_default();
+            blocks.insert(block.name.clone(), block_id);
         }
 
-        for block in function.blocks {
-            self.translate_block(block)?;
+        let entry_block_id = self.blocks[&self.current_function_id()][&function.blocks[0].name];
+
+        let mut parsed_blocks_by_id = function
+            .blocks
+            .into_iter()
+            .map(|block| {
+                let block_id = self.blocks[&self.current_function_id()][&block.name];
+                (block_id, block)
+            })
+            .collect::<HashMap<_, _>>();
+
+        let blocks_order = self.compute_blocks_order(entry_block_id, &parsed_blocks_by_id)?;
+        for block_id in blocks_order {
+            let parsed_block = parsed_blocks_by_id.remove(&block_id).unwrap();
+            self.translate_block(parsed_block)?;
         }
 
         Ok(())
     }
 
+    /// Computes the order in which blocks should be translated. The order will be according
+    /// to the block terminators, starting from the entry block. This is needed because a variable
+    /// in a block might refer to a variable that syntantically happens afterwards, but logically
+    /// happens before.
+    fn compute_blocks_order(
+        &self,
+        entry_block_id: BasicBlockId,
+        parsed_blocks_by_id: &HashMap<BasicBlockId, ParsedBlock>,
+    ) -> Result<Vec<BasicBlockId>, SsaError> {
+        let mut seen = HashSet::new();
+        let mut ordered = Vec::new();
+        let mut stack = VecDeque::new();
+
+        stack.push_back(entry_block_id);
+
+        while let Some(block_id) = stack.pop_front() {
+            if seen.contains(&block_id) {
+                continue;
+            }
+            seen.insert(block_id);
+            ordered.push(block_id);
+
+            let parsed_block = &parsed_blocks_by_id[&block_id];
+            match &parsed_block.terminator {
+                ParsedTerminator::Jmp { destination, .. } => {
+                    stack.push_back(self.lookup_block(destination)?);
+                }
+                ParsedTerminator::Jmpif { then_block, else_block, .. } => {
+                    stack.push_back(self.lookup_block(then_block)?);
+                    stack.push_back(self.lookup_block(else_block)?);
+                }
+                ParsedTerminator::Return(..) => (),
+            }
+        }
+
+        Ok(ordered)
+    }
+
     fn translate_block(&mut self, block: ParsedBlock) -> Result<(), SsaError> {
         let block_id = self.blocks[&self.current_function_id()][&block.name];
         self.builder.switch_to_block(block_id);
@@ -421,7 +475,7 @@ impl Translator {
         Ok(())
     }
 
-    fn lookup_variable(&mut self, identifier: &Identifier) -> Result<ValueId, SsaError> {
+    fn lookup_variable(&self, identifier: &Identifier) -> Result<ValueId, SsaError> {
         if let Some(value_id) = self
             .variables
             .get(&self.current_function_id())
@@ -448,23 +502,23 @@ impl Translator {
         Ok(())
     }
 
-    fn lookup_global(&mut self, identifier: Identifier) -> Result<ValueId, SsaError> {
+    fn lookup_global(&self, identifier: Identifier) -> Result<ValueId, SsaError> {
         if let Some(value_id) = self.global_values.get(&identifier.name) {
             Ok(*value_id)
         } else {
             Err(SsaError::UnknownGlobal(identifier))
         }
     }
 
-    fn lookup_block(&mut self, identifier: &Identifier) -> Result<BasicBlockId, SsaError> {
+    fn lookup_block(&self, identifier: &Identifier) -> Result<BasicBlockId, SsaError> {
         if let Some(block_id) = self.blocks[&self.current_function_id()].get(&identifier.name) {
             Ok(*block_id)
         } else {
             Err(SsaError::UnknownBlock(identifier.clone()))
         }
     }
 
-    fn lookup_function(&mut self, identifier: &Identifier) -> Result<FunctionId, SsaError> {
+    fn lookup_function(&self, identifier: &Identifier) -> Result<FunctionId, SsaError> {
         if let Some(function_id) = self.functions.get(&identifier.name) {
             Ok(*function_id)
         } else {

compiler/noirc_evaluator/src/ssa/parser/tests.rs

@@ -747,3 +747,41 @@ fn test_parses_print() {
         ";
     assert_ssa_roundtrip(src);
 }
+
+#[test]
+fn parses_variable_from_a_syntantically_following_block_but_logically_preceding_block_with_jmp() {
+    let src = "
+        acir(inline) impure fn main f0 {
+          b0():
+            jmp b2()
+          b1():
+            v5 = add v2, v4
+            return
+          b2():
+            v2 = add Field 1, Field 2
+            v4 = add v2, Field 3
+            jmp b1()
+        }
+        ";
+    assert_ssa_roundtrip(src);
+}
+
+#[test]
+fn parses_variable_from_a_syntantically_following_block_but_logically_preceding_block_with_jmpif() {
+    let src = "
+        acir(inline) impure fn main f0 {
+          b0(v0: u1):
+            jmpif v0 then: b2, else: b3
+          b1():
+            v6 = add v3, v5
+            return
+          b2():
+            return
+          b3():
+            v3 = add Field 1, Field 2
+            v5 = add v3, Field 3
+            jmp b1()
+        }
+        ";
+    assert_ssa_roundtrip(src);
+}