From ffb544e4a543be7e9fcad96a3645dd3d79942e96 Mon Sep 17 00:00:00 2001
From: Nicholas Ward <npward@berkeley.edu>
Date: Mon, 4 Oct 2021 14:17:19 -0700
Subject: [PATCH] initial non-native add

---
 src/gadgets/arithmetic_u32.rs | 23 +++++++++++++++++------
 src/gadgets/nonnative.rs      | 22 +++++++++++++++-------
 2 files changed, 32 insertions(+), 13 deletions(-)

diff --git a/src/gadgets/arithmetic_u32.rs b/src/gadgets/arithmetic_u32.rs
index 4a8cff42..d4ff5896 100644
--- a/src/gadgets/arithmetic_u32.rs
+++ b/src/gadgets/arithmetic_u32.rs
@@ -18,6 +18,10 @@ impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
         U32Target(self.add_virtual_target())
     }
 
+    pub fn add_virtual_u32_targets(&self, n: usize) -> Vec<U32Target> {
+        self.add_virtual_targets(n).iter().cloned().map(U32Target).collect()
+    }
+
     pub fn zero_u32(&self) -> U32Target {
         U32Target(self.zero())
     }
@@ -40,17 +44,17 @@ impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
                 let gate_index = self.add_gate(gate.clone(), vec![]);
                 (gate_index, 0)
             },
-            Some((gate_index, copy) => (gate_index, copy),
+            Some((gate_index, copy)) => (gate_index, copy),
         };
 
         let output_low = self.add_virtual_u32_target();
         let output_high = self.add_virtual_u32_target();
 
-        self.connect(Target::wire(gate_index, gate.wire_ith_multiplicand_0(copy)), x);
-        self.connect(Target::wire(gate_index, gate.wire_ith_multiplicand_1(copy)), y);
-        self.connect(Target::wire(gate_index, gate.wire_ith_addend(copy)), z);
-        self.connect(Target::wire(gate_index, gate.wire_ith_output_low_half(copy)), output_low);
-        self.connect(Target::wire(gate_index, gate.wire_ith_output_high_half(copy)), output_high);
+        self.connect(Target::wire(gate_index, U32ArithmeticGate::<F, D>::wire_ith_multiplicand_0(copy)), x.0);
+        self.connect(Target::wire(gate_index, U32ArithmeticGate::<F, D>::wire_ith_multiplicand_1(copy)), y.0);
+        self.connect(Target::wire(gate_index, U32ArithmeticGate::<F, D>::wire_ith_addend(copy)), z.0);
+        self.connect(Target::wire(gate_index, U32ArithmeticGate::<F, D>::wire_ith_output_low_half(copy)), output_low.0);
+        self.connect(Target::wire(gate_index, U32ArithmeticGate::<F, D>::wire_ith_output_high_half(copy)), output_high.0);
 
         self.current_u32_arithmetic_gate = Some((gate_index, 0));
 
@@ -61,6 +65,13 @@ impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
         self.add_mul_u32(a, self.one_u32(), b)
     }
 
+    pub fn add_three_u32(&mut self, a: U32Target, b: U32Target, c: U32Target) -> (U32Target, U32Target) {
+        let (init_low, carry1) = self.add_u32(a, b);
+        let (final_low, carry2) = self.add_u32(c, init_low);
+        let (combined_carry, _zero) = self.add_u32(carry1, carry2);
+        (final_low, combined_carry)
+    }
+
     pub fn mul_u32(&mut self, a: U32Target, b: U32Target) -> (U32Target, U32Target) {
         self.add_mul_u32(a, b, self.zero_u32())
     }
diff --git a/src/gadgets/nonnative.rs b/src/gadgets/nonnative.rs
index 1407360e..01d60c85 100644
--- a/src/gadgets/nonnative.rs
+++ b/src/gadgets/nonnative.rs
@@ -1,3 +1,4 @@
+use num::bigint::BigUint;
 use std::collections::BTreeMap;
 use std::marker::PhantomData;
 
@@ -14,7 +15,7 @@ use crate::util::bimap::bimap_from_lists;
 
 pub struct NonNativeTarget {
     /// The modulus of the field F' being represented.
-    modulus: BigUInt,
+    modulus: BigUint,
     /// These F elements are assumed to contain 32-bit values.
     limbs: Vec<U32Target>,
 }
@@ -26,14 +27,21 @@ impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
         debug_assert!(b.modulus == modulus);
         debug_assert!(b.limbs.len() == num_limbs);
 
-        let mut combined_limbs = self.add_virtual_targets(num_limbs + 1);
-        let mut carry = self.zero();
+        let mut combined_limbs = self.add_virtual_u32_targets(num_limbs + 1);
+        let mut carry = self.zero_u32();
         for i in 0..num_limbs {
-            
+            let (new_limb, carry) = self.add_three_u32(carry, a.limbs[i], b.limbs[i]);
+            combined_limbs[i] = new_limb;
+        }
+        combined_limbs[num_limbs] = carry;
+        
+        NonNativeTarget {
+            modulus,
+            limbs: combined_limbs,
         }
     }
 
-    pub fn reduce_add_result(&mut self, limbs: Vec<Target>, modulus: BigUInt) -> Vec<Target> {
+    pub fn reduce_add_result(&mut self, limbs: Vec<Target>, modulus: BigUint) -> Vec<Target> {
         todo!()
     }
 
@@ -46,8 +54,8 @@ impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
         let mut combined_limbs = self.add_virtual_targets(2 * num_limbs - 1);
         for i in 0..num_limbs {
             for j in 0..num_limbs {
-                let sum = builder.add(a.limbs[i], b.limbs[j]);
-                combined_limbs[i + j] = builder.add(combined_limbs[i + j], sum);
+                let sum = self.add(a.limbs[i], b.limbs[j]);
+                combined_limbs[i + j] = self.add(combined_limbs[i + j], sum);
             }
         }