From 4fd2d2fc97e21efceab849576e544160fd5e3e3d Mon Sep 17 00:00:00 2001
From: Pieter Wuille <pieter.wuille@gmail.com>
Date: Sat, 25 Feb 2017 12:16:58 -0800
Subject: [PATCH] Add a FastRandomContext::randrange and use it

---
 configure.ac              |  2 ++
 src/bench/checkqueue.cpp  |  2 +-
 src/crypto/common.h       | 21 +++++++++++++++++++++
 src/net.h                 |  2 +-
 src/random.h              | 12 ++++++++++++
 src/test/crypto_tests.cpp | 23 +++++++++++++++++++++++
 src/test/random_tests.cpp | 15 +++++++++++++++
 7 files changed, 75 insertions(+), 2 deletions(-)

diff --git a/configure.ac b/configure.ac
index 2a9ee018a0..42b8448cfe 100644
--- a/configure.ac
+++ b/configure.ac
@@ -550,6 +550,8 @@ AC_CHECK_DECLS([bswap_16, bswap_32, bswap_64],,,
                  #include <byteswap.h>
                  #endif])
 
+AC_CHECK_DECLS([__builtin_clz, __builtin_clzl, __builtin_clzll])
+
 dnl Check for MSG_NOSIGNAL
 AC_MSG_CHECKING(for MSG_NOSIGNAL)
 AC_COMPILE_IFELSE([AC_LANG_PROGRAM([[#include <sys/socket.h>]],
diff --git a/src/bench/checkqueue.cpp b/src/bench/checkqueue.cpp
index 6fa9fe4fe8..88a2a570f9 100644
--- a/src/bench/checkqueue.cpp
+++ b/src/bench/checkqueue.cpp
@@ -68,7 +68,7 @@ static void CCheckQueueSpeedPrevectorJob(benchmark::State& state)
         PrevectorJob(){
         }
         PrevectorJob(FastRandomContext& insecure_rand){
-            p.resize(insecure_rand.rand32() % (PREVECTOR_SIZE*2));
+            p.resize(insecure_rand.randrange(PREVECTOR_SIZE*2));
         }
         bool operator()()
         {
diff --git a/src/crypto/common.h b/src/crypto/common.h
index 4a9d1150b6..bcca3d30ea 100644
--- a/src/crypto/common.h
+++ b/src/crypto/common.h
@@ -79,4 +79,25 @@ void static inline WriteBE64(unsigned char* ptr, uint64_t x)
     memcpy(ptr, (char*)&v, 8);
 }
 
+/** Return the smallest number n such that (x >> n) == 0 (or 64 if the highest bit in x is set. */
+uint64_t static inline CountBits(uint64_t x)
+{
+#ifdef HAVE_DECL___BUILTIN_CLZL
+    if (sizeof(unsigned long) >= sizeof(uint64_t)) {
+        return x ? 8 * sizeof(unsigned long) - __builtin_clzl(x) : 0;
+    }
+#endif
+#ifdef HAVE_DECL___BUILTIN_CLZLL
+    if (sizeof(unsigned long long) >= sizeof(uint64_t)) {
+        return x ? 8 * sizeof(unsigned long long) - __builtin_clzll(x) : 0;
+    }
+#endif
+    int ret = 0;
+    while (x) {
+        x >>= 1;
+        ++ret;
+    }
+    return ret;
+}
+
 #endif // BITCOIN_CRYPTO_COMMON_H
diff --git a/src/net.h b/src/net.h
index 3bbe386173..c5360f9b58 100644
--- a/src/net.h
+++ b/src/net.h
@@ -760,7 +760,7 @@ public:
         // after addresses were pushed.
         if (_addr.IsValid() && !addrKnown.contains(_addr.GetKey())) {
             if (vAddrToSend.size() >= MAX_ADDR_TO_SEND) {
-                vAddrToSend[insecure_rand.rand32() % vAddrToSend.size()] = _addr;
+                vAddrToSend[insecure_rand.randrange(vAddrToSend.size())] = _addr;
             } else {
                 vAddrToSend.push_back(_addr);
             }
diff --git a/src/random.h b/src/random.h
index 077f58c4d9..9551e1c461 100644
--- a/src/random.h
+++ b/src/random.h
@@ -7,6 +7,7 @@
 #define BITCOIN_RANDOM_H
 
 #include "crypto/chacha20.h"
+#include "crypto/common.h"
 #include "uint256.h"
 
 #include <stdint.h>
@@ -91,6 +92,17 @@ public:
         }
     }
 
+    /** Generate a random integer in the range [0..range). */
+    uint64_t randrange(uint64_t range)
+    {
+        --range;
+        int bits = CountBits(range);
+        while (true) {
+            uint64_t ret = randbits(bits);
+            if (ret <= range) return ret;
+        }
+    }
+
     /** Generate a random 32-bit integer. */
     uint32_t rand32() { return randbits(32); }
 
diff --git a/src/test/crypto_tests.cpp b/src/test/crypto_tests.cpp
index a20182afb8..72e562808a 100644
--- a/src/test/crypto_tests.cpp
+++ b/src/test/crypto_tests.cpp
@@ -10,6 +10,7 @@
 #include "crypto/sha512.h"
 #include "crypto/hmac_sha256.h"
 #include "crypto/hmac_sha512.h"
+#include "random.h"
 #include "utilstrencodings.h"
 #include "test/test_bitcoin.h"
 #include "test/test_random.h"
@@ -484,4 +485,26 @@ BOOST_AUTO_TEST_CASE(chacha20_testvector)
                  "fab78c9");
 }
 
+BOOST_AUTO_TEST_CASE(countbits_tests)
+{
+    FastRandomContext ctx;
+    for (int i = 0; i <= 64; ++i) {
+        if (i == 0) {
+            // Check handling of zero.
+            BOOST_CHECK_EQUAL(CountBits(0), 0);
+        } else if (i < 10) {
+            for (uint64_t j = 1 << (i - 1); (j >> i) == 0; ++j) {
+                // Exhaustively test up to 10 bits
+                BOOST_CHECK_EQUAL(CountBits(j), i);
+            }
+        } else {
+            for (int k = 0; k < 1000; k++) {
+                // Randomly test 1000 samples of each length above 10 bits.
+                uint64_t j = ((uint64_t)1) << (i - 1) | ctx.randbits(i - 1);
+                BOOST_CHECK_EQUAL(CountBits(j), i);
+            }
+        }
+    }
+}
+
 BOOST_AUTO_TEST_SUITE_END()
diff --git a/src/test/random_tests.cpp b/src/test/random_tests.cpp
index 31b993cd38..8596734226 100644
--- a/src/test/random_tests.cpp
+++ b/src/test/random_tests.cpp
@@ -35,4 +35,19 @@ BOOST_AUTO_TEST_CASE(fastrandom_tests)
     BOOST_CHECK(ctx3.rand64() != ctx4.rand64()); // extremely unlikely to be equal
 }
 
+BOOST_AUTO_TEST_CASE(fastrandom_randbits)
+{
+    FastRandomContext ctx1;
+    FastRandomContext ctx2;
+    for (int bits = 0; bits < 63; ++bits) {
+        for (int j = 0; j < 1000; ++j) {
+            uint64_t rangebits = ctx1.randbits(bits);
+            BOOST_CHECK_EQUAL(rangebits >> bits, 0);
+            uint64_t range = ((uint64_t)1) << bits | rangebits;
+            uint64_t rand = ctx2.randrange(range);
+            BOOST_CHECK(rand < range);
+        }
+    }
+}
+
 BOOST_AUTO_TEST_SUITE_END()