feat: add support for merging input arrays

co-circom/co-circom-snarks/src/lib.rs

@@ -6,7 +6,7 @@ use ark_ff::PrimeField;
 use ark_serialize::{CanonicalDeserialize, CanonicalSerialize};
 use circom_types::Witness;
 use mpc_core::protocols::{
-    rep3::{self, Rep3PrimeFieldShare, Rep3ShareVecType},
+    rep3::{self, MaybeRep3ShareVecType, Rep3PrimeFieldShare, Rep3ShareVecType},
     shamir::{self, ShamirPrimeFieldShare},
 };
 use rand::{distributions::Standard, prelude::Distribution, CryptoRng, Rng, SeedableRng};
@@ -83,6 +83,9 @@ where
     /// A map from variable names to the share of the field element.
     /// This is a BTreeMap because it implements Canonical(De)Serialize.
     pub shared_inputs: BTreeMap<String, Rep3ShareVecType<F, U>>,
+    /// A map from variable names to vecs with maybe unknown elements that need to be merged.
+    /// This is a BTreeMap because it implements Canonical(De)Serialize.
+    pub maybe_shared_inputs: BTreeMap<String, MaybeRep3ShareVecType<F>>,
 }
 
 impl<F: PrimeField, U: Rng + SeedableRng + CryptoRng> Default for SerializeableSharedRep3Input<F, U>
@@ -93,6 +96,7 @@ where
         Self {
             public_inputs: BTreeMap::new(),
             shared_inputs: BTreeMap::new(),
+            maybe_shared_inputs: BTreeMap::new(),
         }
     }
 }
