From c1db2b48666f1f735b977b98ab1ce4e5dd79e862 Mon Sep 17 00:00:00 2001
From: Michael Vines <mvines@gmail.com>
Date: Tue, 26 Oct 2021 23:04:30 -0700
Subject: [PATCH] Wrap a struct around the discrete log precompute hashmap

---
 zk-token-sdk/src/encryption/discrete_log.rs | 20 +++++++++++---------
 zk-token-sdk/src/encryption/elgamal.rs      | 10 +++++-----
 2 files changed, 16 insertions(+), 14 deletions(-)

diff --git a/zk-token-sdk/src/encryption/discrete_log.rs b/zk-token-sdk/src/encryption/discrete_log.rs
index f15bb85fe..1279148d4 100644
--- a/zk-token-sdk/src/encryption/discrete_log.rs
+++ b/zk-token-sdk/src/encryption/discrete_log.rs
@@ -23,8 +23,11 @@ pub struct DiscreteLog {
     pub target: RistrettoPoint,
 }
 
+#[derive(Serialize, Deserialize)]
+pub struct DecodeU32Precomputation(HashMap<[u8; 32], u32>);
+
 /// Builds a HashMap of 2^18 elements
-pub fn decode_u32_precomputation(generator: RistrettoPoint) -> HashMap<[u8; 32], u32> {
+fn decode_u32_precomputation(generator: RistrettoPoint) -> DecodeU32Precomputation {
     let mut hashmap = HashMap::new();
 
     let two12_scalar = Scalar::from(TWO14);
@@ -46,12 +49,12 @@ pub fn decode_u32_precomputation(generator: RistrettoPoint) -> HashMap<[u8; 32],
     });
     println!("     [8/8] completed");
 
-    hashmap
+    DecodeU32Precomputation(hashmap)
 }
 
 /// Pre-compute HashMap needed for decryption. The HashMap is independent of (works for) any key.
 #[allow(non_snake_case)]
-pub fn decode_u32_precomputation_for_G() -> HashMap<[u8; 32], u32> {
+pub fn decode_u32_precomputation_for_G() -> DecodeU32Precomputation {
     decode_u32_precomputation(G)
 }
 
@@ -59,22 +62,21 @@ pub fn decode_u32_precomputation_for_G() -> HashMap<[u8; 32], u32> {
 impl DiscreteLog {
     /// Solves the discrete log problem under the assumption that the solution
     /// is a 32-bit number.
-    pub fn decode_u32(self) -> Option<u32> {
-        let hashmap = decode_u32_precomputation(self.generator);
-        self.decode_u32_online(&hashmap)
+    pub(crate) fn decode_u32(self) -> Option<u32> {
+        self.decode_u32_online(&decode_u32_precomputation(self.generator))
     }
 
     /// Solves the discrete log instance using the pre-computed HashMap by enumerating through 2^14
     /// possible solutions
-    pub fn decode_u32_online(self, hashmap: &HashMap<[u8; 32], u32>) -> Option<u32> {
+    pub fn decode_u32_online(self, hashmap: &DecodeU32Precomputation) -> Option<u32> {
         // iterator for 0G, -1G, -2G, ...
         let ristretto_iter = RistrettoIterator::new(self.target, -self.generator);
 
         let mut decoded = None;
         ristretto_iter.zip(0..TWO14).for_each(|(elem, x_lo)| {
             let key = elem.compress().to_bytes();
-            if hashmap.contains_key(&key) {
-                let x_hi = hashmap[&key];
+            if hashmap.0.contains_key(&key) {
+                let x_hi = hashmap.0[&key];
                 decoded = Some(x_lo + TWO14 * x_hi);
             }
         });
diff --git a/zk-token-sdk/src/encryption/elgamal.rs b/zk-token-sdk/src/encryption/elgamal.rs
index 908595e3a..e7ed5d4ab 100644
--- a/zk-token-sdk/src/encryption/elgamal.rs
+++ b/zk-token-sdk/src/encryption/elgamal.rs
@@ -1,6 +1,6 @@
 use {
     crate::encryption::{
-        discrete_log::DiscreteLog,
+        discrete_log::{DecodeU32Precomputation, DiscreteLog},
         pedersen::{
             Pedersen, PedersenBase, PedersenCommitment, PedersenDecryptHandle, PedersenOpening,
         },
@@ -19,7 +19,7 @@ use {
         signature::Signature,
         signer::{Signer, SignerError},
     },
-    std::{collections::HashMap, convert::TryInto},
+    std::convert::TryInto,
     subtle::{Choice, ConstantTimeEq},
     zeroize::Zeroize,
 };
@@ -120,7 +120,7 @@ impl ElGamal {
     fn decrypt_u32_online(
         secret: &ElGamalSecretKey,
         ct: &ElGamalCiphertext,
-        hashmap: &HashMap<[u8; 32], u32>,
+        hashmap: &DecodeU32Precomputation,
     ) -> Option<u32> {
         let discrete_log_instance = Self::decrypt(secret, ct);
         discrete_log_instance.decode_u32_online(hashmap)
@@ -337,7 +337,7 @@ impl ElGamalSecretKey {
     pub fn decrypt_u32_online(
         &self,
         ct: &ElGamalCiphertext,
-        hashmap: &HashMap<[u8; 32], u32>,
+        hashmap: &DecodeU32Precomputation,
     ) -> Option<u32> {
         ElGamal::decrypt_u32_online(self, ct, hashmap)
     }
@@ -429,7 +429,7 @@ impl ElGamalCiphertext {
     pub fn decrypt_u32_online(
         &self,
         secret: &ElGamalSecretKey,
-        hashmap: &HashMap<[u8; 32], u32>,
+        hashmap: &DecodeU32Precomputation,
     ) -> Option<u32> {
         ElGamal::decrypt_u32_online(secret, self, hashmap)
     }