feat: Extract Brillig VM to allow step debugging (#3259)

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

@@ -14,13 +14,14 @@ use crate::{pwg::OpcodeNotSolvable, OpcodeResolutionError};
 
 use super::{get_value, insert_value};
 
-pub(super) enum BrilligSolverStatus {
+#[derive(Debug)]
+pub enum BrilligSolverStatus {
     Finished,
     InProgress,
     ForeignCallWait(ForeignCallWaitInfo),
 }
 
-pub(super) struct BrilligSolver<'b, B: BlackBoxFunctionSolver> {
+pub struct BrilligSolver<'b, B: BlackBoxFunctionSolver> {
     vm: VM<'b, B>,
     acir_index: usize,
 }
@@ -62,7 +63,7 @@ impl<'b, B: BlackBoxFunctionSolver> BrilligSolver<'b, B> {
     /// Constructs a solver for a Brillig block given the bytecode and initial
     /// witness.
     pub(super) fn new(
-        initial_witness: &mut WitnessMap,
+        initial_witness: &WitnessMap,
         brillig: &'b Brillig,
         bb_solver: &'b B,
         acir_index: usize,
@@ -116,6 +117,15 @@ impl<'b, B: BlackBoxFunctionSolver> BrilligSolver<'b, B> {
         self.handle_vm_status(status)
     }
 
+    pub fn step(&mut self) -> Result<BrilligSolverStatus, OpcodeResolutionError> {
+        let status = self.vm.process_opcode();
+        self.handle_vm_status(status)
+    }
+
+    pub fn program_counter(&self) -> usize {
+        self.vm.program_counter()
+    }
+
     fn handle_vm_status(
         &self,
         vm_status: VMStatus,
@@ -185,7 +195,7 @@ impl<'b, B: BlackBoxFunctionSolver> BrilligSolver<'b, B> {
         Ok(())
     }
 
-    pub(super) fn resolve_pending_foreign_call(&mut self, foreign_call_result: ForeignCallResult) {
+    pub fn resolve_pending_foreign_call(&mut self, foreign_call_result: ForeignCallResult) {
         match self.vm.get_status() {
             VMStatus::ForeignCallWait { .. } => self.vm.resolve_foreign_call(foreign_call_result),
             _ => unreachable!("Brillig VM is not waiting for a foreign call"),

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

@@ -10,12 +10,7 @@ use acir::{
 };
 use acvm_blackbox_solver::BlackBoxResolutionError;
 
-use self::{
-    arithmetic::ArithmeticSolver,
-    brillig::{BrilligSolver, BrilligSolverStatus},
-    directives::solve_directives,
-    memory_op::MemoryOpSolver,
-};
+use self::{arithmetic::ArithmeticSolver, directives::solve_directives, memory_op::MemoryOpSolver};
 use crate::{BlackBoxFunctionSolver, Language};
 
 use thiserror::Error;
@@ -30,6 +25,7 @@ mod directives;
 mod blackbox;
 mod memory_op;
 
+pub use self::brillig::{BrilligSolver, BrilligSolverStatus};
 pub use brillig::ForeignCallWaitInfo;
 
 #[derive(Debug, Clone, PartialEq)]
@@ -63,6 +59,11 @@ impl std::fmt::Display for ACVMStatus {
     }
 }
 
+pub enum StepResult<'a, B: BlackBoxFunctionSolver> {
+    Status(ACVMStatus),
+    IntoBrillig(BrilligSolver<'a, B>),
+}
+
 // This enum represents the different cases in which an
 // opcode can be unsolvable.
 // The most common being that one of its input has not been
@@ -263,6 +264,13 @@ impl<'a, B: BlackBoxFunctionSolver> ACVM<'a, B> {
                 res => res.map(|_| ()),
             },
         };
+        self.handle_opcode_resolution(resolution)
+    }
+
+    fn handle_opcode_resolution(
+        &mut self,
+        resolution: Result<(), OpcodeResolutionError>,
+    ) -> ACVMStatus {
         match resolution {
             Ok(()) => {
                 self.instruction_pointer += 1;
@@ -302,34 +310,64 @@ impl<'a, B: BlackBoxFunctionSolver> ACVM<'a, B> {
         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))
-                }
-                BrilligSolverStatus::InProgress => {
-                    unreachable!("Brillig solver still in progress")
-                }
-                BrilligSolverStatus::Finished => {
-                    // Write execution outputs
-                    solver.finalize(witness, brillig)?;
-                    Ok(None)
-                }
+            return BrilligSolver::<B>::zero_out_brillig_outputs(witness, brillig).map(|_| None);
+        }
+
+        // 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))
+            }
+            BrilligSolverStatus::InProgress => {
+                unreachable!("Brillig solver still in progress")
             }
+            BrilligSolverStatus::Finished => {
+                // Write execution outputs
+                solver.finalize(witness, brillig)?;
+                Ok(None)
+            }
+        }
+    }
+
+    pub fn step_into_brillig_opcode(&mut self) -> StepResult<'a, B> {
+        let Opcode::Brillig(brillig) = &self.opcodes[self.instruction_pointer] else {
+            return StepResult::Status(self.solve_opcode());
+        };
+
+        let witness = &mut self.witness_map;
+        let should_skip = match BrilligSolver::<B>::should_skip(witness, brillig) {
+            Ok(result) => result,
+            Err(err) => return StepResult::Status(self.handle_opcode_resolution(Err(err))),
+        };
+
+        if should_skip {
+            let resolution = BrilligSolver::<B>::zero_out_brillig_outputs(witness, brillig);
+            return StepResult::Status(self.handle_opcode_resolution(resolution));
+        }
+
+        let solver = BrilligSolver::new(witness, brillig, self.backend, self.instruction_pointer);
+        match solver {
+            Ok(solver) => StepResult::IntoBrillig(solver),
+            Err(..) => StepResult::Status(self.handle_opcode_resolution(solver.map(|_| ()))),
+        }
+    }
+
+    pub fn finish_brillig_with_solver(&mut self, solver: BrilligSolver<'a, B>) -> ACVMStatus {
+        if !matches!(&self.opcodes[self.instruction_pointer], Opcode::Brillig(..)) {
+            unreachable!("Not executing a Brillig opcode");
         }
+        self.brillig_solver = Some(solver);
+        self.solve_opcode()
     }
 }
 

