Refactor r1cs backend, finish wtns file implementation

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.1", optional = true }
+wtns-file = { version = "0.1.2", optional = true }
 
 [dependencies.blake2_rfc]
 version = "0.0.1"

fawkes-crypto/src/backend/r1cs/const_tracker.rs

@@ -0,0 +1,78 @@
+use std::io::{self, Read, Write};
+
+use bit_vec::BitVec;
+use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
+
+const MAGIC: &[u8; 4] = b"ZPCT";
+
+pub struct ConstTrackerFile(pub BitVec);
+
+impl ConstTrackerFile {
+    pub fn read<R: Read>(mut r: R) -> io::Result<Self> {
+        let mut magic = [0; MAGIC.len()];
+        r.read_exact(&mut magic)?;
+
+        if magic != *MAGIC {
+            return Err(io::Error::new(
+                io::ErrorKind::InvalidData,
+                "Invalid magic sequence",
+            ));
+        }
+
+        let len = r.read_u64::<LittleEndian>()?;
+        let mut bytes = vec![0; div_ceil(len as usize, 8)];
+        r.read_exact(&mut bytes)?;
+
+        Ok(ConstTrackerFile(BitVec::from_bytes(&bytes)))
+    }
+
+    pub fn write<W: Write>(&self, mut w: W) -> io::Result<()> {
+        w.write_all(MAGIC)?;
+        w.write_u64::<LittleEndian>(self.0.len() as u64)?; // Number of bits
+
+        let bytes = self.0.to_bytes();
+        w.write_all(&bytes)?;
+
+        Ok(())
+    }
+}
+
+#[inline]
+fn div_ceil(a: usize, b: usize) -> usize {
+    let (q, r) = (a / b, a % b);
+    if r == 0 {
+        q
+    } else {
+        q + 1
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    const REFERENCE: &[u8] = &[
+        b'Z', b'P', b'C', b'T', // magic
+        0x0F,0,0,0,0,0,0,0, // length - 15 bits
+        0b10100000, 0b00010010, // bits
+    ];
+
+    #[test]
+    fn test_const_tracker_file_read() {
+        let ct = ConstTrackerFile::read(REFERENCE).unwrap();
+
+        let bytes = ct.0.to_bytes();
+        assert_eq!(&bytes, &[0b10100000, 0b00010010]);
+    }
+
+    #[test]
+    fn test_const_tracker_file_write() {
+        let mut ct = ConstTrackerFile(BitVec::from_bytes(&[0b10100000, 0b00010010]));
+        ct.0.pop();
+
+        let mut buf = Vec::new();
+        ct.write(&mut buf);
+
+        assert_eq!(&buf, REFERENCE);
+    }
+}

fawkes-crypto/src/backend/r1cs/mod.rs

@@ -1,127 +1,13 @@
-use std::io::{self, Read, Write};
-
-use bit_vec::BitVec;
-use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
-pub use r1cs_file::*;
-pub use wtns_file::WtnsFile;
-
-use crate::circuit::cs::Gate;
-use crate::circuit::lc::Index;
-use crate::core::signal::Signal;
-use crate::ff_uint::{Num, PrimeField, Uint};
-
-// TODO: Separate into multiple modules?
-
-const C_MAGIC: &[u8; 8] = b"ZPCCONST";
-
-pub struct ConstTracker(pub BitVec);
-
-impl ConstTracker {
-    pub fn read<R: Read>(mut r: R) -> io::Result<Self> {
-        let mut magic = [u8; 8];
-        r.read_exact(&mut magic)?;
-
-        if magic != *C_MAGIC {
-            return Err(io::Error::new(
-                io::ErrorKind::InvalidData,
-                "Invalid magic sequence",
-            ));
-        }
-
-        let len = r.read_u64::<LittleEndian>()?;
-        let mut bytes = vec![0; len as usize];
-        r.read_exact(&mut bytes);
-
-        Ok(ConstTracker(BitVec::from_bytes(&bytes)))
-    }
-
-    pub fn write<W: Write>(&self, mut w: W) -> io::Result<()> {
-        w.write_all(MAGIC)?;
-        w.write_u64::<LittleEndian>(bytes.len() as u64)?;
-
-        let bytes = self.0.to_bytes();
-        w.write_all(&bytes)?;
-
-        Ok(())
-    }
-}
-
-pub fn get_r1cs_file<Fr: PrimeField, const FS: usize>(
-    gates: &Vec<Gate<Fr>>,
-    consts: &BitVec,
-) -> (R1csFile<FS>, ConstTracker) {
-    use r1cs_file::*;
-
-    let mut n_pub_in = 0;
-    let mut n_prvt_in = 0;
-
-    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;
-                        i
-                    }
-                    Index::Aux(i) => {
-                        n_prvt_in += 1;
-                        i
-                    }
-                };
-
-                let mut c_bytes = [0; FS];
-                c.0.to_uint().put_little_endian(&mut c_bytes);
-
-                (FieldElement::from(c_bytes), i as u32)
-            };
-
-            let a = gate.0.iter().map(&mut map_comb).collect();
-            let b = gate.1.iter().map(&mut map_comb).collect();
-            let c = gate.2.iter().map(&mut map_comb).collect();
-
-            Constraint(a, b, c)
-        })
-        .collect();
-
-    let mut prime_bytes = [0; FS];
-    Fr::MODULUS.put_little_endian(&mut prime_bytes);
-
-    let r1cs_file = R1csFile {
-        header: Header {
-            prime: FieldElement::from(prime_bytes),
-            n_wires: n_pub_in + n_prvt_in,
-            n_pub_out: 0,
-            n_pub_in,
-            n_prvt_in,
-            n_labels: 0,
-            n_constraints: gates.len() as u32,
-        },
-        constraints: Constraints(constraints),
-        map: WireMap(Vec::new()),
-    };
-
-    let consts = ConstTracker(consts.clone());
-
-    (r1cs_file, consts)
-}
-
+#[cfg(feature = "r1cs-file")]
+pub use self::const_tracker::*;
+#[cfg(feature = "r1cs-file")]
+pub use self::r1cs_file::*;
 #[cfg(feature = "wtns-file")]
