From 92e532eda686dfe6be6d5610f4471d3dd4b3899f Mon Sep 17 00:00:00 2001
From: Kevin Jue <kjue235@gmail.com>
Date: Thu, 8 Feb 2024 13:49:13 -0800
Subject: [PATCH] unconstrained

---
 baby-bear/Cargo.toml         |  1 +
 poseidon2/Cargo.toml         |  1 +
 poseidon2/src/babybear.rs    | 29 +++++++++++++++++++++++++----
 symmetric/src/compression.rs |  8 ++++----
 4 files changed, 31 insertions(+), 8 deletions(-)

diff --git a/baby-bear/Cargo.toml b/baby-bear/Cargo.toml
index 6703a742..6de9c84e 100644
--- a/baby-bear/Cargo.toml
+++ b/baby-bear/Cargo.toml
@@ -9,6 +9,7 @@ p3-field = { path = "../field" }
 rand = { version = "0.8.5", optional = true }
 serde = { version = "1.0", default-features = false, features = ["derive"] }
 lazy_static = "1.4"
+succinct-zkvm = { git="https://github.com/succinctlabs/vm.git", branch="kevjue/recursive_verifier_profiler" }
 
 [dev-dependencies]
 p3-field-testing = { path = "../field-testing" }
diff --git a/poseidon2/Cargo.toml b/poseidon2/Cargo.toml
index 134a8343..cdc58d19 100644
--- a/poseidon2/Cargo.toml
+++ b/poseidon2/Cargo.toml
@@ -12,6 +12,7 @@ p3-field = { path = "../field" }
 p3-mds = { path = "../mds" }
 p3-symmetric = { path = "../symmetric" }
 rand = { version = "0.8.5", optional = true }
+succinct-zkvm = { git="https://github.com/succinctlabs/vm.git", branch="kevjue/recursive_verifier_profiler" }
 
 [features]
 rand = ["dep:rand"]
diff --git a/poseidon2/src/babybear.rs b/poseidon2/src/babybear.rs
index d75632c7..3aeaddbd 100644
--- a/poseidon2/src/babybear.rs
+++ b/poseidon2/src/babybear.rs
@@ -4,9 +4,11 @@
 
 use p3_baby_bear::BabyBear;
 // use p3_baby_bear::IN_HASH;
-use p3_field::AbstractField;
+use p3_field::{AbstractField, PrimeField32};
 use p3_symmetric::Permutation;
 
+use succinct_zkvm::{io, unconstrained};
+
 use crate::diffusion::matmul_internal;
 use crate::DiffusionPermutation;
 
@@ -24,13 +26,32 @@ pub const MATRIX_DIAG_24_BABYBEAR: [u64; 24] = [
 #[derive(Debug, Clone, Default)]
 pub struct DiffusionMatrixBabybear;
 
-impl<AF: AbstractField<F = BabyBear>> Permutation<[AF; 16]> for DiffusionMatrixBabybear {
+impl<AF: PrimeField32> Permutation<[AF; 16]> for DiffusionMatrixBabybear {
     fn permute_mut(&self, state: &mut [AF; 16]) {
         // let mut in_hash = IN_HASH.lock().unwrap();
         // *in_hash = true;
         // drop(in_hash);
         // println!("cycle-tracker-start: permute_mut matmul_internal");
-        matmul_internal::<AF, 16>(state, MATRIX_DIAG_16_BABYBEAR);
+
+        unconstrained! {
+            let mut new_state: [AF;16] = [AF::default(); 16];
+            new_state.copy_from_slice(state);
+            matmul_internal::<AF, 16>(&mut new_state, MATRIX_DIAG_16_BABYBEAR);
+            let bytes = state.map(|x| x.as_canonical_u32().to_le_bytes());
+            let mut flat_bytes = Vec::new();
+            for i in 0..16 {
+                flat_bytes.extend_from_slice(&bytes[i]);
+            }
+            io::hint_slice(&flat_bytes);
+        }
+
+        let mut bytes: [u8; 64] = [0; 64];
+        io::read_slice(&mut bytes);
+        let ret = bytes.chunks(4).map(|chunk| AF::from_canonical_u32(u32::from_le_bytes(chunk.try_into().unwrap()))).collect::<Vec<AF>>();
+        for i in 0..16 {
+            state[i] = ret[i];
+        }
+
         // println!("cycle-tracker-end: permute_mut matmul_internal");
         // let mut in_hash = IN_HASH.lock().unwrap();
         // *in_hash = false;
@@ -38,7 +59,7 @@ impl<AF: AbstractField<F = BabyBear>> Permutation<[AF; 16]> for DiffusionMatrixB
     }
 }
 
-impl<AF: AbstractField<F = BabyBear>> DiffusionPermutation<AF, 16> for DiffusionMatrixBabybear {}
+impl<AF: PrimeField32> DiffusionPermutation<AF, 16> for DiffusionMatrixBabybear {}
 
 impl<AF: AbstractField<F = BabyBear>> Permutation<[AF; 24]> for DiffusionMatrixBabybear {
     fn permute_mut(&self, state: &mut [AF; 24]) {
diff --git a/symmetric/src/compression.rs b/symmetric/src/compression.rs
index 31793081..da3091af 100644
--- a/symmetric/src/compression.rs
+++ b/symmetric/src/compression.rs
@@ -33,17 +33,17 @@ where
     InnerP: CryptographicPermutation<[T; WIDTH]>,
 {
     fn compress(&self, input: [[T; CHUNK]; N]) -> [T; CHUNK] {
-        println!("cycle-tracker-start: compress");
+        // println!("cycle-tracker-start: compress");
         debug_assert!(CHUNK * N <= WIDTH);
         let mut pre = [T::default(); WIDTH];
-        println!("cycle-tracker-start: compress_copy_from_slice");
+        // println!("cycle-tracker-start: compress_copy_from_slice");
         for i in 0..N {
             pre[i * CHUNK..(i + 1) * CHUNK].copy_from_slice(&input[i]);
         }
-        println!("cycle-tracker-end: compress_copy_from_slice");
+        // println!("cycle-tracker-end: compress_copy_from_slice");
         let post = self.inner_permutation.permute(pre);
         let ret = post[..CHUNK].try_into().unwrap();
-        println!("cycle-tracker-end: compress");
+        // println!("cycle-tracker-end: compress");
         ret
     }
 }