Attempt to fix r1cs and wtns file generation

zeropoolnetwork · Jun 24, 2021 · 7d41b0b · 7d41b0b
1 parent dd10c08
commit 7d41b0b
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 17 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/fawkes-crypto/Cargo.toml b/fawkes-crypto/Cargo.toml
@@ -26,7 +26,7 @@ getrandom = { version = "0.2", optional = true }
 bit-vec = "0.6.3"
 itertools = "0.10.0"
 r1cs-file = { version = "0.2.1", optional = true }
-wtns-file = { version = "0.1.2", optional = true }
+wtns-file = { version = "0.1.5", optional = true }
 
 [dependencies.blake2_rfc]
 version = "0.0.1"

diff --git a/fawkes-crypto/src/backend/r1cs/r1cs_file.rs b/fawkes-crypto/src/backend/r1cs/r1cs_file.rs
@@ -5,6 +5,7 @@ use crate::circuit::cs::Gate;
 use crate::circuit::lc::Index;
 use crate::ff_uint::{Num, PrimeField, Uint};
 use crate::backend::r1cs::ConstTrackerFile;
+use std::collections::HashSet;
 
 pub fn get_r1cs_file<Fr: PrimeField, const FS: usize>(
     gates: &Vec<Gate<Fr>>,
@@ -15,17 +16,26 @@ pub fn get_r1cs_file<Fr: PrimeField, const FS: usize>(
     let mut n_pub_in = 0;
     let mut n_prvt_in = 0;
 
+    let mut pub_inputs = HashSet::new();
+    let mut prvt_inputs = HashSet::new();
+
     let constraints = gates
         .iter()
         .map(|gate| {
             let mut map_comb = |(c, i): &(Num<Fr>, Index)| {
                 let i = match *i {
                     Index::Input(i) => {
-                        n_pub_in += 1;
+                        if pub_inputs.insert(i) {
+                            n_pub_in += 1;
+                        }
+
                         i
                     }
                     Index::Aux(i) => {
-                        n_prvt_in += 1;
+                        if prvt_inputs.insert(i) {
+                            n_prvt_in += 1;
+                        }
+
                         i
                     }
                 };

diff --git a/fawkes-crypto/src/backend/r1cs/wtns_file.rs b/fawkes-crypto/src/backend/r1cs/wtns_file.rs
@@ -3,20 +3,21 @@ use wtns_file::{WtnsFile, FieldElement};
 
 use crate::ff_uint::{Uint, PrimeField};
 use crate::core::signal::Signal;
-use crate::circuit::cs::{WitnessCS, Gate, CS};
-use crate::circuit::lc::Index;
+use crate::circuit::cs::{WitnessCS, Gate};
 
 pub fn get_witness<
     'a,
     Fr: PrimeField,
     Pub: Signal<WitnessCS<'a, Fr>>,
     Sec: Signal<WitnessCS<'a, Fr>>,
+    C: Fn(Pub, Sec),
     const FS: usize,
 >(
     gates: &'a Vec<Gate<Fr>>,
     consts: &'a BitVec,
     input_pub: &'a Pub::Value,
     input_sec: &'a Sec::Value,
+    circuit: C,
 ) -> WtnsFile<FS> {
     let cs = WitnessCS::rc_new(gates, consts);
 
@@ -25,19 +26,21 @@ pub fn get_witness<
 
     let signal_pub = Pub::alloc(&cs, Some(input_pub));
     signal_pub.inputize();
-    let _signal_sec = Sec::alloc(&cs, Some(input_sec));
+    let signal_sec = Sec::alloc(&cs, Some(input_sec));
 
-    let cs = cs.borrow_mut();
-    let mut witness = Vec::with_capacity(cs.num_aux());
+    circuit(signal_pub, signal_sec);
 
-    for i in (cs.num_input() + 1)..cs.num_aux() {
-        let num = cs.get_value(Index::Aux(i)).unwrap();
-        let mut bytes = [0; FS];
-        num.0.to_uint().put_little_endian(&mut bytes);
-        let fe = FieldElement::from(bytes);
+    let cs = cs.borrow();
 
-        witness.push(fe);
-    }
+    let witness = cs.values_input
+        .iter()
+        .chain(cs.values_aux.iter())
+        .map(|num| {
+            let mut bytes = [0; FS];
+            num.0.to_uint().put_little_endian(&mut bytes);
+            FieldElement::from(bytes)
+        })
+        .collect();
 
     WtnsFile::from_vec(witness, FieldElement::from(prime_bytes))
 }