feat(brillig): Add unique labeling and remove relative jumps (#1652)

* feat: added unique labeling

* refactor: remove relative jumps

* feat: add section/global labeling

* pr review

* Update crates/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs

* style: clippy

---------

Co-authored-by: kevaundray <kevtheappdev@gmail.com>

crates/noirc_evaluator/src/brillig/brillig_gen/brillig_fn.rs

@@ -0,0 +1,41 @@
+use std::collections::HashMap;
+
+use acvm::acir::brillig_vm::RegisterIndex;
+
+use crate::{
+    brillig::brillig_ir::BrilligContext,
+    ssa_refactor::ir::{function::FunctionId, value::ValueId},
+};
+
+pub(crate) struct FunctionContext {
+    pub(crate) function_id: FunctionId,
+    /// Map from SSA values to Register Indices.
+    pub(crate) ssa_value_to_register: HashMap<ValueId, RegisterIndex>,
+}
+
+impl FunctionContext {
+    /// Gets a `RegisterIndex` for a `ValueId`, if one already exists
+    /// or creates a new `RegisterIndex` using the latest available
+    /// free register.
+    pub(crate) fn get_or_create_register(
+        &mut self,
+        brillig_context: &mut BrilligContext,
+        value: ValueId,
+    ) -> RegisterIndex {
+        if let Some(register_index) = self.ssa_value_to_register.get(&value) {
+            return *register_index;
+        }
+
+        let register = brillig_context.create_register();
+
+        // Cache the `ValueId` so that if we call it again, it will
+        // return the register that has just been created.
+        //
+        // WARNING: This assumes that a register has not been
+        // modified. If a MOV instruction has overwritten the value
+        // at a register, then this cache will be invalid.
+        self.ssa_value_to_register.insert(value, register);
+
+        register
+    }
+}

crates/noirc_evaluator/src/brillig/brillig_ir.rs

@@ -40,6 +40,11 @@ pub(crate) struct BrilligContext {
     latest_register: usize,
     /// Tracks memory allocations
     memory: BrilligMemory,
+    /// Context label, must be unique with respect to the function
+    /// being linked.
+    context_label: String,
+    /// Section label, used to separate sections of code
+    section_label: usize,
 }
 
 impl BrilligContext {
@@ -64,16 +69,41 @@ impl BrilligContext {
     }
 
     /// Adds a label to the next opcode
-    pub(crate) fn add_label_to_next_opcode<T: ToString>(&mut self, label: T) {
+    pub(crate) fn enter_context<T: ToString>(&mut self, label: T) {
+        self.context_label = label.to_string();
+        self.section_label = 0;
+        // Add a context label to the next opcode
         self.obj.add_label_at_position(label.to_string(), self.obj.index_of_next_opcode());
+        // Add a section label to the next opcode
+        self.obj
+            .add_label_at_position(self.current_section_label(), self.obj.index_of_next_opcode());
+    }
+
+    /// Increments the section label and adds a section label to the next opcode
+    fn enter_next_section(&mut self) {
+        self.section_label += 1;
+        self.obj
+            .add_label_at_position(self.current_section_label(), self.obj.index_of_next_opcode());
+    }
+
+    /// Internal function used to compute the section labels
+    fn compute_section_label(&self, section: usize) -> String {
+        format!("{}-{}", self.context_label, section)
+    }
+
+    /// Returns the next section label
+    fn next_section_label(&self) -> String {
+        self.compute_section_label(self.section_label + 1)
+    }
+
+    /// Returns the current section label
+    fn current_section_label(&self) -> String {
+        self.compute_section_label(self.section_label)
     }
 
     /// Adds a unresolved `Jump` instruction to the bytecode.
     pub(crate) fn jump_instruction<T: ToString>(&mut self, target_label: T) {
-        self.add_unresolved_jump(
-            BrilligOpcode::Jump { location: 0 },
-            UnresolvedJumpLocation::Label(target_label.to_string()),
-        );
+        self.add_unresolved_jump(BrilligOpcode::Jump { location: 0 }, target_label.to_string());
     }
 
     /// Adds a unresolved `JumpIf` instruction to the bytecode.
@@ -84,7 +114,7 @@ impl BrilligContext {
     ) {
         self.add_unresolved_jump(
             BrilligOpcode::JumpIf { condition, location: 0 },
-            UnresolvedJumpLocation::Label(target_label.to_string()),
+            target_label.to_string(),
         );
     }
 
@@ -109,12 +139,12 @@ impl BrilligContext {
     /// Emits brillig bytecode to jump to a trap condition if `condition`
     /// is false.
     pub(crate) fn constrain_instruction(&mut self, condition: RegisterIndex) {
-        // Jump to the relative location after the trap
         self.add_unresolved_jump(
             BrilligOpcode::JumpIf { condition, location: 0 },
-            UnresolvedJumpLocation::Relative(2),
+            self.next_section_label(),
         );
         self.push_opcode(BrilligOpcode::Trap);
+        self.enter_next_section();
     }
 
     /// Processes a return instruction.

crates/noirc_evaluator/src/brillig/brillig_ir/artifact.rs

@@ -26,25 +26,14 @@ pub(crate) type JumpInstructionPosition = OpcodeLocation;
 
 /// When constructing the bytecode, there may be instructions
 /// which require one to jump to a specific region of code (function)
-/// or a position relative to the current instruction.
 ///
 /// The position of a function cannot always be known
 /// at this point in time, so Jumps are unresolved
 /// until all functions/all of the bytecode has been processed.
 /// `Label` is used as the jump location and once all of the bytecode
 /// has been processed, the jumps are resolved using a map from Labels
 /// to their position in the bytecode.
-///
-/// Sometimes the jump destination may be relative to the jump instruction.
-/// Since the absolute position in the bytecode cannot be known until
-/// all internal and external functions have been linked, jumps of this
-/// nature cannot be fully resolved while building the bytecode either.
-/// We add relative jumps into the `Relative` variant of this enum.
-#[derive(Debug, Clone)]
-pub(crate) enum UnresolvedJumpLocation {
-    Label(String),
-    Relative(i32),
-}
+pub(crate) type UnresolvedJumpLocation = Label;
 
 impl BrilligArtifact {
     /// Link two Brillig artifacts together and resolve all unresolved jump instructions.
@@ -66,7 +55,8 @@ impl BrilligArtifact {
         }
 
         for (label_id, position_in_bytecode) in &obj.labels {
-            self.labels.insert(label_id.clone(), position_in_bytecode + offset);
+            let old_value = self.labels.insert(label_id.clone(), position_in_bytecode + offset);
+            assert!(old_value.is_none(), "overwriting label {label_id} {old_value:?}");
         }
 
         self.byte_code.extend_from_slice(&obj.byte_code);
@@ -126,12 +116,7 @@ impl BrilligArtifact {
     /// linkage with other bytecode has happened.
     fn resolve_jumps(&mut self) {
         for (location_of_jump, unresolved_location) in &self.unresolved_jumps {
-            let resolved_location = match unresolved_location {
-                UnresolvedJumpLocation::Label(label) => self.labels[label],
-                UnresolvedJumpLocation::Relative(offset) => {
-                    (offset + *location_of_jump as i32) as usize
-                }
-            };
+            let resolved_location = self.labels[unresolved_location];
 
             let jump_instruction = self.byte_code[*location_of_jump].clone();
             match jump_instruction {

crates/noirc_evaluator/src/brillig/mod.rs

@@ -1,7 +1,7 @@
 pub(crate) mod brillig_gen;
 pub(crate) mod brillig_ir;
 
-use self::{brillig_gen::BrilligGen, brillig_ir::artifact::BrilligArtifact};
+use self::{brillig_gen::convert_ssa_function, brillig_ir::artifact::BrilligArtifact};
 use crate::ssa_refactor::{
     ir::function::{Function, FunctionId, RuntimeType},
     ssa_gen::Ssa,
@@ -19,7 +19,7 @@ pub struct Brillig {
 impl Brillig {
     /// Compiles a function into brillig and store the compilation artifacts
     pub(crate) fn compile(&mut self, func: &Function) {
-        let obj = BrilligGen::compile(func);
+        let obj = convert_ssa_function(func);
         self.ssa_function_to_brillig.insert(func.id(), obj);
     }
 }

crates/noirc_evaluator/src/ssa_refactor/ir/map.rs

@@ -99,6 +99,12 @@ impl std::fmt::Display for Id<super::function::Function> {
     }
 }
 
+impl std::fmt::Display for Id<super::instruction::Instruction> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "f{}", self.index)
+    }
+}
+
 /// A DenseMap is a Vec wrapper where each element corresponds
 /// to a unique ID that can be used to access the element. No direct
 /// access to indices is provided. Since IDs must be stable and correspond