remove read_sponge_state columns (#13)

Author: kunxian-xia

extensions/native/circuit/src/poseidon2/air.rs

@@ -1,5 +1,6 @@
 use std::{array::from_fn, borrow::Borrow, sync::Arc};
 
+use itertools::Itertools;
 use openvm_circuit::{
     arch::{ExecutionBridge, ExecutionState},
     system::memory::{offline_checker::MemoryBridge, MemoryAddress, CHUNK},
@@ -724,7 +725,6 @@ impl<AB: InteractionBuilder, const SBOX_REGISTERS: usize> Air<AB>
             write_data,
             data,
             should_permute,
-            read_sponge_state,
             write_sponge_state,
             write_final_idx,
             final_idx,
@@ -856,29 +856,30 @@ impl<AB: InteractionBuilder, const SBOX_REGISTERS: usize> Air<AB>
                 end_idx,
             );
 
-        let full_sponge_input = from_fn::<_, { CHUNK * 2 }, _>(|i| local.inner.inputs[i]);
         let full_sponge_output = from_fn::<_, { CHUNK * 2 }, _>(|i| {
             local.inner.ending_full_rounds[BABY_BEAR_POSEIDON2_HALF_FULL_ROUNDS - 1].post[i]
         });
 
-        self.memory_bridge
-            .read(
-                MemoryAddress::new(self.address_space, state_ptr),
-                full_sponge_input,
-                start_timestamp + (end_idx - start_idx) * AB::F::TWO,
-                &read_sponge_state,
-            )
-            .eval(builder, multi_observe_row * should_permute);
-
         self.memory_bridge
             .write(
                 MemoryAddress::new(self.address_space, state_ptr),
                 full_sponge_output,
-                start_timestamp + (end_idx - start_idx) * AB::F::TWO + AB::F::ONE,
+                start_timestamp + (end_idx - start_idx) * AB::F::TWO,
                 &write_sponge_state,
             )
             .eval(builder, multi_observe_row * should_permute);
 
+        // enforce that prev_data is permutation input
+        write_sponge_state
+            .prev_data()
+            .iter()
+            .zip_eq(local.inner.inputs.iter())
+            .for_each(|(a, b)| {
+                builder
+                    .when(multi_observe_row * should_permute)
+                    .assert_eq(*a, *b);
+            });
+
         /*
         self.memory_bridge
             .write(
@@ -897,7 +898,10 @@ impl<AB: InteractionBuilder, const SBOX_REGISTERS: usize> Air<AB>
         builder
             .when(next.multi_observe_row)
             .when(not(next_multi_observe_specific.is_first))
-            .assert_eq(next_multi_observe_specific.curr_len, multi_observe_specific.curr_len + end_idx - start_idx);
+            .assert_eq(
+                next_multi_observe_specific.curr_len,
+                multi_observe_specific.curr_len + end_idx - start_idx,
+            );
 
         // Boundary conditions
         builder

extensions/native/circuit/src/poseidon2/chip.rs

@@ -676,7 +676,7 @@ where
                 if len >= (CHUNK - pos) {
                     chunks.push((pos.clone(), CHUNK.clone()));
                     len -= CHUNK - pos;
-                    final_timestamp_inc += 2 * (CHUNK - pos + 1);
+                    final_timestamp_inc += 2 * (CHUNK - pos) + 1;
                     pos = 0;
                 } else {
                     chunks.push((pos.clone(), pos + len));
@@ -766,11 +766,7 @@ where
                 multi_observe_cols.end_idx = F::from_canonical_usize(chunk_end);
 
                 multi_observe_cols.is_first = F::ZERO;
-                multi_observe_cols.is_last = if i == num_chunks - 1 {
-                    F::ONE
-                } else {
-                    F::ZERO
-                };
+                multi_observe_cols.is_last = if i == num_chunks - 1 { F::ONE } else { F::ZERO };
                 multi_observe_cols.curr_len = F::from_canonical_usize(input_idx);
 
                 for j in chunk_start..CHUNK {
@@ -796,13 +792,10 @@ where
                     cur_timestamp += 2;
                 }
 
+                let permutation_input: [F; 16] =
+                    memory_read_native(state.memory.data(), state_ptr_u32);
                 if chunk_end >= CHUNK {
                     multi_observe_cols.should_permute = F::ONE;
-                    let permutation_input: [F; 16] = tracing_read_native_helper(
-                        state.memory,
-                        state_ptr_u32,
-                        multi_observe_cols.read_sponge_state.as_mut(),
-                    );
                     cols.inner.inputs.clone_from_slice(&permutation_input);
                     let output = self.subchip.permute(permutation_input);
                     tracing_write_native_inplace(
@@ -811,12 +804,10 @@ where
                         std::array::from_fn(|i| output[i]),
                         &mut multi_observe_cols.write_sponge_state,
                     );
-                    cur_timestamp += 2;
+                    cur_timestamp += 1;
                 } else {
                     multi_observe_cols.should_permute = F::ZERO;
-                    let sponge_state: [F; 16] =
-                        memory_read_native(state.memory.data(), state_ptr_u32);
-                    cols.inner.inputs.clone_from_slice(&sponge_state);
+                    cols.inner.inputs.clone_from_slice(&permutation_input);
                 }
             }
         } else {
@@ -1219,11 +1210,6 @@ impl<F: PrimeField32, const SBOX_REGISTERS: usize> NativePoseidon2Filler<F, SBOX
                 mem_fill_helper(
                     mem_helper,
                     start_timestamp_u32,
-                    multi_observe_cols.read_sponge_state.as_mut(),
-                );
-                mem_fill_helper(
-                    mem_helper,
-                    start_timestamp_u32 + 1,
                     multi_observe_cols.write_sponge_state.as_mut(),
                 );
             }

extensions/native/circuit/src/poseidon2/columns.rs

@@ -237,7 +237,6 @@ pub struct MultiObserveCols<T> {
 
     // Permutation
     pub should_permute: T,
-    pub read_sponge_state: MemoryReadAuxCols<T>,
     pub write_sponge_state: MemoryWriteAuxCols<T, { CHUNK * 2 }>,
 
     // Final write back and registers