From 7d41b0be9eb49dac5473ee36bc8c9a97c24d79c7 Mon Sep 17 00:00:00 2001
From: Dmitry Vdovin <voidxnull@gmail.com>
Date: Thu, 24 Jun 2021 18:13:00 +0200
Subject: [PATCH] Attempt to fix r1cs and wtns file generation

---
 Cargo.lock                                  |  4 +--
 fawkes-crypto/Cargo.toml                    |  2 +-
 fawkes-crypto/src/backend/r1cs/r1cs_file.rs | 14 +++++++++--
 fawkes-crypto/src/backend/r1cs/wtns_file.rs | 27 ++++++++++++---------
 4 files changed, 30 insertions(+), 17 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 59a386d..7c14b69 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -807,9 +807,9 @@ checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f"
 
 [[package]]
 name = "wtns-file"
-version = "0.1.2"
+version = "0.1.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "97af06ea70a8f40f1c0e3db66fbae790455197bb3b9df3396329cdb0ccd22793"
+checksum = "3d3b856452298f68a5879e3901918bac5d753ca9fa4be8a983a37a3d25dabf0a"
 dependencies = [
  "byteorder",
 ]
diff --git a/fawkes-crypto/Cargo.toml b/fawkes-crypto/Cargo.toml
index 3ed6c32..f5cf141 100644
--- 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
index 7833969..639222e 100644
--- 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
index 2618c2a..650b03f 100644
--- 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))
 }