write final pos back

Author: kunxian-xia

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

@@ -728,7 +728,6 @@ impl<AB: InteractionBuilder, const SBOX_REGISTERS: usize> Air<AB>
             should_permute,
             write_sponge_state,
             write_final_idx,
-            final_idx,
             input_register_1,
             input_register_2,
             input_register_3,
@@ -830,6 +829,16 @@ impl<AB: InteractionBuilder, const SBOX_REGISTERS: usize> Air<AB>
         }
 
         for i in 0..CHUNK {
+            builder
+                .when(multi_observe_row)
+                .assert_bool(aux_after_start[i]);
+            builder
+                .when(multi_observe_row)
+                .assert_bool(aux_before_end[i]);
+            builder
+                .when(multi_observe_row)
+                .when(is_first)
+                .assert_zero(aux_read_enabled[i]);
             builder
                 .when(multi_observe_row)
                 .assert_eq(aux_after_start[i] * aux_before_end[i], aux_read_enabled[i]);
@@ -889,19 +898,22 @@ impl<AB: InteractionBuilder, const SBOX_REGISTERS: usize> Air<AB>
                     .assert_eq(*a, *b);
             });
 
-        /*
+        builder
+            .when(multi_observe_row)
+            .when(aux_read_enabled[CHUNK - 1])
+            .assert_one(should_permute);
+
+        // final_idx = aux_read_enabled[CHUNK-1] * 0 + (1 - aux_read_enabled[CHUNK-1]) * end_idx
+        let final_idx = aux_read_enabled[CHUNK - 1] * AB::Expr::ZERO
+            + (AB::Expr::ONE - aux_read_enabled[CHUNK - 1]) * end_idx;
         self.memory_bridge
             .write(
-                MemoryAddress::new(
-                    self.address_space,
-                    input_register_1,
-                ),
+                MemoryAddress::new(self.address_space, input_register_1),
                 [final_idx],
-                start_timestamp + is_first * AB::F::from_canonical_usize(4) + (end_idx - start_idx) * AB::F::TWO + should_permute * AB::F::TWO,
-                &write_final_idx
+                start_timestamp + (end_idx - start_idx) * AB::F::TWO + should_permute,
+                &write_final_idx,
             )
             .eval(builder, multi_observe_row * is_last);
-        */
 
         // Field transitions
         builder

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

@@ -685,6 +685,7 @@ where
                     pos += len;
                 }
             }
+            final_timestamp_inc += 1; // write back to init_pos_register
 
             let allocated_rows = arena
                 .alloc(MultiRowLayout::new(NativePoseidon2Metadata {
@@ -810,6 +811,15 @@ where
                     multi_observe_cols.should_permute = F::ZERO;
                     cols.inner.inputs.clone_from_slice(&permutation_input);
                 }
+                if i == num_chunks - 1 {
+                    let final_idx = F::from_canonical_usize(chunk_end % CHUNK);
+                    tracing_write_native_inplace(
+                        state.memory,
+                        init_pos_register.as_canonical_u32(),
+                        [final_idx],
+                        &mut multi_observe_cols.write_final_idx,
+                    );
+                }
             }
         } else {
             unreachable!()
@@ -1213,6 +1223,14 @@ impl<F: PrimeField32, const SBOX_REGISTERS: usize> NativePoseidon2Filler<F, SBOX
                     start_timestamp_u32,
                     multi_observe_cols.write_sponge_state.as_mut(),
                 );
+                start_timestamp_u32 += 1;
+            }
+            if row_idx == num_rows - 1 {
+                mem_fill_helper(
+                    mem_helper,
+                    start_timestamp_u32,
+                    multi_observe_cols.write_final_idx.as_mut(),
+                );
             }
         }
     }

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

@@ -242,5 +242,4 @@ pub struct MultiObserveCols<T> {
 
     // Final write back and registers
     pub write_final_idx: MemoryWriteAuxCols<T, 1>,
-    pub final_idx: T,
 }

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

@@ -331,8 +331,9 @@ impl<F: PrimeField32, const SBOX_REGISTERS: usize> MeteredExecutor<F>
 {
     #[inline(always)]
     fn metered_pre_compute_size(&self) -> usize {
-        std::cmp::max(
+        max3(
             size_of::<E2PreCompute<Pos2PreCompute<F, SBOX_REGISTERS>>>(),
+            size_of::<E2PreCompute<MultiObservePreCompute<F, SBOX_REGISTERS>>>(),
             size_of::<E2PreCompute<VerifyBatchPreCompute<F, SBOX_REGISTERS>>>(),
         )
     }