Add sponge ops

circuit-prover/Cargo.toml

@@ -49,3 +49,7 @@ path = "examples/fibonacci.rs"
 [[example]]
 name = "mmcs_verify"
 path = "examples/mmcs_verify.rs"
+
+[[example]]
+name = "hash"
+path = "examples/hash.rs"

circuit-prover/examples/hash.rs

@@ -0,0 +1,62 @@
+use core::array;
+
+/// Hash verification circuit: prove the correct computation of a hash.
+/// Public inputs: inputs, output
+use p3_baby_bear::{BabyBear, Poseidon2BabyBear};
+use p3_circuit::Challenger;
+use p3_circuit::builder::CircuitBuilder;
+use p3_circuit_prover::MultiTableProver;
+use p3_circuit_prover::config::babybear_config::build_standard_config_babybear;
+use p3_circuit_prover::prover::{ProverError, TablePacking};
+use p3_symmetric::{CryptographicHasher, PaddingFreeSponge};
+use rand::rngs::SmallRng;
+use rand::{Rng, SeedableRng};
+
+type F = BabyBear;
+
+const HASH_RATE: usize = 8;
+const HASH_CAPACITY: usize = 8;
+const HASH_STATE_SIZE: usize = HASH_RATE + HASH_CAPACITY;
+
+fn main() -> Result<(), ProverError> {
+    let mut rng = SmallRng::seed_from_u64(1);
+    let perm = Poseidon2BabyBear::<HASH_STATE_SIZE>::new_from_rng_128(&mut rng);
+
+    let mut builder = CircuitBuilder::<F>::new();
+
+    let mut challenger = Challenger::default();
+
+    // Public inputs: hash inputs and hash output
+    let inputs: [_; HASH_STATE_SIZE] = array::from_fn(|_| builder.add_public_input());
+    let outputs: [_; HASH_RATE] = array::from_fn(|_| builder.add_public_input());
+
+    // Add hash operations
+    challenger.add_inputs(&inputs);
+    challenger.squeeze(&mut builder, &outputs);
+
+    let circuit = builder.build()?;
+    let mut runner = circuit.runner();
+
+    // Generate random inputs
+    let input_values: [_; 16] = array::from_fn(|_| rng.random::<F>());
+
+    // Set public inputs
+    let hasher = PaddingFreeSponge::<_, HASH_STATE_SIZE, HASH_RATE, HASH_RATE>::new(perm.clone());
+    let output_values = hasher.hash_iter(input_values);
+    let public_inputs = input_values
+        .into_iter()
+        .chain(output_values)
+        .collect::<Vec<_>>();
+    runner.set_public_inputs(&public_inputs)?;
+
+    let traces = runner.run()?;
+    let config = build_standard_config_babybear();
+    let table_packing = TablePacking::from_counts(4, 1);
+    let multi_prover = MultiTableProver::new(config).with_table_packing(table_packing);
+    let proof = multi_prover.prove_all_tables(&traces)?;
+    multi_prover.verify_all_tables(&proof)?;
+
+    println!("✅ Verified hash({:?}) = {:?}", input_values, output_values);
+
+    Ok(())
+}

circuit-prover/src/prover.rs

@@ -267,6 +267,9 @@ where
         let mmcs_air = MmcsVerifyAir::<F>::new(self.mmcs_config);
         let mmcs_proof = prove(&self.config, &mmcs_air, mmcs_matrix, pis);
 
