feat(brillig): Added locations for brillig artifacts (#2415)

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

@@ -1,6 +1,7 @@
 use crate::brillig::brillig_ir::{
     BrilligBinaryOp, BrilligContext, BRILLIG_INTEGER_ARITHMETIC_BIT_SIZE,
 };
+use crate::ssa::ir::dfg::CallStack;
 use crate::ssa::ir::{
     basic_block::{BasicBlock, BasicBlockId},
     dfg::DataFlowGraph,
@@ -202,6 +203,7 @@ impl<'block> BrilligBlock<'block> {
     /// Converts an SSA instruction into a sequence of Brillig opcodes.
     fn convert_ssa_instruction(&mut self, instruction_id: InstructionId, dfg: &DataFlowGraph) {
         let instruction = &dfg[instruction_id];
+        self.brillig_context.set_call_stack(dfg.get_call_stack(instruction_id));
 
         match instruction {
             Instruction::Binary(binary) => {
@@ -479,6 +481,8 @@ impl<'block> BrilligBlock<'block> {
             }
             _ => todo!("ICE: Instruction not supported {instruction:?}"),
         };
+
+        self.brillig_context.set_call_stack(CallStack::new());
     }
 
     fn convert_ssa_function_call(

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

@@ -2,8 +2,10 @@ use acvm::acir::brillig::{
     BinaryFieldOp, BinaryIntOp, Opcode as BrilligOpcode, RegisterIndex, Value,
 };
 
+use crate::brillig::brillig_ir::artifact::GeneratedBrillig;
+
 /// Generates brillig bytecode which computes the inverse of its input if not null, and zero else.
-pub(crate) fn directive_invert() -> Vec<BrilligOpcode> {
+pub(crate) fn directive_invert() -> GeneratedBrillig {
     //  We generate the following code:
     // fn invert(x : Field) -> Field {
     //    1/ x
@@ -16,20 +18,23 @@ pub(crate) fn directive_invert() -> Vec<BrilligOpcode> {
     // Location of the stop opcode
     let stop_location = 3;
 
-    vec![
-        // If the input is zero, then we jump to the stop opcode
-        BrilligOpcode::JumpIfNot { condition: input, location: stop_location },
-        // Put value one in register (1)
-        BrilligOpcode::Const { destination: one_const, value: Value::from(1_usize) },
-        // Divide 1 by the input, and set the result of the division into register (0)
-        BrilligOpcode::BinaryFieldOp {
-            op: BinaryFieldOp::Div,
-            lhs: one_const,
-            rhs: input,
-            destination: input,
-        },
-        BrilligOpcode::Stop,
-    ]
+    GeneratedBrillig {
+        byte_code: vec![
+            // If the input is zero, then we jump to the stop opcode
+            BrilligOpcode::JumpIfNot { condition: input, location: stop_location },
+            // Put value one in register (1)
+            BrilligOpcode::Const { destination: one_const, value: Value::from(1_usize) },
+            // Divide 1 by the input, and set the result of the division into register (0)
+            BrilligOpcode::BinaryFieldOp {
+                op: BinaryFieldOp::Div,
+                lhs: one_const,
+                rhs: input,
+                destination: input,
+            },
+            BrilligOpcode::Stop,
+        ],
+        locations: Default::default(),
+    }
 }
 
 /// Generates brillig bytecode which computes `a / b` and returns the quotient and remainder.
@@ -47,43 +52,55 @@ pub(crate) fn directive_invert() -> Vec<BrilligOpcode> {
 ///    }
 /// }
 /// ```
-pub(crate) fn directive_quotient(bit_size: u32) -> Vec<BrilligOpcode> {
+pub(crate) fn directive_quotient(bit_size: u32) -> GeneratedBrillig {
     // `a` is (0) (i.e register index 0)
     // `b` is (1)
     // `predicate` is (2)
-    vec![
-        // If the predicate is zero, we jump to the exit segment
-        BrilligOpcode::JumpIfNot { condition: RegisterIndex::from(2), location: 6 },
-        //q = a/b is set into register (3)
-        BrilligOpcode::BinaryIntOp {
-            op: BinaryIntOp::UnsignedDiv,
-            lhs: RegisterIndex::from(0),
-            rhs: RegisterIndex::from(1),
-            destination: RegisterIndex::from(3),
-            bit_size,
-        },
-        //(1)= q*b
-        BrilligOpcode::BinaryIntOp {
-            op: BinaryIntOp::Mul,
-            lhs: RegisterIndex::from(3),
-            rhs: RegisterIndex::from(1),
-            destination: RegisterIndex::from(1),
-            bit_size,
-        },
-        //(1) = a-q*b
-        BrilligOpcode::BinaryIntOp {
-            op: BinaryIntOp::Sub,
-            lhs: RegisterIndex::from(0),
-            rhs: RegisterIndex::from(1),
-            destination: RegisterIndex::from(1),
-            bit_size,
-        },
-        //(0) = q
-        BrilligOpcode::Mov { destination: RegisterIndex::from(0), source: RegisterIndex::from(3) },
-        BrilligOpcode::Stop,
-        // Exit segment: we return 0,0
-        BrilligOpcode::Const { destination: RegisterIndex::from(0), value: Value::from(0_usize) },
-        BrilligOpcode::Const { destination: RegisterIndex::from(1), value: Value::from(0_usize) },
-        BrilligOpcode::Stop,
-    ]
+    GeneratedBrillig {
+        byte_code: vec![
+            // If the predicate is zero, we jump to the exit segment
+            BrilligOpcode::JumpIfNot { condition: RegisterIndex::from(2), location: 6 },
+            //q = a/b is set into register (3)
+            BrilligOpcode::BinaryIntOp {
+                op: BinaryIntOp::UnsignedDiv,
+                lhs: RegisterIndex::from(0),
+                rhs: RegisterIndex::from(1),
+                destination: RegisterIndex::from(3),
+                bit_size,
+            },
+            //(1)= q*b
+            BrilligOpcode::BinaryIntOp {
+                op: BinaryIntOp::Mul,
+                lhs: RegisterIndex::from(3),
+                rhs: RegisterIndex::from(1),
+                destination: RegisterIndex::from(1),
+                bit_size,
+            },
+            //(1) = a-q*b
+            BrilligOpcode::BinaryIntOp {
+                op: BinaryIntOp::Sub,
+                lhs: RegisterIndex::from(0),
+                rhs: RegisterIndex::from(1),
+                destination: RegisterIndex::from(1),
+                bit_size,
+            },
+            //(0) = q
+            BrilligOpcode::Mov {
+                destination: RegisterIndex::from(0),
+                source: RegisterIndex::from(3),
+            },
+            BrilligOpcode::Stop,
+            // Exit segment: we return 0,0
+            BrilligOpcode::Const {
+                destination: RegisterIndex::from(0),
+                value: Value::from(0_usize),
+            },
+            BrilligOpcode::Const {
+                destination: RegisterIndex::from(1),
+                value: Value::from(0_usize),
+            },
+            BrilligOpcode::Stop,
+        ],
+        locations: Default::default(),
+    }
 }

crates/noirc_evaluator/src/brillig/brillig_ir.rs

@@ -10,6 +10,8 @@ pub(crate) mod registers;
 
 mod entry_point;
 
+use crate::ssa::ir::dfg::CallStack;
+
 use self::{
     artifact::{BrilligArtifact, UnresolvedJumpLocation},
     registers::BrilligRegistersContext,
@@ -104,7 +106,7 @@ impl BrilligContext {
 
     /// Adds a brillig instruction to the brillig byte code
     pub(crate) fn push_opcode(&mut self, opcode: BrilligOpcode) {
-        self.obj.byte_code.push(opcode);
+        self.obj.push_opcode(opcode);
     }
 
     /// Returns the artifact
@@ -945,6 +947,11 @@ impl BrilligContext {
             _ => unreachable!("ICE: Expected vector, got {variable:?}"),
         }
     }
+
+    /// Sets a current call stack that the next pushed opcodes will be associated with.
+    pub(crate) fn set_call_stack(&mut self, call_stack: CallStack) {
+        self.obj.set_call_stack(call_stack);
+    }
 }
 
 /// Type to encapsulate the binary operation types in Brillig
@@ -970,7 +977,7 @@ pub(crate) mod tests {
 
     use crate::brillig::brillig_ir::BrilligContext;
 
-    use super::artifact::BrilligParameter;
+    use super::artifact::{BrilligParameter, GeneratedBrillig};
     use super::{BrilligOpcode, ReservedRegisters};
 
     pub(crate) struct DummyBlackBoxSolver;
@@ -1010,7 +1017,7 @@ pub(crate) mod tests {
         context: BrilligContext,
         arguments: Vec<BrilligParameter>,
         returns: Vec<BrilligParameter>,
-    ) -> Vec<BrilligOpcode> {
+    ) -> GeneratedBrillig {
         let artifact = context.artifact();
         let mut entry_point_artifact =
             BrilligContext::new_entry_point_artifact(arguments, returns, "test".to_string());
@@ -1028,7 +1035,7 @@ pub(crate) mod tests {
         let mut vm = VM::new(
             Registers { inner: param_registers },
             memory,
-            create_entry_point_bytecode(context, arguments, returns),
+            create_entry_point_bytecode(context, arguments, returns).byte_code,
             vec![],
             &DummyBlackBoxSolver,
         );
@@ -1079,7 +1086,7 @@ pub(crate) mod tests {
 
         context.stop_instruction();
 
-        let bytecode = context.artifact().byte_code;
+        let bytecode = context.artifact().finish().byte_code;
         let number_sequence: Vec<Value> = (0_usize..12_usize).map(Value::from).collect();
         let mut vm = VM::new(
             Registers { inner: vec![] },

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

@@ -1,5 +1,7 @@
 use acvm::acir::brillig::Opcode as BrilligOpcode;
-use std::collections::HashMap;
+use std::collections::{BTreeMap, HashMap};
+
+use crate::ssa::ir::dfg::CallStack;
 
 /// Represents a parameter or a return value of a function.
 #[derive(Debug, Clone)]
@@ -9,11 +11,19 @@ pub(crate) enum BrilligParameter {
     Slice(Vec<BrilligParameter>),
 }
 
+/// The result of compiling and linking brillig artifacts.
+/// This is ready to run bytecode with attached metadata.
+#[derive(Debug)]
+pub(crate) struct GeneratedBrillig {
+    pub(crate) byte_code: Vec<BrilligOpcode>,
+    pub(crate) locations: BTreeMap<OpcodeLocation, CallStack>,
+}
+
 #[derive(Default, Debug, Clone)]
 /// Artifacts resulting from the compilation of a function into brillig byte code.
 /// Currently it is just the brillig bytecode of the function.
 pub(crate) struct BrilligArtifact {
-    pub(crate) byte_code: Vec<BrilligOpcode>,
+    byte_code: Vec<BrilligOpcode>,
     /// The set of jumps that need to have their locations
     /// resolved.
     unresolved_jumps: Vec<(JumpInstructionPosition, UnresolvedJumpLocation)>,
@@ -26,6 +36,10 @@ pub(crate) struct BrilligArtifact {
     /// TODO: perhaps we should combine this with the `unresolved_jumps` field
     /// TODO: and have an enum which indicates whether the jump is internal or external
     unresolved_external_call_labels: Vec<(JumpInstructionPosition, UnresolvedJumpLocation)>,
+    /// Maps the opcodes that are associated with a callstack to it.
+    locations: BTreeMap<OpcodeLocation, CallStack>,
+    /// The current call stack. All opcodes that are pushed will be associated with this call stack.
+    call_stack: CallStack,
 }
 
 /// A pointer to a location in the opcode.
@@ -52,9 +66,9 @@ pub(crate) type UnresolvedJumpLocation = Label;
 
 impl BrilligArtifact {
     /// Resolves all jumps and generates the final bytecode
-    pub(crate) fn finish(mut self) -> Vec<BrilligOpcode> {
+    pub(crate) fn finish(mut self) -> GeneratedBrillig {
         self.resolve_jumps();
-        self.byte_code
+        GeneratedBrillig { byte_code: self.byte_code, locations: self.locations }
     }
 
     /// Gets the first unresolved function call of this artifact.
@@ -116,10 +130,17 @@ impl BrilligArtifact {
             self.unresolved_external_call_labels
                 .push((position_in_bytecode + offset, label_id.clone()));
         }
+
+        for (position_in_bytecode, call_stack) in obj.locations.iter() {
+            self.locations.insert(position_in_bytecode + offset, call_stack.clone());
+        }
     }
 
     /// Adds a brillig instruction to the brillig byte code
     pub(crate) fn push_opcode(&mut self, opcode: BrilligOpcode) {
+        if !self.call_stack.is_empty() {
+            self.locations.insert(self.index_of_next_opcode(), self.call_stack.clone());
+        }
         self.byte_code.push(opcode);
     }
 
@@ -217,4 +238,8 @@ impl BrilligArtifact {
             }
         }
     }
+
+    pub(crate) fn set_call_stack(&mut self, call_stack: CallStack) {
+        self.call_stack = call_stack;
+    }
 }

crates/noirc_evaluator/src/ssa/acir_gen/acir_ir/acir_variable.rs

@@ -1,5 +1,6 @@
 use super::generated_acir::GeneratedAcir;
 use crate::brillig::brillig_gen::brillig_directive;
+use crate::brillig::brillig_ir::artifact::GeneratedBrillig;
 use crate::errors::{InternalError, RuntimeError};
 use crate::ssa::acir_gen::{AcirDynamicArray, AcirValue};
 use crate::ssa::ir::dfg::CallStack;
@@ -910,7 +911,7 @@ impl AcirContext {
     pub(crate) fn brillig(
         &mut self,
         predicate: AcirVar,
-        code: Vec<BrilligOpcode>,
+        generated_brillig: GeneratedBrillig,
         inputs: Vec<AcirValue>,
         outputs: Vec<AcirType>,
     ) -> Result<Vec<AcirValue>, InternalError> {
@@ -932,7 +933,9 @@ impl AcirContext {
 
         // Optimistically try executing the brillig now, if we can complete execution they just return the results.
         // This is a temporary measure pending SSA optimizations being applied to Brillig which would remove constant-input opcodes (See #2066)
-        if let Some(brillig_outputs) = self.execute_brillig(code.clone(), &b_inputs, &outputs) {
+        if let Some(brillig_outputs) =
+            self.execute_brillig(generated_brillig.byte_code.clone(), &b_inputs, &outputs)
+        {
             return Ok(brillig_outputs);
         }
 
@@ -953,7 +956,7 @@ impl AcirContext {
             }
         });
         let predicate = self.var_to_expression(predicate)?;
-        self.acir_ir.brillig(Some(predicate), code, b_inputs, b_outputs);
+        self.acir_ir.brillig(Some(predicate), generated_brillig, b_inputs, b_outputs);
 
         Ok(outputs_var)
     }

crates/noirc_evaluator/src/ssa/acir_gen/acir_ir/generated_acir.rs

@@ -3,13 +3,12 @@
 use std::collections::BTreeMap;
 
 use crate::{
-    brillig::brillig_gen::brillig_directive,
+    brillig::{brillig_gen::brillig_directive, brillig_ir::artifact::GeneratedBrillig},
     errors::{InternalError, RuntimeError},
     ssa::ir::dfg::CallStack,
 };
 
 use acvm::acir::{
-    brillig::Opcode as BrilligOpcode,
     circuit::{
         brillig::{Brillig as AcvmBrillig, BrilligInputs, BrilligOutputs},
         opcodes::{BlackBoxFuncCall, FunctionInput, Opcode as AcirOpcode},
@@ -788,18 +787,24 @@ impl GeneratedAcir {
     pub(crate) fn brillig(
         &mut self,
         predicate: Option<Expression>,
-        code: Vec<BrilligOpcode>,
+        generated_brillig: GeneratedBrillig,
         inputs: Vec<BrilligInputs>,
         outputs: Vec<BrilligOutputs>,
     ) {
         let opcode = AcirOpcode::Brillig(AcvmBrillig {
             inputs,
             outputs,
             foreign_call_results: Vec::new(),
-            bytecode: code,
+            bytecode: generated_brillig.byte_code,
             predicate,
         });
         self.push_opcode(opcode);
+        for (brillig_index, call_stack) in generated_brillig.locations {
+            self.locations.insert(
+                OpcodeLocation::Brillig { acir_index: self.opcodes.len() - 1, brillig_index },
+                call_stack,
+            );
+        }
     }
 
     /// Generate gates and control bits witnesses which ensure that out_expr is a permutation of in_expr