From e5f7c49d55d668aca9d3690c14e1fe372a7f4a93 Mon Sep 17 00:00:00 2001
From: Sean Bowe <ewillbefull@gmail.com>
Date: Wed, 4 May 2016 18:26:05 -0600
Subject: [PATCH] zkSNARK: Ensure that values balance correctly.

---
 src/zcash/circuit/gadget.tcc | 36 ++++++++++++++++++++++++++++++++++--
 src/zcash/circuit/utils.tcc  |  9 ++++++++-
 2 files changed, 42 insertions(+), 3 deletions(-)

diff --git a/src/zcash/circuit/gadget.tcc b/src/zcash/circuit/gadget.tcc
index 9ce142143..7872a54ff 100644
--- a/src/zcash/circuit/gadget.tcc
+++ b/src/zcash/circuit/gadget.tcc
@@ -144,6 +144,28 @@ public:
             // Constrain the JoinSplit output constraints.
             zk_output_notes[i]->generate_r1cs_constraints();
         }
+
+        // Value balance
+        {
+            linear_combination<FieldT> left_side = packed_addition(zk_vpub_old);
+            for (size_t i = 0; i < NumInputs; i++) {
+                left_side = left_side + packed_addition(zk_input_notes[i]->value);
+            }
+
+            linear_combination<FieldT> right_side = packed_addition(zk_vpub_new);
+            for (size_t i = 0; i < NumOutputs; i++) {
+                right_side = right_side + packed_addition(zk_output_notes[i]->value);
+            }
+
+            // Ensure that both sides are equal
+            this->pb.add_r1cs_constraint(r1cs_constraint<FieldT>(
+                1,
+                left_side,
+                right_side
+            ));
+
+            // TODO: #854
+        }
     }
 
     void generate_r1cs_witness(
@@ -158,6 +180,16 @@ public:
         // Witness `zero`
         this->pb.val(ZERO) = FieldT::zero();
 
+        // Witness public balance values
+        zk_vpub_old.fill_with_bits(
+            this->pb,
+            uint64_to_bool_vector(vpub_old)
+        );
+        zk_vpub_new.fill_with_bits(
+            this->pb,
+            uint64_to_bool_vector(vpub_new)
+        );
+
         // Witness phi
         zk_phi->bits.fill_with_bits(
             this->pb,
@@ -211,8 +243,8 @@ public:
             insert_uint256(verify_inputs, commitments[i]);
         }
 
-        insert_uint64(verify_inputs, 0); // TODO: vpub_old
-        insert_uint64(verify_inputs, 0); // TODO: vpub_new
+        insert_uint64(verify_inputs, vpub_old);
+        insert_uint64(verify_inputs, vpub_new);
 
         assert(verify_inputs.size() == verifying_input_bit_size());
         auto verify_field_elements = pack_bit_vector_into_field_element_vector<FieldT>(verify_inputs);
diff --git a/src/zcash/circuit/utils.tcc b/src/zcash/circuit/utils.tcc
index 3088362cb..efde63bc2 100644
--- a/src/zcash/circuit/utils.tcc
+++ b/src/zcash/circuit/utils.tcc
@@ -45,4 +45,11 @@ void insert_uint256(std::vector<bool>& into, uint256 from) {
 void insert_uint64(std::vector<bool>& into, uint64_t from) {
     std::vector<bool> num = uint64_to_bool_vector(from);
     into.insert(into.end(), num.begin(), num.end());
-}
\ No newline at end of file
+}
+
+template<typename FieldT>
+linear_combination<FieldT> packed_addition(pb_variable_array<FieldT>& input) {
+    return pb_packing_sum<FieldT>(pb_variable_array<FieldT>(
+        input.rbegin(), input.rend()
+    ));
+}