From 8a32e77949b12b94f8da46b6c4014cec389b0162 Mon Sep 17 00:00:00 2001
From: Sean Bowe <ewillbefull@gmail.com>
Date: Sat, 17 Sep 2016 12:02:05 -0600
Subject: [PATCH] K query in parallel.

---
 src/coordinator.rs        |  4 ++--
 src/protocol/mod.rs       | 25 +++++++++++++------------
 src/protocol/multicore.rs | 11 +++++++++++
 src/verifier.rs           |  2 +-
 4 files changed, 27 insertions(+), 15 deletions(-)

diff --git a/src/coordinator.rs b/src/coordinator.rs
index bd1f2bd..80745ab 100644
--- a/src/coordinator.rs
+++ b/src/coordinator.rs
@@ -200,9 +200,9 @@ impl ConnectionHandler {
             }
         }
 
-        info!("Initializing stage3 with stage2");
+        info!("Initializing stage3 with constraint system and stage2");
 
-        let mut stage3 = Stage3Contents::new(&stage2);
+        let mut stage3 = Stage3Contents::new(&cs, &stage2);
         for (pubkey, peerid) in pubkeys.iter().zip(peers.iter()) {
             info!("Sending stage3 to peerid={}", peerid.to_hex());
 
diff --git a/src/protocol/mod.rs b/src/protocol/mod.rs
index 2861566..937d623 100644
--- a/src/protocol/mod.rs
+++ b/src/protocol/mod.rs
@@ -236,20 +236,21 @@ pub struct Stage3Contents {
 }
 
 impl Stage3Contents {
-    pub fn new(stage2: &Stage2Contents) -> Self {
-        let mut pk_k = Vec::with_capacity(stage2.pk_a.len()+3);
+    pub fn new(cs: &CS, stage2: &Stage2Contents) -> Self {
+        assert_eq!(stage2.pk_a.len(), cs.num_vars + 1);
+        assert_eq!(stage2.pk_b_temp.len(), cs.num_vars + 1);
+        assert_eq!(stage2.pk_c.len(), cs.num_vars + 1);
 
-        for ((a, b), c) in stage2.pk_a.iter().take(stage2.pk_a.len() - 1)
-                           .zip(stage2.pk_b_temp.iter().take(stage2.pk_b_temp.len() - 1))
-                           .zip(stage2.pk_c.iter().take(stage2.pk_c.len() - 1))
-        {
-            pk_k.push(*a + *b + *c);
-        }
+        let mut pk_k = Vec::with_capacity(cs.num_vars + 3);
 
         // Perform Z extention as libsnark does.
-        pk_k.push(stage2.pk_a[stage2.pk_a.len() - 1]);
-        pk_k.push(stage2.pk_b_temp[stage2.pk_b_temp.len() - 1]);
-        pk_k.push(stage2.pk_c[stage2.pk_c.len() - 1]);
+        pk_k.extend_from_slice(&stage2.pk_a);
+        pk_k.push(stage2.pk_b_temp[cs.num_vars]);
+        pk_k.push(stage2.pk_c[cs.num_vars]);
+
+        // Add B and C
+        add_all_to(&mut pk_k[0..cs.num_vars], &stage2.pk_b_temp[0..cs.num_vars]);
+        add_all_to(&mut pk_k[0..cs.num_vars], &stage2.pk_c[0..cs.num_vars]);
 
         Stage3Contents {
             vk_gamma: G2::one(),
@@ -349,7 +350,7 @@ fn compare_to_libsnark_generate() {
     }
 
     // Stage 3
-    let mut stage3 = Stage3Contents::new(&stage2);
+    let mut stage3 = Stage3Contents::new(&cs, &stage2);
     for (private, public) in privkeys.iter().zip(pubkeys.iter()) {
         let prev = stage3.clone();
         stage3.transform(private);
diff --git a/src/protocol/multicore.rs b/src/protocol/multicore.rs
index 96efe85..120411a 100644
--- a/src/protocol/multicore.rs
+++ b/src/protocol/multicore.rs
@@ -25,3 +25,14 @@ pub fn mul_all_by<G: Group>(v: &mut [G], c: Fr) {
         }
     }, ::THREADS);
 }
+
+pub fn add_all_to<G: Group>(v: &mut [G], other: &[G]) {
+    assert_eq!(v.len(), other.len());
+
+    parallel(v, |mut i, v| {
+        for a in v {
+            *a = *a + other[i];
+            i += 1;
+        }
+    }, ::THREADS);
+}
diff --git a/src/verifier.rs b/src/verifier.rs
index ba4fb47..2150b5b 100644
--- a/src/verifier.rs
+++ b/src/verifier.rs
@@ -68,7 +68,7 @@ fn main() {
         stage2 = new_stage;
     }
 
-    let mut stage3 = Stage3Contents::new(&stage2);
+    let mut stage3 = Stage3Contents::new(&cs, &stage2);
 
     for i in 0..num_players {
         let new_stage: Stage3Contents = decode_from(&mut f, Infinite).unwrap();