gnosisguild · ozgurarmanc · May 27, 2025 · May 16, 2025 · May 18, 2025 · May 18, 2025
diff --git a/circuits/src/lib.nr b/circuits/src/lib.nr
@@ -1,2 +1,3 @@
 pub mod polynomial;
 pub mod pk_encryption;
+pub mod safe;
diff --git a/circuits/src/pk_encryption.nr b/circuits/src/pk_encryption.nr
@@ -3,7 +3,7 @@ use super::pk_enc_constants_1024_2x52_2048::{
     R1_LOW_BOUNDS, R1_UP_BOUNDS, R2_BOUNDS, U_BOUND,
 };
 use super::polynomial::Polynomial;
-use poseidon::poseidon2::Poseidon2;
+use super::safe::SafeSponge;
 
 // Total size of the all polynomial's coefficients.
 global size: u32 = (10 * L + 4) * N - 8;
@@ -65,15 +65,7 @@ impl BfvPkEncryptionCircuit {
     }
 
     pub fn correct_encryption(self) {
-        // TO
10000
DO: Add domain separation
-        // We assign all the coefficients to an array to generate a challenge value with Poseidon hash function
-        let inputs = self.payload();
-        let gamma = Poseidon2::hash(inputs, size);
-
-        // cyclo poly is equal to x^N + 1
-        let cyclo_at_gamma = gamma.pow_32(N as Field) + 1;
-
-        // Bound checks
+        // Bound check
         self.u.range_check_1bound(U_BOUND);
         self.e0.range_check_1bound(E_BOUND);
         self.e1.range_check_1bound(E_BOUND);
@@ -88,25 +80,31 @@ impl BfvPkEncryptionCircuit {
             self.p2is[i].range_check_1bound(P2_BOUNDS[i]);
         }
 
-        // Gamma evaluation
-        let u_at_gamma = self.u.eval(gamma);
-        let e0_at_gamma = self.e0.eval(gamma);
-        let e1_at_gamma = self.e1.eval(gamma);
-        let k1_at_gamma = self.k1.eval(gamma);
-        let pk0is_at_gamma = self.pk0is.map(|pk| pk.eval(gamma));
-        let pk1is_at_gamma = self.pk1is.map(|pk| pk.eval(gamma));
-        let r1i_at_gamma = self.r1is.map(|r1| r1.eval(gamma));
-        let r2i_at_gamma = self.r2is.map(|r2| r2.eval(gamma));
-        let p1is_at_gamma = self.p1is.map(|p1| p1.eval(gamma));
-        let p2is_at_gamma = self.p2is.map(|p2| p2.eval(gamma));
+        // We assign all the coefficients to an array to generate challenge values
+        let inputs = self.payload();
+        let mut safe = SafeSponge::start([size, 2 * L]);
+        safe = safe.absorb(inputs);
+        let gammas = safe.squeeze();
 
+        // CORRECT ENCRYPTION CONSTRAINT
         // For each `i` Prove that LHS(gamma) = RHS(gamma)
         // pk0_u = pk0is(gamma) * u(gamma) + e0(gamma)
         // LHS = ct0i(gamma)
         // RHS = pk0_u  + k1(gamma) * k0i + r1i(gamma) * qi + r2i(gamma) * cyclo(gamma)
-
-        // CORRECT ENCRYPTION CONSTRAINT
         for i in 0..L {
+            let gamma = gammas.get(i);
+
+            // cyclo poly is equal to x^N + 1
+            let cyclo_at_gamma = gamma.pow_32(N as Field) + 1;
+
+            // Gamma evaluation
+            let u_at_gamma = self.u.eval(gamma);
+            let e0_at_gamma = self.e0.eval(gamma);
+            let k1_at_gamma = self.k1.eval(gamma);
+            let pk0is_at_gamma = self.pk0is.map(|pk| pk.eval(gamma));
+            let r1i_at_gamma = self.r1is.map(|r1| r1.eval(gamma));
+            let r2i_at_gamma = self.r2is.map(|r2| r2.eval(gamma));
+
             // First step
             let pk0_u = (pk0is_at_gamma[i] * u_at_gamma) + e0_at_gamma;
 
@@ -123,6 +121,18 @@ impl BfvPkEncryptionCircuit {
             // LHS(gamma) = RHS(gamma)
             assert_eq(lhs, rhs);
 
+            // Gamma evaluation
+            let gamma = gammas.get(i + L);
+
+            // cyclo poly is equal to x^N + 1
+            let cyclo_at_gamma = gamma.pow_32(N as Field) + 1;
+
+            let u_at_gamma = self.u.eval(gamma);
+            let e1_at_gamma = self.e1.eval(gamma);
+            let pk1is_at_gamma = self.pk1is.map(|pk| pk.eval(gamma));
+            let p1is_at_gamma = self.p1is.map(|p1| p1.eval(gamma));
+            let p2is_at_gamma = self.p2is.map(|p2| p2.eval(gamma));
+
             // Second step
             let pk1_u = (pk1is_at_gamma[i] * u_at_gamma) + e1_at_gamma;
 
@@ -136,6 +146,7 @@ impl BfvPkEncryptionCircuit {
             // LHS(gamma) = RHS(gamma)
             assert_eq(lhs, rhs);
         }
+        safe.finish();
     }
 }
 

diff --git a/circuits/src/safe.nr b/circuits/src/safe.nr
@@ -0,0 +1,81 @@
+use super::pk_enc_constants_1024_2x52_2048::TAG;
+use poseidon::poseidon2_permutation;
+
+global RATE: u32 = 3;
+global CAPACITY: u32 = 1;
+global WIDTH: u32 = 4;
+
+pub struct SafeSponge<let L: u32, let S: u32> {
+    state: [Field; WIDTH],
+    out: Vec<Field>,
+    absorb_pos: u32,
+    squeeze_pos: u32,
+    io_pattern: [u32; L],
+    io_count: u32,
+}
+
+impl<let L: u32, let S: u32> SafeSponge<L, S> {
+    pub fn start(pattern: [u32; L]) -> SafeSponge<L, S> {
+        let mut sponge = SafeSponge::<L, S> {
+            state: [0; WIDTH],
+            out: Vec::new(),
+            absorb_pos: 0,
+            squeeze_pos: 0,
+            io_pattern: pattern,
+            io_count: 0,
+        };
+        sponge.state[0] = TAG;
+        sponge
+    }
+
+    pub fn absorb(mut self, input: [Field; S]) -> SafeSponge<L, S> {
+        assert(self.io_pattern[self.io_count] as u32 == S);
+
+        for i in 0..self.io_pattern[self.io_count] {
+            if self.absorb_pos == RATE {
+                self.state = poseidon2_permutation(self.state, self.state.len());
+                self.absorb_pos = 0;
+            }
+            let pos = self.absorb_pos + CAPACITY;
+            self.state[pos] = self.state[pos] + input[i];
+            self.absorb_pos += 1;
+        }
+        self.io_count += 1;
+        self.squeeze_pos = RATE;
+        self
+    }
+
+    pub fn squeeze(mut self) -> Vec<Field> {
+        for _ in 0..self.io_pattern[self.io_count] {
+            if self.squeeze_pos == RATE {
+                self.state = poseidon2_permutation(self.state, self.state.len());
+                self.squeeze_pos = 0;
+                self.absorb_pos = 0;
+            }
+            self.out.push(self.state[self.squeeze_pos + CAPACITY]);
+            self.squeeze_pos += 1;
+        }
+        self.io_count += 1;
+        self.out
+    }
+
+    pub fn finish(mut self) {
+        // Clear the state
+        self.state = [0; WIDTH];
+        self.out = Vec::new();
+        self.io_count = 0;
+        self.io_pattern = [0; L];
+        self.squeeze_pos = 0;
+        self.absorb_pos = 0;
+    }
+}
+
+#[test]
+fn test_safe() {
+    let pattern = [5, 4];
+    let mut safe = SafeSponge::start(pattern);
+    safe = safe.absorb([1, 2, 3, 4, 5]);
+    let res = safe.squeeze();
+    safe.finish();
+    assert_eq(res.get(0), 0x0ef061c5b88d5d81e9c7c306701e01e03c5040a3d3c4b051ada3f150b44dc595);
+}
diff --git a/rs-script/Cargo.lock b/rs-script/Cargo.lock
diff --git a/rs-script/Cargo.toml b/rs-script/Cargo.toml
@@ -17,4 +17,5 @@ num-traits = "0.2"
 ndarray = "0.15.6"
 itertools = "0.13.0"
 rayon = "1.10.0"
-toml = "0.8"
+toml = "0.8"
+blake3 = "1.8.2"
diff --git a/rs-script/src/main.rs b/rs-script/src/main.rs
@@ -1,5 +1,6 @@
 mod poly;
 
+use blake3::Hasher;
 use fhe::bfv::{
     BfvParameters, BfvParametersBuilder, Ciphertext, Encoding, Plaintext, PublicKey, SecretKey,
 };
@@ -8,20 +9,20 @@ use fhe_math::{
     zq::Modulus,
 };
 use fhe_traits::*;
-use itertools::izip;
+use itertools::{izip, Itertools};
 
-use num_bigint::BigInt;
-use num_traits::{Num, Signed, ToPrimitive, Zero};
+use num_bigint::{BigInt, BigUint, Sign};
+use num_traits::{FromPrimitive, Num, Signed, ToPrimitive, Zero};
 use rand::rngs::StdRng;
 use rand::SeedableRng;
 use rayon::iter::{ParallelBridge, ParallelIterator};
 use serde_json::json;
-use std::fs::File;
 use std::io::Write;
 use std::ops::Deref;
 use std::path::Path;
 use std::sync::Arc;
 use std::vec;
+use std::{fs::File, str::FromStr};
 
 use poly::*;
 
@@ -833,6 +834,43 @@ impl InputValidationBounds {
             k0i_constants
         )?;
 
+        let pk_bound_u64 = self
+            .pk
+            .iter()
+            .map(|x| x.to_u64().unwrap())
+            .collect::<Vec<u64>>();
+
+        let mut hasher = Hasher::new();
+        hasher.update(params.degree().to_le_bytes().as_slice());
+        hasher.update(self.pk.len().to_le_bytes().as_slice());
+        hasher.update(
+            &pk_bound_u64
+                .iter()
+                .flat_map(|num| num.to_le_bytes())
+                .collect::<Vec<u8>>(),
+        );
+        hasher.update(
+            &ctx.moduli()
+                .iter()
+                .flat_map(|num| num.to_le_bytes())
+                .collect::<Vec<u8>>(),
+        );
+        let _domain_seperator = BigUint::from_bytes_le(hasher.finalize().as_bytes());
+        let size = (10 * self.pk.len() + 4) * params.degree() - 8;
+        let io_pattern = [
+            BigUint::from_usize(size).unwrap(),
+            BigUint::from_usize(2 * self.pk.len()).unwrap(),
+        ]
+        .map(|x| x.to_bytes_le());
+        hasher.update(io_pattern[0].as_slice());
+        hasher.update(io_pattern[1].as_slice());
+
+        // TODO: Match the TAG calculation with the one in SAFE
+        let tag = BigUint::from_bytes_le(hasher.finalize().as_bytes()) % ctx.modulus().clone();
+
+        writeln!(file, "/// Constant value for the SAFE sponge algorithm.")?;
+        writeln!(file, "pub global TAG: Field = {:?};", tag)?;
+
         Ok(())
     }
 }
@@ -943,24 +981,6 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
     let file_path = output_path.join("Prover.toml");
     std::fs::write(file_path, toml_string).expect("Failed to write TOML file");
 
-    let zeroes = InputValidationVectors::new(moduli.len(), params.degree());
-    let zeroes_data = to_prover_toml_format(&zeroes);
-    let zeroes_toml = toml::to_string_pretty(&zeroes_data).expect("Failed to serialize TOML");
-
-    let file_path_zeroes = output_path.join("Prover0.toml");
-    std::fs::write(file_path_zeroes, zeroes_toml).expect("Failed to write zeroes TOML");
-
-    // Generate zeros filename and write file
-    let filename_zeroes = format!(
-        "pk_enc_{}_{}x{}_{}_zeroes.json",
-        N,
-        moduli.len(),
-        moduli_bitsize,
-        t.modulus()
-    );
-    let zeroes_json = InputValidationVectors::new(moduli.len(), params.degree()).to_json();
-    write_json_to_file(&output_path, &filename_zeroes, &zeroes_json);
-
     let filename_constants = format!(
         "pk_enc_constants_{}_{}x{}_{}.nr",
         N,