Don't unwrap when not synthesizing witness.

.gitignore

@@ -3,3 +3,4 @@ Cargo.lock
 /scratch
 
 .vscode/settings.json
+.idea

Cargo.toml

@@ -16,7 +16,7 @@ blstrs = { workspace = true, optional = true }
 byteorder = { workspace = true }
 ec-gpu = { workspace = true, optional = true }
 ec-gpu-gen = { workspace = true, optional = true }
-ff ={ workspace = true }
+ff = { workspace = true }
 generic-array = { workspace = true }
 pasta_curves = { workspace = true, features = ["serde"] }
 serde = { version = "1.0", features = ["derive"] }
@@ -80,13 +80,13 @@ abomonation = ["dep:abomonation", "dep:abomonation_derive"]
 [workspace]
 resolver = "2"
 members = [
-  "gbench",
+    "gbench",
 ]
 
 # Dependencies that should be kept in sync through the whole workspace
 [workspace.dependencies]
 bellpepper-core = { version = "0.4.0", default-features = false }
-bellpepper = { version = "0.4.0", default-features = false }
+bellpepper = { git = "https://github.com/lurk-lab/bellpepper.git", branch = "dev" }
 blake2s_simd = "1.0.1"
 blstrs = { version = "0.7.0" }
 ff = "0.13.0"

src/sponge/api.rs

@@ -1,9 +1,10 @@
+use ff::PrimeField;
+
 /// This module implements a variant of the 'Secure Sponge API for Field Elements':  https://hackmd.io/bHgsH6mMStCVibM_wYvb2w
 ///
 /// The API is defined by the `SpongeAPI` trait, which is implemented in terms of the `InnerSpongeAPI` trait.
 /// `Neptune` provides implementations of `InnerSpongeAPI` for both `sponge::Sponge` and `sponge_circuit::SpongeCircuit`.
-use crate::poseidon::{Arity, Poseidon, PoseidonConstants};
-use ff::PrimeField;
+use crate::poseidon::{Arity, PoseidonConstants};
 
 #[derive(Debug)]
 pub enum Error {
@@ -259,6 +260,20 @@ impl<F: PrimeField, A: Arity<F>, S: InnerSpongeAPI<F, A>> SpongeAPI<F, A> for S
 
 #[cfg(test)]
 mod test {
+    use bellpepper::util_cs::test_shape_cs::TestShapeCS;
+    use bellpepper_core::num::AllocatedNum;
+    use bellpepper_core::test_cs::TestConstraintSystem;
+    use bellpepper_core::ConstraintSystem;
+    use blstrs::Scalar as Fr;
+    use ff::{Field, PrimeFieldBits};
+    use generic_array::typenum::U24;
+    use serde::{Deserialize, Serialize};
+
+    use crate::circuit2::Elt;
+    use crate::sponge::circuit::SpongeCircuit;
+    use crate::sponge::vanilla::Mode::Simplex;
+    use crate::sponge::vanilla::SpongeTrait;
+
     use super::*;
 
     #[test]
@@ -308,4 +323,70 @@ mod test {
             0,
         );
     }
+
+    #[test]
+    fn test_sponge_api_multiple_cs() {
+        fn sponge_cycle<Scalar, CS: ConstraintSystem<Scalar>>(
+            cs: &mut CS,
+            elts: &[AllocatedNum<Scalar>],
+        ) -> Vec<Elt<Scalar>>
+        where
+            Scalar: PrimeField + PrimeFieldBits + Serialize + for<'de> Deserialize<'de>,
+        {
+            let constant: PoseidonConstants<Scalar, U24> = PoseidonConstants::new();
+            let mut ns = cs.namespace(|| "ns");
+            let mut sponge = SpongeCircuit::new_with_constants(&constant, Simplex);
+            let acc = &mut ns;
+            let parameter = IOPattern(vec![
+                SpongeOp::Absorb(elts.len() as u32),
+                SpongeOp::Squeeze(1u32),
+            ]);
+
+            sponge.start(parameter, None, acc);
+            SpongeAPI::absorb(
+                &mut sponge,
+                elts.len() as u32,
+                &(0..elts.len())
+                    .map(|i| Elt::Allocated(elts[i].clone()))
+                    .collect::<Vec<Elt<Scalar>>>(),
+                acc,
+            );
+            let output = SpongeAPI::squeeze(&mut sponge, 1, acc);
+            sponge.finish(acc).unwrap();
+            output
+        }
+
+        /*********************************
+         * Test absorb w/ ShapeCS
+         *********************************/
+        let mut cs: TestShapeCS<Fr> = TestShapeCS::new();
+
+        let elts = (0..10)
+            .map(|i| {
+                AllocatedNum::alloc(cs.namespace(|| format!("elt_{i}")), || Ok(Fr::ONE)).unwrap()
+            })
+            .collect::<Vec<_>>();
+
+        // Values are `None`, we should get a None back
+        let hash = sponge_cycle(&mut cs, &elts);
+        assert!(hash[0].val().is_none());
+
+        /*********************************
+         * Test absorb w/ TestConstraintSystem
+         *********************************/
+        let mut cs: TestConstraintSystem<Fr> = TestConstraintSystem::new();
+
+        let elts = (0..10)
+            .map(|i| {
+                AllocatedNum::alloc(cs.namespace(|| format!("elt_{i}")), || Ok(Fr::ONE)).unwrap()
+            })
+            .collect::<Vec<_>>();
+
+        // We have values, we expect one in return
+        let hash = sponge_cycle(&mut cs, &elts);
+        assert_eq!(
+            "Scalar(0x4d1f7863ee494536a938bd87d761a30828eeeeebfbc160117135dc6766f6e16c)",
+            hash[0].val().unwrap().to_string()
+        );
+    }
 }