@@ -109,45 +113,76 @@ where
         seeded: bool,
         additive: bool,
     ) -> [Rep3ShareVecType<F, U>; 3] {
-        let (share1, share2, share3) = match (seeded, additive) {
+        match (seeded, additive) {
             (true, true) => {
                 let [share1, share2, share3] =
                     rep3::share_field_elements_additive_seeded::<_, _, U>(input, rng);
                 let share1 = Rep3ShareVecType::SeededAdditive(share1);
                 let share2 = Rep3ShareVecType::SeededAdditive(share2);
                 let share3 = Rep3ShareVecType::SeededAdditive(share3);
-                (share1, share2, share3)
+                [share1, share2, share3]
             }
             (true, false) => {
                 let [share1, share2, share3] =
                     rep3::share_field_elements_seeded::<_, _, U>(input, rng);
                 let share1 = Rep3ShareVecType::SeededReplicated(share1);
                 let share2 = Rep3ShareVecType::SeededReplicated(share2);
                 let share3 = Rep3ShareVecType::SeededReplicated(share3);
-                (share1, share2, share3)
+                [share1, share2, share3]
             }
             (false, true) => {
                 let [share1, share2, share3] = rep3::share_field_elements_additive(input, rng);
                 let share1 = Rep3ShareVecType::Additive(share1);
                 let share2 = Rep3ShareVecType::Additive(share2);
                 let share3 = Rep3ShareVecType::Additive(share3);
-                (share1, share2, share3)
+                [share1, share2, share3]
             }
             (false, false) => {
                 let [share1, share2, share3] = rep3::share_field_elements(input, rng);
                 let share1 = Rep3ShareVecType::Replicated(share1);
                 let share2 = Rep3ShareVecType::Replicated(share2);
                 let share3 = Rep3ShareVecType::Replicated(share3);
-                (share1, share2, share3)
+                [share1, share2, share3]
             }
-        };
-        [share1, share2, share3]
+        }
+    }
+
+    /// Shares a given input with unknown elements into a [MaybeRep3ShareVecType] type.
+    pub fn maybe_share_rep3<R: Rng + CryptoRng>(
+        input: &[Option<F>],
+        rng: &mut R,
+        additive: bool,
+    ) -> [MaybeRep3ShareVecType<F>; 3] {
+        if additive {
+            let [share1, share2, share3] = rep3::share_maybe_field_elements_additive(input, rng);
+            let share1 = MaybeRep3ShareVecType::Additive(share1);
+            let share2 = MaybeRep3ShareVecType::Additive(share2);
+            let share3 = MaybeRep3ShareVecType::Additive(share3);
+            [share1, share2, share3]
+        } else {
+            let [share1, share2, share3] = rep3::share_maybe_field_elements(input, rng);
+            let share1 = MaybeRep3ShareVecType::Replicated(share1);
+            let share2 = MaybeRep3ShareVecType::Replicated(share2);
+            let share3 = MaybeRep3ShareVecType::Replicated(share3);
+            [share1, share2, share3]
+        }
     }
 
     /// Merges two [SerializeableSharedRep3Input]s into one, performing basic sanity checks.
     pub fn merge(self, other: Self) -> eyre::Result<Self> {
         let mut shared_inputs = self.shared_inputs;
+        let maybe_shared_inputs = self.maybe_shared_inputs;
         let public_inputs = self.public_inputs;
+
+        for (key, value) in other.public_inputs.iter() {
+            if !public_inputs.contains_key(key) {
+                eyre::bail!("Public input \"{key}\" must be present in all files");
+            }
+            if public_inputs.get(key).expect("is there we checked") != value {
+                eyre::bail!("Public input \"{key}\" must be same in all files");
+            }
+        }
+
         for (key, value) in other.shared_inputs {
             if shared_inputs.contains_key(&key) {
                 eyre::bail!("Input with name {} present in multiple input shares", key);
@@ -159,18 +194,70 @@ where
             }
             shared_inputs.insert(key, value);
         }
-        for (key, value) in other.public_inputs {
-            if !public_inputs.contains_key(&key) {
-                eyre::bail!("Public input \"{key}\" must be present in all files");
+
+        let mut merged_maybe_shared_inputs = BTreeMap::new();
+        if maybe_shared_inputs.len() != other.maybe_shared_inputs.len() {
+            eyre::bail!("Both inputs must have the same number of unmerged entries");
+        }
+        for ((k1, v1), (k2, v2)) in maybe_shared_inputs
+            .into_iter()
+            .zip(other.maybe_shared_inputs.into_iter())
+        {
+            if k1 != k2 {
+                eyre::bail!("Both inputs must have the same keys for unmerged elements");
             }
-            if public_inputs.get(&key).expect("is there we checked") != &value {
-                eyre::bail!("Public input \"{key}\" must be same in all files");
+
+            match (v1, v2) {
+                (
+                    MaybeRep3ShareVecType::Replicated(shares),
+                    MaybeRep3ShareVecType::Replicated(other_shares),
+                ) => {
+                    let merged = shares
+                        .into_iter()
+                        .zip(other_shares.into_iter())
+                        .map(|(a, b)| match (a, b) {
+                            (None, None) => Ok(None),
+                            (a @ Some(_), None) | (None, a @ Some(_)) => Ok(a),
+                            _ => Err(eyre::eyre!("Input {} present in both unmerged inputs", k1)),
+                        })
+                        .collect::<Result<Vec<_>, eyre::Report>>()?;
+                    if let Some(merged) = merged.iter().cloned().collect::<Option<Vec<_>>>() {
+                        shared_inputs.insert(k1.clone(), Rep3ShareVecType::Replicated(merged));
+                    } else {
+                        merged_maybe_shared_inputs
+                            .insert(k1.clone(), MaybeRep3ShareVecType::Replicated(merged));
+                    }
+                }
+                (
+                    MaybeRep3ShareVecType::Additive(shares),
+                    MaybeRep3ShareVecType::Additive(other_shares),
+                ) => {
+                    let merged = shares
+                        .into_iter()
+                        .zip(other_shares.into_iter())
+                        .map(|(a, b)| match (a, b) {
+                            (None, None) => Ok(None),
+                            (a @ Some(_), None) | (None, a @ Some(_)) => Ok(a),
+                            _ => Err(eyre::eyre!("Input {} present in both unmerged inputs", k1)),
+                        })
+                        .collect::<Result<Vec<_>, eyre::Report>>()?;
+                    if let Some(merged) = merged.iter().cloned().collect::<Option<Vec<_>>>() {
+                        shared_inputs.insert(k1.clone(), Rep3ShareVecType::Additive(merged));
+                    } else {
+                        merged_maybe_shared_inputs
+                            .insert(k1.clone(), MaybeRep3ShareVecType::Additive(merged));
+                    }
+                }
+                _ => {
+                    eyre::bail!("Input {} cannot be merged, the share type is different", k1);
+                }
             }
         }
 
         Ok(Self {
-            shared_inputs,
             public_inputs,
+            shared_inputs,
+            maybe_shared_inputs: merged_maybe_shared_inputs,
         })
     }
 }

co-circom/co-circom/examples/groth16/run_full_sum_arrays_with_merge.sh

@@ -0,0 +1,23 @@
+# split input into shares
+cargo run --release --bin co-circom -- split-input  --circuit test_vectors/sum_arrays/circuit.circom --input test_vectors/sum_arrays/input0.json --protocol REP3 --curve BN254 --out-dir test_vectors/sum_arrays
+cargo run --release --bin co-circom -- split-input  --circuit test_vectors/sum_arrays/circuit.circom --input test_vectors/sum_arrays/input1.json --protocol REP3 --curve BN254 --out-dir test_vectors/sum_arrays
+cargo run --release --bin co-circom -- split-input  --circuit test_vectors/sum_arrays/circuit.circom --input test_vectors/sum_arrays/input2.json --protocol REP3 --curve BN254 --out-dir test_vectors/sum_arrays
+# merge inputs into single input file
+cargo run --release --bin co-circom -- merge-input-shares --inputs test_vectors/sum_arrays/input0.json.0.shared --inputs test_vectors/sum_arrays/input1.json.0.shared --protocol REP3 --curve BN254 --out test_vectors/sum_arrays/input01.json.0.shared
+cargo run --release --bin co-circom -- merge-input-shares --inputs test_vectors/sum_arrays/input0.json.1.shared --inputs test_vectors/sum_arrays/input1.json.1.shared --protocol REP3 --curve BN254 --out test_vectors/sum_arrays/input01.json.1.shared
+cargo run --release --bin co-circom -- merge-input-shares --inputs test_vectors/sum_arrays/input0.json.2.shared --inputs test_vectors/sum_arrays/input1.json.2.shared --protocol REP3 --curve BN254 --out test_vectors/sum_arrays/input01.json.2.shared
+cargo run --release --bin co-circom -- merge-input-shares --inputs test_vectors/sum_arrays/input01.json.0.shared --inputs test_vectors/sum_arrays/input2.json.0.shared --protocol REP3 --curve BN254 --out test_vectors/sum_arrays/input.json.0.shared
+cargo run --release --bin co-circom -- merge-input-shares --inputs test_vectors/sum_arrays/input01.json.1.shared --inputs test_vectors/sum_arrays/input2.json.1.shared --protocol REP3 --curve BN254 --out test_vectors/sum_arrays/input.json.1.shared
+cargo run --release --bin co-circom -- merge-input-shares --inputs test_vectors/sum_arrays/input01.json.2.shared --inputs test_vectors/sum_arrays/input2.json.2.shared --protocol REP3 --curve BN254 --out test_vectors/sum_arrays/input.json.2.shared
+# run witness extension in MPC
+cargo run --release --bin co-circom -- generate-witness -O2 --input test_vectors/sum_arrays/input.json.0.shared --circuit test_vectors/sum_arrays/circuit.circom --protocol REP3 --curve BN254 --config ../configs/party1.toml --out test_vectors/sum_arrays/witness.wtns.0.shared &
+cargo run --release --bin co-circom -- generate-witness -O2 --input test_vectors/sum_arrays/input.json.1.shared --circuit test_vectors/sum_arrays/circuit.circom --protocol REP3 --curve BN254 --config ../configs/party2.toml --out test_vectors/sum_arrays/witness.wtns.1.shared &
+cargo run --release --bin co-circom -- generate-witness -O2 --input test_vectors/sum_arrays/input.json.2.shared --circuit test_vectors/sum_arrays/circuit.circom --protocol REP3 --curve BN254 --config ../configs/party3.toml --out test_vectors/sum_arrays/witness.wtns.2.shared
+wait $(jobs -p)
+# run proving in MPC
+cargo run --release --bin co-circom -- generate-proof groth16 --witness test_vectors/sum_arrays/witness.wtns.0.shared --zkey test_vectors/sum_arrays/sum_arrays.zkey --protocol REP3 --curve BN254 --config ../configs/party1.toml --out proof.0.json --public-input public_input.json &
+cargo run --release --bin co-circom -- generate-proof groth16 --witness test_vectors/sum_arrays/witness.wtns.1.shared --zkey test_vectors/sum_arrays/sum_arrays.zkey --protocol REP3 --curve BN254 --config ../configs/party2.toml --out proof.1.json &
+cargo run --release --bin co-circom -- generate-proof groth16 --witness test_vectors/sum_arrays/witness.wtns.2.shared --zkey test_vectors/sum_arrays/sum_arrays.zkey --protocol REP3 --curve BN254 --config ../configs/party3.toml --out proof.2.json
+wait $(jobs -p)
+# verify proof
+cargo run --release --bin co-circom -- verify groth16 --proof proof.0.json --vk test_vectors/sum_arrays/verification_key.json --public-input public_input.json --curve BN254

co-circom/co-circom/examples/groth16/test_vectors/sum_arrays/input0.json

@@ -0,0 +1,5 @@
+{
+  "a": ["1", "?", "?"],
+  "b": ["4", "5", "6"],
+  "c": ["7", "8", "9"]
+}

co-circom/co-circom/examples/groth16/test_vectors/sum_arrays/input1.json

@@ -0,0 +1,5 @@
+{
+  "a": ["?", "2", "?"],
+  "b": ["4", "5", "6"],
+  "c": ["7", "8", "9"]
+}

co-circom/co-circom/examples/groth16/test_vectors/sum_arrays/input2.json

@@ -0,0 +1,5 @@
+{
+  "a": ["?", "?", "3"],
+  "b": ["4", "5", "6"],
+  "c": ["7", "8", "9"]
+}

co-circom/co-circom/src/file_utils.rs

@@ -90,16 +90,20 @@ where
     }
 }
 
