From dc77ff484715eac8bbd5afcbc4e0c46a9d115342 Mon Sep 17 00:00:00 2001
From: Duncan Tebbs <duncan.tebbs@gmail.com>
Date: Thu, 14 Jan 2021 17:50:56 +0000
Subject: [PATCH] libzeth: make joinsplit a bit more composable

---
 libzeth/circuits/circuit_wrapper.tcc |  3 +++
 libzeth/circuits/joinsplit.tcc       | 15 ++++++++-------
 2 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/libzeth/circuits/circuit_wrapper.tcc b/libzeth/circuits/circuit_wrapper.tcc
index 5d6379894..d30a93c8c 100644
--- a/libzeth/circuits/circuit_wrapper.tcc
+++ b/libzeth/circuits/circuit_wrapper.tcc
@@ -27,8 +27,11 @@ circuit_wrapper<
     NumOutputs,
     TreeDepth>::circuit_wrapper()
 {
+    // Joinsplit gadget internally allocates its public data first.
     // TODO: joinsplit_gadget should be refactored to be properly composable.
     joinsplit = std::make_shared<joinsplit_type>(pb);
+    const size_t num_public_elements = joinsplit->get_num_public_elements();
+    pb.set_input_sizes(num_public_elements);
 
     // Generate constraints
     joinsplit->generate_r1cs_constraints();
diff --git a/libzeth/circuits/joinsplit.tcc b/libzeth/circuits/joinsplit.tcc
index edbd35ed4..b4c41f3fb 100644
--- a/libzeth/circuits/joinsplit.tcc
+++ b/libzeth/circuits/joinsplit.tcc
@@ -192,10 +192,10 @@ public:
             // represented. The aggregation of these bits plus of value_pub_in,
             // and value_pub_out take `nb_field_residual` field element(s) to be
             // represented
-            const size_t nb_packed_inputs =
+            const size_t num_packed_public_elements =
                 2 * NumInputs + 1 + nb_field_residual;
-            const size_t nb_inputs = 1 + NumOutputs + nb_packed_inputs;
-            pb.set_input_sizes(nb_inputs);
+            const size_t num_public_elements =
+                1 + NumOutputs + num_packed_public_elements;
             // ---------------------------------------------------------------
 
             ZERO.allocate(pb, FMT(this->annotation_prefix, " ZERO"));
@@ -277,14 +277,15 @@ public:
             // since we are packing all the inputs nullifiers + the h_is +
             // + the h_sig + the residual bits
             assert(packed_inputs.size() == NumInputs + 1 + NumInputs + 1);
-            assert(nb_packed_inputs == [this]() {
+            assert(num_packed_public_elements == [this]() {
                 size_t sum = 0;
                 for (const auto &i : packed_inputs) {
                     sum = sum + i.size();
                 }
                 return sum;
             }());
-            assert(nb_inputs == get_inputs_field_element_size());
+            assert(num_public_elements == get_num_public_elements());
+            (void)num_public_elements;
 
             // [SANITY CHECK] Total size of unpacked inputs
             size_t total_size_unpacked_inputs = 0;
@@ -599,8 +600,8 @@ public:
         return get_inputs_bit_size() - (1 + NumOutputs) * FieldT::capacity();
     }
 
-    // Computes the number of field elements in the primary inputs
-    static size_t get_inputs_field_element_size()
+    // Computes the number of field elements in the public data
+    static size_t get_num_public_elements()
     {
         size_t nb_elements = 0;