src/sponge/circuit.rs

@@ -34,6 +34,7 @@ where
     pattern: IOPattern,
     io_count: usize,
     poseidon: Poseidon<'a, F, A>,
+    shape_only: bool, // must be set to true if any value required for accurate witness-generation is missing.
     _c: PhantomData<C>,
 }
 
@@ -57,6 +58,7 @@ impl<'a, F: PrimeField, A: Arity<F>, CS: 'a + ConstraintSystem<F>> SpongeTrait<'
             queue: VecDeque::with_capacity(A::to_usize()),
             pattern: IOPattern(Vec::new()),
             poseidon: Poseidon::new(constants),
+            shape_only: false,
             io_count: 0,
             _c: Default::default(),
         }
@@ -101,7 +103,14 @@ impl<'a, F: PrimeField, A: Arity<F>, CS: 'a + ConstraintSystem<F>> SpongeTrait<'
     }
 
     fn set_element(&mut self, index: usize, elt: Self::Elt) {
-        self.poseidon.elements[index] = elt.val().unwrap();
+        // If `elt` has no value, we are synthesizing. `self.poseidon.elements` is used only for witness-generation, so we
+        // don't need to set in that case.
+        if let Some(f) = elt.val() {
+            self.poseidon.elements[index] = f;
+        } else {
+            // Since we failed to update the poseidon state, record that we must only be synthesizing shape.
+            self.shape_only = true;
+        }
         self.state.elements[index] = elt;
     }
 
@@ -134,6 +143,9 @@ impl<'a, F: PrimeField, A: Arity<F>, CS: 'a + ConstraintSystem<F>> SpongeTrait<'
         self.permutation_count += 1;
 
         if ns.is_witness_generator() {
+            // Do not risk generating an inaccurate witness if a missing value has tainted our tracking of the concrete
+            // state.
+            assert!(!self.shape_only);
             self.poseidon.generate_witness_into_cs(ns);
 
             for (elt, scalar) in self