From 639c40047f074f424df1740522fd5c6031c133dc Mon Sep 17 00:00:00 2001
From: Jack Grigg <jack@z.cash>
Date: Fri, 6 May 2016 12:17:42 +1200
Subject: [PATCH] Use comparator object for sorting StepRows

---
 src/crypto/equihash.cpp | 18 +++++++++++++-----
 src/crypto/equihash.h   | 16 +++++++++++++---
 2 files changed, 26 insertions(+), 8 deletions(-)

diff --git a/src/crypto/equihash.cpp b/src/crypto/equihash.cpp
index e66e8e058..b03258025 100644
--- a/src/crypto/equihash.cpp
+++ b/src/crypto/equihash.cpp
@@ -227,6 +227,7 @@ std::set<std::vector<eh_index>> Equihash<N,K>::BasicSolve(const eh_HashState& ba
 
     // 1) Generate first list
     LogPrint("pow", "Generating first list\n");
+    size_t hashLen = N/8;
     std::vector<FullStepRow> X;
     X.reserve(init_size);
     for (eh_index i = 0; i < init_size; i++) {
@@ -238,7 +239,7 @@ std::set<std::vector<eh_index>> Equihash<N,K>::BasicSolve(const eh_HashState& ba
         LogPrint("pow", "Round %d:\n", r);
         // 2a) Sort the list
         LogPrint("pow", "- Sorting list\n");
-        std::sort(X.begin(), X.end());
+        std::sort(X.begin(), X.end(), CompareSR(hashLen));
 
         LogPrint("pow", "- Finding collisions\n");
         int i = 0;
@@ -284,6 +285,8 @@ std::set<std::vector<eh_index>> Equihash<N,K>::BasicSolve(const eh_HashState& ba
             X.erase(X.begin()+posFree, X.end());
             X.shrink_to_fit();
         }
+
+        hashLen -= CollisionByteLength;
     }
 
     // k+1) Find a collision on last 2n(k+1) bits
@@ -291,7 +294,7 @@ std::set<std::vector<eh_index>> Equihash<N,K>::BasicSolve(const eh_HashState& ba
     std::set<std::vector<eh_index>> solns;
     if (X.size() > 1) {
         LogPrint("pow", "- Sorting list\n");
-        std::sort(X.begin(), X.end());
+        std::sort(X.begin(), X.end(), CompareSR(hashLen));
         LogPrint("pow", "- Finding collisions\n");
         for (int i = 0; i < X.size() - 1; i++) {
             FullStepRow res(X[i], X[i+1], 0);
@@ -369,6 +372,7 @@ std::set<std::vector<eh_index>> Equihash<N,K>::OptimisedSolve(const eh_HashState
 
         // 1) Generate first list
         LogPrint("pow", "Generating first list\n");
+        size_t hashLen = N/8;
         std::vector<TruncatedStepRow> Xt;
         Xt.reserve(init_size);
         for (eh_index i = 0; i < init_size; i++) {
@@ -380,7 +384,7 @@ std::set<std::vector<eh_index>> Equihash<N,K>::OptimisedSolve(const eh_HashState
             LogPrint("pow", "Round %d:\n", r);
             // 2a) Sort the list
             LogPrint("pow", "- Sorting list\n");
-            std::sort(Xt.begin(), Xt.end());
+            std::sort(Xt.begin(), Xt.end(), CompareSR(hashLen));
 
             LogPrint("pow", "- Finding collisions\n");
             int i = 0;
@@ -425,13 +429,15 @@ std::set<std::vector<eh_index>> Equihash<N,K>::OptimisedSolve(const eh_HashState
                 Xt.erase(Xt.begin()+posFree, Xt.end());
                 Xt.shrink_to_fit();
             }
+
+            hashLen -= CollisionByteLength;
         }
 
         // k+1) Find a collision on last 2n(k+1) bits
         LogPrint("pow", "Final round:\n");
         if (Xt.size() > 1) {
             LogPrint("pow", "- Sorting list\n");
-            std::sort(Xt.begin(), Xt.end());
+            std::sort(Xt.begin(), Xt.end(), CompareSR(hashLen));
             LogPrint("pow", "- Finding collisions\n");
             for (int i = 0; i < Xt.size() - 1; i++) {
                 TruncatedStepRow res(Xt[i], Xt[i+1], 0);
@@ -453,6 +459,7 @@ std::set<std::vector<eh_index>> Equihash<N,K>::OptimisedSolve(const eh_HashState
     int invalidCount = 0;
     for (eh_trunc* partialSoln : partialSolns) {
         // 1) Generate first list of possibilities
+        size_t hashLen = N/8;
         std::vector<std::vector<FullStepRow>> X;
         X.reserve(soln_size);
         for (eh_index i = 0; i < soln_size; i++) {
@@ -476,7 +483,7 @@ std::set<std::vector<eh_index>> Equihash<N,K>::OptimisedSolve(const eh_HashState
                 std::vector<FullStepRow> ic(X[v]);
                 ic.reserve(X[v].size() + X[v+1].size());
                 ic.insert(ic.end(), X[v+1].begin(), X[v+1].end());
-                std::sort(ic.begin(), ic.end());
+                std::sort(ic.begin(), ic.end(), CompareSR(hashLen));
                 CollideBranches(ic, CollisionByteLength, CollisionBitLength + 1, partialSoln[(1<<r)*v], partialSoln[(1<<r)*(v+1)]);
 
                 // 2v) Check if this has become an invalid solution
@@ -487,6 +494,7 @@ std::set<std::vector<eh_index>> Equihash<N,K>::OptimisedSolve(const eh_HashState
             }
 
             X = Xc;
+            hashLen -= CollisionByteLength;
         }
 
         // We are at the top of the tree
diff --git a/src/crypto/equihash.h b/src/crypto/equihash.h
index 22b8cc898..642c43d62 100644
--- a/src/crypto/equihash.h
+++ b/src/crypto/equihash.h
@@ -23,6 +23,8 @@ typedef uint8_t eh_trunc;
 
 class StepRow
 {
+    friend class CompareSR;
+
 protected:
     unsigned char* hash;
     unsigned int len;
@@ -36,12 +38,20 @@ public:
     bool IsZero();
     std::string GetHex() { return HexStr(hash, hash+len); }
 
-    friend inline bool operator==(const StepRow& a, const StepRow& b) { return memcmp(a.hash, b.hash, a.len) == 0; }
-    friend inline bool operator<(const StepRow& a, const StepRow& b) { return memcmp(a.hash, b.hash, a.len) < 0; }
-
     friend bool HasCollision(StepRow& a, StepRow& b, int l);
 };
 
+class CompareSR
+{
+private:
+    size_t len;
+
+public:
+    CompareSR(size_t l) : len {l} { }
+
+    inline bool operator()(const StepRow& a, const StepRow& b) { return memcmp(a.hash, b.hash, len) < 0; }
+};
+
 bool HasCollision(StepRow& a, StepRow& b, int l);
 
 class FullStepRow : public StepRow