From 20abe2083c60bc46d2d3d410d5ffb5fe82f3513d Mon Sep 17 00:00:00 2001
From: Jack Grigg <jack@z.cash>
Date: Sun, 14 Aug 2016 01:58:08 +1200
Subject: [PATCH] Extend byte array expansion and compression methods with
 optional padding

---
 src/crypto/equihash.cpp     | 15 +++++++++------
 src/crypto/equihash.h       |  4 ++--
 src/gtest/test_equihash.cpp | 18 ++++++++++++------
 3 files changed, 23 insertions(+), 14 deletions(-)

diff --git a/src/crypto/equihash.cpp b/src/crypto/equihash.cpp
index 9bf5a9c11..85d979355 100644
--- a/src/crypto/equihash.cpp
+++ b/src/crypto/equihash.cpp
@@ -53,12 +53,12 @@ void GenerateHash(const eh_HashState& base_state, eh_index g,
 
 void ExpandArray(const unsigned char* in, size_t in_len,
                  unsigned char* out, size_t out_len,
-                 size_t bit_len)
+                 size_t bit_len, size_t byte_pad)
 {
     assert(bit_len >= 8);
     assert(8*sizeof(uint32_t) >= 7+bit_len);
 
-    size_t out_width { (bit_len+7)/8 };
+    size_t out_width { (bit_len+7)/8 + byte_pad };
     assert(out_len == 8*out_width*in_len/bit_len);
 
     uint32_t bit_len_mask { ((uint32_t)1 << bit_len) - 1 };
@@ -77,7 +77,10 @@ void ExpandArray(const unsigned char* in, size_t in_len,
         // output element.
         if (acc_bits >= bit_len) {
             acc_bits -= bit_len;
-            for (size_t x = 0; x < out_width; x++) {
+            for (size_t x = 0; x < byte_pad; x++) {
+                out[j+x] = 0;
+            }
+            for (size_t x = byte_pad; x < out_width; x++) {
                 out[j+x] = (
                     // Big-endian
                     acc_value >> (acc_bits+(8*(out_width-x-1)))
@@ -93,12 +96,12 @@ void ExpandArray(const unsigned char* in, size_t in_len,
 
 void CompressArray(const unsigned char* in, size_t in_len,
                    unsigned char* out, size_t out_len,
-                   size_t bit_len)
+                   size_t bit_len, size_t byte_pad)
 {
     assert(bit_len >= 8);
     assert(8*sizeof(uint32_t) >= 7+bit_len);
 
-    size_t in_width { (bit_len+7)/8 };
+    size_t in_width { (bit_len+7)/8 + byte_pad };
     assert(out_len == bit_len*in_len/(8*in_width));
 
     uint32_t bit_len_mask { ((uint32_t)1 << bit_len) - 1 };
@@ -114,7 +117,7 @@ void CompressArray(const unsigned char* in, size_t in_len,
         // input element.
         if (acc_bits < 8) {
             acc_value = acc_value << bit_len;
-            for (size_t x = 0; x < in_width; x++) {
+            for (size_t x = byte_pad; x < in_width; x++) {
                 acc_value = acc_value | (
                     (
                         // Apply bit_len_mask across byte boundaries
diff --git a/src/crypto/equihash.h b/src/crypto/equihash.h
index c3b91d67c..69d03cf8e 100644
--- a/src/crypto/equihash.h
+++ b/src/crypto/equihash.h
@@ -26,10 +26,10 @@ typedef uint8_t eh_trunc;
 
 void ExpandArray(const unsigned char* in, size_t in_len,
                  unsigned char* out, size_t out_len,
-                 size_t bit_len);
+                 size_t bit_len, size_t byte_pad=0);
 void CompressArray(const unsigned char* in, size_t in_len,
                    unsigned char* out, size_t out_len,
-                   size_t bit_len);
+                   size_t bit_len, size_t byte_pad=0);
 
 eh_index ArrayToEhIndex(const unsigned char* array);
 eh_trunc TruncateIndex(const eh_index i, const unsigned int ilen);
diff --git a/src/gtest/test_equihash.cpp b/src/gtest/test_equihash.cpp
index 7289c1b4f..5afb6af77 100644
--- a/src/gtest/test_equihash.cpp
+++ b/src/gtest/test_equihash.cpp
@@ -4,31 +4,37 @@
 #include "crypto/equihash.h"
 #include "uint256.h"
 
-void TestExpandAndCompress(const std::string &scope, size_t bit_len,
+void TestExpandAndCompress(const std::string &scope, size_t bit_len, size_t byte_pad,
                            std::vector<unsigned char> compact,
                            std::vector<unsigned char> expanded)
 {
     SCOPED_TRACE(scope);
 
     std::vector<unsigned char> out(expanded.size());
-    ExpandArray(compact.data(), compact.size(), out.data(), out.size(), bit_len);
+    ExpandArray(compact.data(), compact.size(),
+                out.data(), out.size(), bit_len, byte_pad);
     EXPECT_EQ(expanded, out);
 
     out.resize(compact.size());
-    CompressArray(expanded.data(), expanded.size(), out.data(), out.size(), bit_len);
+    CompressArray(expanded.data(), expanded.size(),
+                  out.data(), out.size(), bit_len, byte_pad);
     EXPECT_EQ(compact, out);
 }
 
 TEST(equihash_tests, expand_and_contract_arrays) {
-    TestExpandAndCompress("8 11-bit chunks, all-ones", 11,
+    TestExpandAndCompress("8 11-bit chunks, all-ones", 11, 0,
                           ParseHex("ffffffffffffffffffffff"),
                           ParseHex("07ff07ff07ff07ff07ff07ff07ff07ff"));
-    TestExpandAndCompress("8 21-bit chunks, alternating 1s and 0s", 21,
+    TestExpandAndCompress("8 21-bit chunks, alternating 1s and 0s", 21, 0,
                           ParseHex("aaaaad55556aaaab55555aaaaad55556aaaab55555"),
                           ParseHex("155555155555155555155555155555155555155555155555"));
-    TestExpandAndCompress("16 14-bit chunks, alternating 11s and 00s", 14,
+    TestExpandAndCompress("16 14-bit chunks, alternating 11s and 00s", 14, 0,
                           ParseHex("cccf333cccf333cccf333cccf333cccf333cccf333cccf333cccf333"),
                           ParseHex("3333333333333333333333333333333333333333333333333333333333333333"));
+
+    TestExpandAndCompress("8 11-bit chunks, all-ones, 2-byte padding", 11, 2,
+                          ParseHex("ffffffffffffffffffffff"),
+                          ParseHex("000007ff000007ff000007ff000007ff000007ff000007ff000007ff000007ff"));
 }
 
 TEST(equihash_tests, is_probably_duplicate) {