+        // TODO: Add sponge AIR
+        // let sponge_matrix = SpongeAir::<F, Perm>::trace_to_matrix(&traces.sponge_trace);
+
         Ok(MultiTableProof {
             witness: TableProof {
                 proof: witness_proof,
@@ -398,7 +401,6 @@ mod tests {
         let x_val = BabyBear::from_u64(7);
         let expected_val = BabyBear::from_u64(13); // 7 + 10 - 3 - 1 = 13
         runner.set_public_inputs(&[x_val, expected_val])?;
-
         let traces = runner.run()?;
 
         // Create BabyBear prover and prove all tables
@@ -660,6 +662,7 @@ mod tests {
 
         let expected_val = x_val * y_val + z_val;
         runner.set_public_inputs(&[x_val, y_val, z_val, expected_val])?;
+
         let traces = runner.run()?;
 
         // Build Goldilocks config with challenge degree 2 (Poseidon2)

circuit/Cargo.toml

@@ -21,10 +21,10 @@ p3-koala-bear.workspace = true
 p3-matrix.workspace = true
 p3-symmetric.workspace = true
 p3-uni-stark.workspace = true
-rand.workspace = true
 
 # Other common dependencies
 hashbrown = "0.16.0"
+rand.workspace = true
 itertools.workspace = true
 serde.workspace = true
 thiserror.workspace = true

circuit/src/builder.rs

@@ -277,6 +277,26 @@ where
         op_id
     }
 
+    /// Add a hash absorb operation (non-primitive operation)
+    pub fn add_hash_absorb(&mut self, input_exprs: &[ExprId], reset: bool) -> NonPrimitiveOpId {
+        let op_id = NonPrimitiveOpId(self.non_primitive_ops.len() as u32);
+        let witness_exprs = input_exprs.to_vec();
+        self.non_primitive_ops
+            .push((op_id, NonPrimitiveOpType::HashAbsorb(reset), witness_exprs));
+
+        op_id
+    }
+
+    /// Add a hash squeeze operation (non-primitive operation)
+    pub fn add_hash_squeeze(&mut self, output_exprs: &[ExprId]) -> NonPrimitiveOpId {
+        let op_id = NonPrimitiveOpId(self.non_primitive_ops.len() as u32);
+        let witness_exprs = output_exprs.to_vec();
+        self.non_primitive_ops
+            .push((op_id, NonPrimitiveOpType::HashSqueeze, witness_exprs));
+
+        op_id
+    }
+
     /// Check whether an op type is enabled on this builder.
     fn is_op_enabled(&self, op: &NonPrimitiveOpType) -> bool {
         self.enabled_ops.contains_key(op)
@@ -574,6 +594,27 @@ where
                         root: root_widx,
                     });
                 }
+                NonPrimitiveOpType::HashAbsorb(reset) => {
+                    let inputs_widx = witness_exprs
+                        .iter()
+                        .map(|expr| Self::get_witness_id(expr_to_widx, *expr, "HashAbsorb input"))
+                        .collect::<Vec<_>>();
+
+                    lowered_ops.push(NonPrimitiveOp::HashAbsorb {
+                        reset_flag: *reset,
+                        inputs: inputs_widx.into_iter().collect::<Result<Vec<_>, _>>()?,
+                    });
+                }
+                NonPrimitiveOpType::HashSqueeze => {
+                    let outputs_widx = witness_exprs
+                        .iter()
+                        .map(|expr| Self::get_witness_id(expr_to_widx, *expr, "HashSqueeze output"))
+                        .collect::<Vec<_>>();
+
+                    lowered_ops.push(NonPrimitiveOp::HashSqueeze {
+                        outputs: outputs_widx.into_iter().collect::<Result<Vec<_>, _>>()?,
+                    });
+                }
                 NonPrimitiveOpType::FriVerify => {
                     todo!() // TODO: Add FRIVerify when it lands
                 } // Add more variants here as needed
@@ -1000,6 +1041,7 @@ mod tests {
         assert_eq!(circuit.non_primitive_ops.len(), 1);
         match &circuit.non_primitive_ops[0] {
             NonPrimitiveOp::MmcsVerify { .. } => {}
+            _ => panic!("Expected MmcsVerify operation"),
         }
     }
 

circuit/src/challenger.rs

@@ -0,0 +1,64 @@
+use alloc::vec::Vec;
+
+use p3_field::PrimeCharacteristicRing;
+
+use crate::{CircuitBuilder, ExprId};
+
+pub struct Challenger<const R: usize> {
+    pub inputs: Vec<ExprId>,
+    pub reset_flag: bool,
+}
+
+impl<const R: usize> Default for Challenger<R> {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl<const R: usize> Challenger<R> {
+    pub fn new() -> Self {
+        Self {
+            inputs: Vec::new(),
+            reset_flag: false,
+        }
+    }
+
+    pub fn add_input(&mut self, expr_id: ExprId) {
+        self.inputs.push(expr_id);
+    }
+
+    pub fn add_inputs(&mut self, expr_ids: &[ExprId]) {
+        self.inputs.extend_from_slice(expr_ids);
+    }
+
+    pub fn clear(&mut self) {
+        self.inputs.clear();
+        self.reset_flag = true;
+    }
+
+    pub fn squeeze<F: Clone + PrimeCharacteristicRing + Eq + core::hash::Hash>(
+        &mut self,
+        builder: &mut CircuitBuilder<F>,
+        outputs: &[ExprId; R],
+    ) {
+        assert!(
+            self.inputs.len().is_multiple_of(R),
+            "Number of inputs must be a multiple of R"
+        );
+
+        let input_chunks = self.inputs.chunks_exact(R);
+
+        for chunk in input_chunks {
+            if self.reset_flag {
+                builder.add_hash_absorb(chunk, true);
+                self.reset_flag = false;
+            } else {
+                builder.add_hash_absorb(chunk, false);
+            }
+        }
+        self.inputs.clear();
+        self.reset_flag = false;
+
+        builder.add_hash_squeeze(outputs);
+    }
+}

circuit/src/lib.rs

@@ -2,6 +2,7 @@
 extern crate alloc;
 
 pub mod builder;
+pub mod challenger;
 pub mod circuit;
 pub mod errors;
 pub mod expr;
@@ -15,6 +16,7 @@ pub mod utils;
 
 // Re-export public API
 pub use builder::{CircuitBuilder, CircuitBuilderError};
+pub use challenger::Challenger;
 pub use circuit::{Circuit, CircuitField};
 pub use errors::CircuitError;
 pub use expr::{Expr, ExpressionGraph};

circuit/src/op.rs

@@ -55,7 +55,8 @@ pub enum NonPrimitiveOpType {
     // Mmcs Verify gate with the argument is the size of the path
     MmcsVerify,
     FriVerify,
-    // Future: FriVerify, HashAbsorb, etc.
+    HashAbsorb(bool), // bool: If true, reset the capacity before absorbing
+    HashSqueeze,
 }
 
 /// Non-primitive operation types
@@ -98,6 +99,17 @@ pub enum NonPrimitiveOp {
         index: WitnessId,
         root: MmcsWitnessId,
     },
+    /// Sponge hash absorb operation. Absorbs a single chunk of input.
+    ///
+    /// Public interface (on witness bus):
+    /// - `reset_flag`: Whether to reset the capacity before absorbing
+    /// - `inputs`: The input chunk to be absorbed.
+    HashAbsorb {
+        reset_flag: bool,
+        inputs: Vec<WitnessId>,
+    },
+    /// Sponge hash squeeze operation. Produces a single chunk of output.
+    HashSqueeze { outputs: Vec<WitnessId> },
 }
 
 /// Configuration parameters for Mmcs verification operations. When