From 881b6c0d8fbc9e2838bc8eaa08639798170dd346 Mon Sep 17 00:00:00 2001
From: Jan Ferdinand Sauer <ferdinand@neptune.cash>
Date: Fri, 1 Dec 2023 15:43:21 +0100
Subject: [PATCH] fix: fail Sponge instructions if Sponge state is
 uninitialized

---
 triton-vm/src/error.rs                 |  3 +++
 triton-vm/src/instruction.rs           |  2 +-
 triton-vm/src/table/processor_table.rs |  2 +-
 triton-vm/src/vm.rs                    | 32 +++++++++++++++-----------
 4 files changed, 24 insertions(+), 15 deletions(-)

diff --git a/triton-vm/src/error.rs b/triton-vm/src/error.rs
index ea4dba23..ba4abfda 100644
--- a/triton-vm/src/error.rs
+++ b/triton-vm/src/error.rs
@@ -76,6 +76,9 @@ pub enum InstructionError {
     #[error("division by 0 is impossible")]
     DivisionByZero,
 
+    #[error("the Sponge state must be initialized before it can be used")]
+    SpongeNotInitialized,
+
     #[error("the logarithm of 0 does not exist")]
     LogarithmOfZero,
 
diff --git a/triton-vm/src/instruction.rs b/triton-vm/src/instruction.rs
index c8dfeec6..19674e56 100644
--- a/triton-vm/src/instruction.rs
+++ b/triton-vm/src/instruction.rs
@@ -964,7 +964,7 @@ mod tests {
     ) -> (Program, usize) {
         let num_push_instructions = 10;
         let push_instructions = triton_asm![push 1; num_push_instructions];
-        let program = triton_program!({&push_instructions} {test_instruction} nop halt);
+        let program = triton_program!(sponge_init {&push_instructions} {test_instruction} nop halt);
 
         let stack_size_when_reaching_test_instruction =
             NUM_OP_STACK_REGISTERS + num_push_instructions;
diff --git a/triton-vm/src/table/processor_table.rs b/triton-vm/src/table/processor_table.rs
index cc4bc73d..c52012ee 100644
--- a/triton-vm/src/table/processor_table.rs
+++ b/triton-vm/src/table/processor_table.rs
@@ -3396,7 +3396,7 @@ pub(crate) mod tests {
     #[test]
     /// helps identifying whether the printing causes an infinite loop
     fn print_simple_processor_table_row() {
-        let program = triton_program!(push 2 assert halt);
+        let program = triton_program!(push 2 sponge_init assert halt);
         let err = program.run([].into(), [].into()).unwrap_err();
         println!("\n{}", err.vm_state);
     }
diff --git a/triton-vm/src/vm.rs b/triton-vm/src/vm.rs
index b66430ed..5cb976bc 100644
--- a/triton-vm/src/vm.rs
+++ b/triton-vm/src/vm.rs
@@ -78,7 +78,7 @@ pub struct VMState {
     /// Sponge state.
     /// Note that this is the _full_ state, including capacity. The capacity should never be
     /// exposed outside of the VM.
-    pub sponge_state: [BFieldElement; tip5::STATE_SIZE],
+    pub sponge_state: Option<[BFieldElement; tip5::STATE_SIZE]>,
 
     /// Indicates whether the terminating instruction `halt` has been executed.
     pub halting: bool,
@@ -453,18 +453,19 @@ impl VMState {
     }
 
     fn sponge_init(&mut self) -> Vec<CoProcessorCall> {
-        self.sponge_state = Tip5::init().state;
+        self.sponge_state = Some(Tip5::init().state);
         self.instruction_pointer += 1;
         vec![SpongeStateReset]
     }
 
     fn sponge_absorb(&mut self) -> Result<Vec<CoProcessorCall>> {
+        let Some(mut state) = self.sponge_state else {
+            return Err(SpongeNotInitialized);
+        };
         let to_absorb = self.op_stack.pop_multiple::<{ tip5::RATE }>()?;
-        self.sponge_state[..tip5::RATE].copy_from_slice(&to_absorb);
-        let tip5_trace = Tip5::trace(&mut Tip5State {
-            state: self.sponge_state,
-        });
-        self.sponge_state = tip5_trace.last().unwrap().to_owned();
+        state[..tip5::RATE].copy_from_slice(&to_absorb);
+        let tip5_trace = Tip5::trace(&mut Tip5State { state });
+        self.sponge_state = Some(tip5_trace.last().unwrap().to_owned());
 
         let co_processor_calls = vec![Tip5Trace(SpongeAbsorb, Box::new(tip5_trace))];
 
@@ -473,13 +474,14 @@ impl VMState {
     }
 
     fn sponge_squeeze(&mut self) -> Result<Vec<CoProcessorCall>> {
+        let Some(state) = self.sponge_state else {
+            return Err(SpongeNotInitialized);
+        };
         for i in (0..tip5::RATE).rev() {
-            self.op_stack.push(self.sponge_state[i]);
+            self.op_stack.push(state[i]);
         }
-        let tip5_trace = Tip5::trace(&mut Tip5State {
-            state: self.sponge_state,
-        });
-        self.sponge_state = tip5_trace.last().unwrap().to_owned();
+        let tip5_trace = Tip5::trace(&mut Tip5State { state });
+        self.sponge_state = Some(tip5_trace.last().unwrap().to_owned());
 
         let co_processor_calls = vec![Tip5Trace(SpongeSqueeze, Box::new(tip5_trace))];
 
@@ -973,7 +975,11 @@ impl Display for VMState {
             .join(" | ");
         print_row(f, format!("ib6-0:    [ {ib_registers} ]",))?;
 
-        let sponge_state_register = |i: usize| self.sponge_state[i].value();
+        let Some(sponge_state) = self.sponge_state else {
+            return writeln!(f, "╰─{:─<total_width$}─╯", "");
+        };
+
+        let sponge_state_register = |i: usize| sponge_state[i].value();
         let sponge_state_slice = |idxs: [usize; 4]| {
             idxs.map(sponge_state_register)
                 .map(|ss| format!("{ss:>register_width$}"))