From 5f702d471eb94c2b1a7943125b4f97f0a5b90c3d Mon Sep 17 00:00:00 2001
From: Jan Ferdinand Sauer <ferdinand@neptune.cash>
Date: Fri, 15 Dec 2023 12:05:54 +0100
Subject: [PATCH] refactor: allow tracing program execution from a given
 starting state

---
 triton-vm/src/program.rs | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/triton-vm/src/program.rs b/triton-vm/src/program.rs
index 0fbc5dc4..b96b2092 100644
--- a/triton-vm/src/program.rs
+++ b/triton-vm/src/program.rs
@@ -425,9 +425,19 @@ impl Program {
         &self,
         public_input: PublicInput,
         non_determinism: NonDeterminism<BFieldElement>,
+    ) -> Result<(AlgebraicExecutionTrace, Vec<BFieldElement>)> {
+        let state = VMState::new(self, public_input, non_determinism);
+        self.trace_execution_of_state(state)
+    }
+
+    /// Trace the execution of a [`Program`] from a given [`VMState`]. Consider using
+    /// [`trace_execution`][Self::trace_execution], unless you know this is what you want.
+    pub fn trace_execution_of_state(
+        &self,
+        mut state: VMState,
     ) -> Result<(AlgebraicExecutionTrace, Vec<BFieldElement>)> {
         let mut aet = AlgebraicExecutionTrace::new(self.clone());
-        let mut state = VMState::new(self, public_input, non_determinism);
+        assert_eq!(self.instructions, state.program);
         assert_eq!(self.len_bwords(), aet.instruction_multiplicities.len());
 
         while !state.halting {