feat: Allow step debugging inside Brillig blocks

acvm-repo/acvm/src/pwg/brillig.rs

@@ -111,11 +111,20 @@ impl<'b, B: BlackBoxFunctionSolver> BrilligSolver<'b, B> {
         Ok(Self { vm, acir_index })
     }
 
+    pub(super) fn program_counter(&self) -> usize {
+        self.vm.program_counter()
+    }
+
     pub(super) fn solve(&mut self) -> Result<BrilligSolverStatus, OpcodeResolutionError> {
         let status = self.vm.process_opcodes();
         self.handle_vm_status(status)
     }
 
+    pub(super) fn step(&mut self) -> Result<BrilligSolverStatus, OpcodeResolutionError> {
+        let status = self.vm.process_opcode();
+        self.handle_vm_status(status)
+    }
+
     fn handle_vm_status(
         &self,
         vm_status: VMStatus,

acvm-repo/acvm/src/pwg/mod.rs

@@ -177,6 +177,21 @@ impl<'a, B: BlackBoxFunctionSolver> ACVM<'a, B> {
         self.instruction_pointer
     }
 
+    /// Returns the location for the next opcode to execute, or None if execution is finished
+    pub fn location(&self) -> Option<OpcodeLocation> {
+        if self.instruction_pointer >= self.opcodes.len() {
+            // evaluation finished
+            None
+        } else if let Some(solver) = &self.brillig_solver {
+            Some(OpcodeLocation::Brillig {
+                acir_index: self.instruction_pointer,
+                brillig_index: solver.program_counter(),
+            })
+        } else {
+            Some(OpcodeLocation::Acir(self.instruction_pointer))
+        }
+    }
+
     /// Finalize the ACVM execution, returning the resulting [`WitnessMap`].
     pub fn finalize(self) -> WitnessMap {
         if self.status != ACVMStatus::Solved {
@@ -242,6 +257,10 @@ impl<'a, B: BlackBoxFunctionSolver> ACVM<'a, B> {
     }
 
     pub fn solve_opcode(&mut self) -> ACVMStatus {
+        self.step_opcode(false)
+    }
+
+    pub fn step_opcode(&mut self, step_into_brillig: bool) -> ACVMStatus {
         let opcode = &self.opcodes[self.instruction_pointer];
 
         let resolution = match opcode {
@@ -258,11 +277,15 @@ impl<'a, B: BlackBoxFunctionSolver> ACVM<'a, B> {
                 let solver = self.block_solvers.entry(*block_id).or_default();
                 solver.solve_memory_op(op, &mut self.witness_map, predicate)
             }
-            Opcode::Brillig(_) => match self.solve_brillig_opcode() {
-                Ok(Some(foreign_call)) => return self.wait_for_foreign_call(foreign_call),
+            Opcode::Brillig(_) => match self.step_brillig_opcode(step_into_brillig) {
+                Ok(BrilligSolverStatus::ForeignCallWait(foreign_call)) => {
+                    return self.wait_for_foreign_call(foreign_call)
+                }
+                Ok(BrilligSolverStatus::InProgress) => return self.status(ACVMStatus::InProgress),
                 res => res.map(|_| ()),
             },
         };
+
         match resolution {
             Ok(()) => {
                 self.instruction_pointer += 1;
@@ -296,40 +319,49 @@ impl<'a, B: BlackBoxFunctionSolver> ACVM<'a, B> {
         }
     }
 
-    fn solve_brillig_opcode(
+    fn step_brillig_opcode(
         &mut self,
-    ) -> Result<Option<ForeignCallWaitInfo>, OpcodeResolutionError> {
+        step_into: bool,
+    ) -> Result<BrilligSolverStatus, OpcodeResolutionError> {
         let Opcode::Brillig(brillig) = &self.opcodes[self.instruction_pointer] else {
             unreachable!("Not executing a Brillig opcode");
         };
         let witness = &mut self.witness_map;
-        if BrilligSolver::<B>::should_skip(witness, brillig)? {
-            BrilligSolver::<B>::zero_out_brillig_outputs(witness, brillig).map(|_| None)
-        } else {
-            // If we're resuming execution after resolving a foreign call then
-            // there will be a cached `BrilligSolver` to avoid recomputation.
-            let mut solver: BrilligSolver<'_, B> = match self.brillig_solver.take() {
-                Some(solver) => solver,
-                None => {
-                    BrilligSolver::new(witness, brillig, self.backend, self.instruction_pointer)?
-                }
-            };
-            match solver.solve()? {
-                BrilligSolverStatus::ForeignCallWait(foreign_call) => {
-                    // Cache the current state of the solver
-                    self.brillig_solver = Some(solver);
-                    Ok(Some(foreign_call))
+
+        // Try to use the cached `BrilligSolver` which we will have if:
+        // - stepping into a Brillig block
+        // - resuming execution from a foreign call
+        let mut solver: BrilligSolver<'_, B> = match self.brillig_solver.take() {
+            Some(solver) => solver,
+            None => {
+                if BrilligSolver::<B>::should_skip(witness, brillig)? {
+                    // Exit early if the block doesn't need to be executed (false predicate)
+                    return BrilligSolver::<B>::zero_out_brillig_outputs(witness, brillig)
+                        .map(|_| BrilligSolverStatus::Finished);
                 }
-                BrilligSolverStatus::InProgress => {
+                BrilligSolver::new(witness, brillig, self.backend, self.instruction_pointer)?
+            }
+        };
+
+        let status = if step_into { solver.step()? } else { solver.solve()? };
+        match status {
+            BrilligSolverStatus::ForeignCallWait(_) => {
+                // Cache the current state of the solver
+                self.brillig_solver = Some(solver);
+            }
+            BrilligSolverStatus::InProgress => {
+                if !step_into {
                     unreachable!("Brillig solver still in progress")
                 }
-                BrilligSolverStatus::Finished => {
-                    // Write execution outputs
-                    solver.finalize(witness, brillig)?;
-                    Ok(None)
-                }
+                // Cache the current state of the solver
+                self.brillig_solver = Some(solver);
             }
-        }
+            BrilligSolverStatus::Finished => {
+                // Write execution outputs
+                solver.finalize(witness, brillig)?;
+            }
+        };
+        Ok(status)
     }
 }
 

tooling/debugger/src/lib.rs

@@ -29,9 +29,13 @@ struct DebugContext<'backend, B: BlackBoxFunctionSolver> {
 
 impl<'backend, B: BlackBoxFunctionSolver> DebugContext<'backend, B> {
     fn step_opcode(&mut self) -> Result<SolveResult, NargoError> {
-        let solver_status = self.acvm.solve_opcode();
+        let solver_status = self.acvm.step_opcode(true);
 
-        match solver_status {
+        self.handle_acvm_status(solver_status)
+    }
+
+    fn handle_acvm_status(&mut self, status: ACVMStatus) -> Result<SolveResult, NargoError> {
+        match status {
             ACVMStatus::Solved => Ok(SolveResult::Done),
             ACVMStatus::InProgress => Ok(SolveResult::Ok),
             ACVMStatus::Failure(error) => {
@@ -68,13 +72,23 @@ impl<'backend, B: BlackBoxFunctionSolver> DebugContext<'backend, B> {
     }
 
     fn show_current_vm_status(&self) {
-        let ip = self.acvm.instruction_pointer();
+        let location = self.acvm.location();
         let opcodes = self.acvm.opcodes();
-        if ip >= opcodes.len() {
-            println!("Finished execution");
-        } else {
-            println!("Stopped at opcode {}: {}", ip, opcodes[ip]);
-            Self::show_source_code_location(&OpcodeLocation::Acir(ip), &self.debug_artifact);
+
+        match location {
+            None => println!("Finished execution"),
+            Some(location) => {
+                match location {
+                    OpcodeLocation::Acir(ip) => {
+                        println!("Stopped at opcode {}: {}", ip, opcodes[ip])
+                    }
+                    OpcodeLocation::Brillig { acir_index: ip, brillig_index } => println!(
+                        "Stopped at opcode {} in Brillig block {}: {}",
+                        brillig_index, ip, opcodes[ip]
+                    ),
+                }
+                Self::show_source_code_location(&location, &self.debug_artifact);
+            }
         }
     }