tooling/debugger/src/lib.rs

@@ -1,5 +1,8 @@
 use acvm::acir::circuit::OpcodeLocation;
-use acvm::pwg::{ACVMStatus, ErrorLocation, OpcodeResolutionError, ACVM};
+use acvm::pwg::{
+    ACVMStatus, BrilligSolver, BrilligSolverStatus, ErrorLocation, OpcodeResolutionError,
+    StepResult, ACVM,
+};
 use acvm::BlackBoxFunctionSolver;
 use acvm::{acir::circuit::Circuit, acir::native_types::WitnessMap};
 
@@ -26,17 +29,60 @@ enum SolveResult {
 
 struct DebugContext<'backend, B: BlackBoxFunctionSolver> {
     acvm: ACVM<'backend, B>,
+    brillig_solver: Option<BrilligSolver<'backend, B>>,
     debug_artifact: DebugArtifact,
     foreign_call_executor: ForeignCallExecutor,
     circuit: &'backend Circuit,
     show_output: bool,
 }
 
 impl<'backend, B: BlackBoxFunctionSolver> DebugContext<'backend, B> {
+    fn step_brillig_opcode(&mut self) -> Result<SolveResult, NargoError> {
+        let Some(mut solver) = self.brillig_solver.take() else {
+            unreachable!("Missing Brillig solver");
+        };
+        match solver.step() {
+            Ok(status) => {
+                println!("Brillig step result: {:?}", status);
+                println!("Brillig program counter: {:?}", solver.program_counter());
+                match status {
+                    BrilligSolverStatus::InProgress => {
+                        self.brillig_solver = Some(solver);
+                        Ok(SolveResult::Ok)
+                    }
+                    BrilligSolverStatus::Finished => {
+                        let status = self.acvm.finish_brillig_with_solver(solver);
+                        self.handle_acvm_status(status)
+                    }
+                    BrilligSolverStatus::ForeignCallWait(foreign_call) => {
+                        let foreign_call_result =
+                            self.foreign_call_executor.execute(&foreign_call, self.show_output)?;
+                        solver.resolve_pending_foreign_call(foreign_call_result);
+                        self.brillig_solver = Some(solver);
+                        Ok(SolveResult::Ok)
+                    }
+                }
+            }
+            Err(err) => self.handle_acvm_status(ACVMStatus::Failure(err)),
+        }
+    }
+
     fn step_opcode(&mut self) -> Result<SolveResult, NargoError> {
-        let solver_status = self.acvm.solve_opcode();
+        if matches!(self.brillig_solver, Some(_)) {
+            self.step_brillig_opcode()
+        } else {
+            match self.acvm.step_into_brillig_opcode() {
+                StepResult::IntoBrillig(solver) => {
+                    self.brillig_solver = Some(solver);
+                    self.step_brillig_opcode()
+                }
+                StepResult::Status(status) => self.handle_acvm_status(status),
+            }
+        }
+    }
 
-        match 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) => {
@@ -189,6 +235,7 @@ pub fn debug_circuit<B: BlackBoxFunctionSolver>(
     let context = RefCell::new(DebugContext {
         acvm: ACVM::new(blackbox_solver, &circuit.opcodes, initial_witness),
         foreign_call_executor: ForeignCallExecutor::default(),
+        brillig_solver: None,
         circuit,
         debug_artifact,
         show_output,