-pub fn get_witness<
-    'a,
-    Fr: PrimeField,
-    Pub: Signal<WitnessCS<'a, Fr>>,
-    Sec: Signal<WitnessCS<'a, Fr>>,
-    const FS: usize,
->(
-    gates: &Vec<Gate<Fr>>,
-    consts: &BitVec,
-    input_pub: &Pub::Value,
-    input_sec: &Pub::Value,
-) -> WtnsFile {
-    let cs = WitnessCS::rc_new(gates, consts);
+pub use self::wtns_file::*;
 
-    let signal_pub = Pub::alloc(cs, Some(input_pub));
-    signal_pub.inputize();
-    let signal_sec = Sec::alloc(cs, Some(input_sec));
-}
+#[cfg(feature = "r1cs-file")]
+mod const_tracker;
+#[cfg(feature = "r1cs-file")]
+mod r1cs_file;
+#[cfg(feature = "wtns-file")]
+mod wtns_file;

fawkes-crypto/src/backend/r1cs/r1cs_file.rs

@@ -0,0 +1,98 @@
+use bit_vec::BitVec;
+pub use r1cs_file::*;
+
+use crate::circuit::cs::Gate;
+use crate::circuit::lc::Index;
+use crate::ff_uint::{Num, PrimeField, Uint};
+use crate::backend::r1cs::ConstTrackerFile;
+
+pub fn get_r1cs_file<Fr: PrimeField, const FS: usize>(
+    gates: &Vec<Gate<Fr>>,
+    consts: &BitVec,
+) -> (R1csFile<FS>, ConstTrackerFile) {
+    use r1cs_file::*;
+
+    let mut n_pub_in = 0;
+    let mut n_prvt_in = 0;
+
+    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;
+                        i
+                    }
+                    Index::Aux(i) => {
+                        n_prvt_in += 1;
+                        i
+                    }
+                };
+
+                let mut c_bytes = [0; FS];
+                c.0.to_uint().put_little_endian(&mut c_bytes);
+
+                (FieldElement::from(c_bytes), i as u32)
+            };
+
+            let a = gate.0.iter().map(&mut map_comb).collect();
+            let b = gate.1.iter().map(&mut map_comb).collect();
+            let c = gate.2.iter().map(&mut map_comb).collect();
+
+            Constraint(a, b, c)
+        })
+        .collect();
+
+    let mut prime_bytes = [0; FS];
+    Fr::MODULUS.put_little_endian(&mut prime_bytes);
+
+    let r1cs_file = R1csFile {
+        header: Header {
+            prime: FieldElement::from(prime_bytes),
+            n_wires: n_pub_in + n_prvt_in,
+            n_pub_out: 0,
+            n_pub_in,
+            n_prvt_in,
+            n_labels: 0,
+            n_constraints: gates.len() as u32,
+        },
+        constraints: Constraints(constraints),
+        map: WireMap(Vec::new()),
+    };
+
+    let consts = ConstTrackerFile(consts.clone());
+
+    (r1cs_file, consts)
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::{
+        backend::bellman_groth16::{
+            engines::Bn256,
+            setup::setup
+        },
+        circuit::cs::CS,
+        circuit::num::CNum,
+        circuit::poseidon::{c_poseidon_merkle_proof_root, CMerkleProof},
+        core::signal::Signal,
+        native::poseidon::{PoseidonParams},
+    };
+    use crate::backend::r1cs::get_r1cs_file;
+    use crate::backend::bellman_groth16::engines::Engine;
+
+    #[test]
+    fn test_parameters_get_r1cs_file() {
+        fn circuit<C:CS>(public: CNum<C>, secret: (CNum<C>, CMerkleProof<C, 32>)) {
+            let poseidon_params = PoseidonParams::<C::Fr>::new(3, 8, 53);
+            let res = c_poseidon_merkle_proof_root(&secret.0, &secret.1, &poseidon_params);
+            res.assert_eq(&public);
+        }
+
+        let params = setup::<Bn256, _, _, _>(circuit);
+        let file = get_r1cs_file::<<Bn256 as Engine>::Fr, 32>(&params.1, &params.2);
+
+        assert!(true)
+    }
+}

fawkes-crypto/src/backend/r1cs/wtns_file.rs

@@ -0,0 +1,43 @@
+use bit_vec::BitVec;
+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;
+
+pub fn get_witness<
+    'a,
+    Fr: PrimeField,
+    Pub: Signal<WitnessCS<'a, Fr>>,
+    Sec: Signal<WitnessCS<'a, Fr>>,
+    const FS: usize,
+>(
+    gates: &'a Vec<Gate<Fr>>,
+    consts: &'a BitVec,
+    input_pub: &'a Pub::Value,
+    input_sec: &'a Sec::Value,
+) -> WtnsFile<FS> {
+    let cs = WitnessCS::rc_new(gates, consts);
+
+    let mut prime_bytes = [0; FS];
+    Fr::MODULUS.put_little_endian(&mut prime_bytes);
+
+    let signal_pub = Pub::alloc(&cs, Some(input_pub));
+    signal_pub.inputize();
+    let _signal_sec = Sec::alloc(&cs, Some(input_sec));
+
+    let cs = cs.borrow_mut();
+    let mut witness = Vec::with_capacity(cs.num_aux());
+
+    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);
+
+        witness.push(fe);
+    }
+
+    WtnsFile::from_vec(witness, FieldElement::from(prime_bytes))
+}