-pub(crate) fn parse_array<F: PrimeField>(val: &serde_json::Value) -> color_eyre::Result<Vec<F>> {
+pub(crate) fn parse_array<F: PrimeField>(
+    val: &serde_json::Value,
+) -> color_eyre::Result<Vec<Option<F>>> {
     let json_arr = val.as_array().expect("is an array");
     let mut field_elements = vec![];
     for ele in json_arr {
         if ele.is_array() {
             field_elements.extend(parse_array::<F>(ele)?);
         } else if ele.is_boolean() {
-            field_elements.push(parse_boolean(ele)?);
+            field_elements.push(Some(parse_boolean(ele)?));
+        } else if ele.as_str().is_some_and(|e| e == "?") {
+            field_elements.push(None);
         } else {
-            field_elements.push(parse_field(ele)?);
+            field_elements.push(Some(parse_field(ele)?));
         }
     }
     Ok(field_elements)

co-circom/co-circom/src/lib.rs

@@ -23,7 +23,7 @@ use co_circom_snarks::{
     SerializeableSharedRep3Input, SerializeableSharedRep3Witness, SharedInput, SharedWitness,
 };
 use co_groth16::Rep3CoGroth16;
-use color_eyre::eyre::Context;
+use color_eyre::eyre::{bail, Context, ContextCompat};
 use figment::{
     providers::{Env, Format, Serialized, Toml},
     Figment,
@@ -659,11 +659,15 @@ where
         let parsed_vals = if val.is_array() {
             file_utils::parse_array(&val)?
         } else if val.is_boolean() {
-            vec![file_utils::parse_boolean(&val)?]
+            vec![Some(file_utils::parse_boolean(&val)?)]
         } else {
-            vec![file_utils::parse_field(&val)?]
+            vec![Some(file_utils::parse_field(&val)?)]
         };
         if public_inputs.contains(&name) {
+            let parsed_vals = parsed_vals
+                .into_iter()
+                .collect::<Option<Vec<P::ScalarField>>>()
+                .context("Public inputs must not be unkown")?;
             shares[0]
                 .public_inputs
                 .insert(name.clone(), parsed_vals.clone());
@@ -672,11 +676,33 @@ where
                 .insert(name.clone(), parsed_vals.clone());
             shares[2].public_inputs.insert(name.clone(), parsed_vals);
         } else {
-            let [share0, share1, share2] =
-                SerializeableSharedRep3Input::share_rep3(&parsed_vals, &mut rng, seeded, additive);
-            shares[0].shared_inputs.insert(name.clone(), share0);
-            shares[1].shared_inputs.insert(name.clone(), share1);
-            shares[2].shared_inputs.insert(name.clone(), share2);
+            // if all elements are Some, then we can share normally
+            // else we can only share as Vec<Option<T>> and we have to merge unknown inputs later
+            if parsed_vals.iter().all(Option::is_some) {
+                let parsed_vals = parsed_vals
+                    .into_iter()
+                    .collect::<Option<Vec<_>>>()
+                    .expect("all are Some");
+                let [share0, share1, share2] = SerializeableSharedRep3Input::share_rep3(
+                    &parsed_vals,
+                    &mut rng,
+                    seeded,
+                    additive,
+                );
+                shares[0].shared_inputs.insert(name.clone(), share0);
+                shares[1].shared_inputs.insert(name.clone(), share1);
+                shares[2].shared_inputs.insert(name.clone(), share2);
+            } else {
+                let [share0, share1, share2] =
+                    SerializeableSharedRep3Input::<_, SeedRng>::maybe_share_rep3(
+                        &parsed_vals,
+                        &mut rng,
+                        additive,
+                    );
+                shares[0].maybe_shared_inputs.insert(name.clone(), share0);
+                shares[1].maybe_shared_inputs.insert(name.clone(), share1);
+                shares[2].maybe_shared_inputs.insert(name.clone(), share2);
+            };
         }
     }
     Ok(shares)
@@ -690,6 +716,10 @@ pub fn parse_shared_input<R: Read, F: PrimeField, N: Rep3Network>(
     let deserialized: SerializeableSharedRep3Input<F, SeedRng> =
         bincode::deserialize_from(reader).context("trying to parse input share file")?;
 
+    if !deserialized.maybe_shared_inputs.is_empty() {
+        bail!("still unmerged elements left");
+    }
+
     let public_inputs = deserialized.public_inputs;
     let shared_inputs_ = deserialized.shared_inputs;