From 4c44da0ab76cdcbda611b2ffcebb02a5b8a72087 Mon Sep 17 00:00:00 2001
From: Ivan Smirnov <ifsmirnov@yandex.ru>
Date: Mon, 20 Mar 2017 13:12:27 +0300
Subject: [PATCH] added wnext

---
 build.py         |  35 +++++++----
 jngen.h          | 154 +++++++++++++++++++++++++++++++++++++++--------
 random.h         |  92 +++++++++++++++++++++-------
 sequence_ops.h   |   8 ++-
 tests/random.cpp |  42 +++++++++++++
 5 files changed, 272 insertions(+), 59 deletions(-)

diff --git a/build.py b/build.py
index e6d6742..9ffcac1 100755
--- a/build.py
+++ b/build.py
@@ -6,18 +6,18 @@
 HEADER_REGEX = re.compile('#include "(.*)"')
 
 
-def is_header(line):
+def extract_header(line):
     res = HEADER_REGEX.match(line)
     if res:
         return res.groups()[0]
 
 
-def direct_deps(filename):
+def get_direct_deps(filename):
     res = set()
     with open(filename) as fin:
         for line in fin.readlines():
-            t = is_header(line)
-            if t:
+            t = extract_header(line)
+            if t and not t.endswith("_inl.h"):
                 res.add(t)
     return res
 
@@ -25,15 +25,19 @@ def direct_deps(filename):
 try:
     unused_files = set(map(str.strip, open(".unused_files").readlines()))
 except IOError:
-    unused_files = {}
+    unused_files = set()
+
+
+unused_files.add("jngen.h")
 
 
 deps = {}
 
 
 for filename in os.listdir('.'):
-    if filename.endswith('.h') and filename != 'jngen.h' and filename not in unused_files:
-        deps[filename] = direct_deps(filename)
+    if filename.endswith('.h') and not filename.endswith("_inl.h") and\
+            filename not in unused_files:
+        deps[filename] = get_direct_deps(filename)
 
 order = []
 
@@ -46,9 +50,18 @@ def direct_deps(filename):
                 deps[other].discard(item)
             break
 
+
+def write_file(filename, stream):
+    with open(filename) as fin:
+        for line in fin.readlines():
+            include_or_not = HEADER_REGEX.match(line)
+            if include_or_not:
+                if include_or_not.groups()[0].endswith("_inl.h"):
+                    write_file(include_or_not.groups()[0], stream)
+            elif '#pragma once' not in line:
+                stream.write(line)
+
+
 with open("jngen.h", "w") as fout:
     for filename in order:
-        with open(filename) as fin:
-            for line in fin.readlines():
-                if '#pragma once' not in line and not HEADER_REGEX.match(line):
-                    fout.write(line)
+        write_file(filename, fout)
diff --git a/jngen.h b/jngen.h
index 573f4e5..e51fdaf 100644
--- a/jngen.h
+++ b/jngen.h
@@ -415,6 +415,10 @@ std::string Pattern::next(std::function<int(int)> rnd) const {
 using jngen::Pattern;
 
 
+#include <algorithm>
+#include <cmath>
+#include <cstdlib>
+#include <iterator>
 #include <limits>
 #include <random>
 #include <string>
@@ -508,11 +512,6 @@ class Random {
         return uniformRandom(n, *this, (uint32_t (Random::*)())&Random::next);
     }
 
-    long next(long n) {
-        ensure(n > 0);
-        return uniformRandom(n, *this, &Random::next64);
-    }
-
     long long next(long long n) {
         ensure(n > 0);
         return uniformRandom(n, *this, &Random::next64);
@@ -524,28 +523,24 @@ class Random {
     }
 
     double next(double n) {
+        ensure(n >= 0);
         return nextf() * n;
     }
 
-    int next(int l, int r) {
-        return l + next(r-l+1);
-    }
-
-    long next(long l, long r) {
-        return l + next(r-l+1);
-    }
-
-    long long next(long long l, long long r) {
-        return l + next(r-l+1);
-    }
+    int next(int l, int r);
+    long long next(long long l, long long r);
+    size_t next(size_t l, size_t r);
+    double next(double l, double r);
 
-    size_t next(size_t l, size_t r) {
-        return l + next(r-l+1);
-    }
+    int wnext(int n, int w);
+    long long wnext(long long n, int w);
+    size_t wnext(size_t n, int w);
+    double wnext(double n, int w);
 
-    double next(double l, double r) {
-        return l + next(r-l);
-    }
+    int wnext(int l, int r, int w);
+    long long wnext(long long l, long long r, int w);
+    size_t wnext(size_t l, size_t r, int w);
+    double wnext(double l, double r, int w);
 
     std::string next(const std::string& pattern) {
         return Pattern(pattern).next([this](int n) { return next(n); });
@@ -561,8 +556,57 @@ class Random {
         return tnext<std::pair<int, int>>(args...);
     }
 
+    template<typename Iterator>
+    typename Iterator::value_type choice(Iterator begin, Iterator end) {
+        auto length = std::distance(begin, end);
+        ensure(length > 0, "Cannot select from a range of negative length");
+        size_t index = tnext<size_t>(length);
+        std::advance(begin, index);
+        return *begin;
+    }
+
+    template<typename Container>
+    typename Container::value_type choice(const Container& container) {
+        ensure(!container.empty(), "Cannot select from an empty container");
+        return choice(container.begin(), container.end());
+    }
+
 private:
+    template<typename T>
+    T baseWnext(T n, int w) {
+        static_assert(std::is_arithmetic<T>::value,
+            "Only numeric types allowed for baseWnext<T>(T n, int w)");
+        if (std::abs(w) <= WNEXT_LIMIT) {
+            T result = next(n);
+            while (w > 0) {
+                result = std::max(result, next(n));
+                --w;
+            }
+            while (w < 0) {
+                result = std::min(result, next(n));
+                ++w;
+            }
+            return result;
+        }
+
+        if (w < 0) {
+            if (std::is_integral<T>::value) {
+                return n - 1 - baseWnext(n, -w);
+            } else {
+                return n - baseWnext(n, -w);
+            }
+        }
+
+        T upperLimit =
+            std::is_integral<T>::value ? n-1 : n;
+
+        double val = std::pow(nextf(), 1.0 / (w + 1));
+        T result = val * n;
+        return std::max(T(0), std::min(result, upperLimit));
+    }
+
     std::mt19937 randomEngine_;
+    constexpr static int WNEXT_LIMIT = 8;
 };
 
 Random rnd;
@@ -660,6 +704,64 @@ void registerGen(int argc, char *argv[], int version = 1) {
     rnd.seed(seed);
 }
 
+#define JNGEN_INCLUDE_RANDOM_INL_H
+#ifndef JNGEN_INCLUDE_RANDOM_INL_H
+#error File "random_inl.h" must not be included directly.
+#endif
+
+namespace jngen {
+
+int Random::next(int l, int r) {
+    return l + next(r-l+1);
+}
+
+long long Random::next(long long l, long long r) {
+    return l + next(r-l+1);
+}
+
+size_t Random::next(size_t l, size_t r) {
+    return l + next(r-l+1);
+}
+
+double Random::next(double l, double r) {
+    return l + next(r-l);
+}
+
+int Random::wnext(int n, int w) {
+    return baseWnext(n, w);
+}
+
+long long Random::wnext(long long n, int w) {
+    return baseWnext(n, w);
+}
+
+size_t Random::wnext(size_t n, int w) {
+    return baseWnext(n, w);
+}
+
+double Random::wnext(double n, int w) {
+    return baseWnext(n, w);
+}
+
+int Random::wnext(int l, int r, int w) {
+    return l + wnext(r-l+1, w);
+}
+
+long long Random::wnext(long long l, long long r, int w) {
+    return l + wnext(r-l+1, w);
+}
+
+size_t Random::wnext(size_t l, size_t r, int w) {
+    return l + wnext(r-l+1, w);
+}
+
+double Random::wnext(double l, double r, int w) {
+    return l + wnext(r-l, w);
+}
+
+} // namespace jngen
+#undef JNGEN_INCLUDE_RANDOM_INL_H
+
 #include <iostream>
 #include <type_traits>
 
@@ -1042,8 +1144,12 @@ void shuffle(Iterator begin, Iterator end) {
 
 template<typename Iterator>
 typename Iterator::value_type choice(Iterator begin, Iterator end) {
-    ensure(end > begin, "Cannot select from a range of negative length");
-    return *(begin + rnd.next(end - begin));
+    return rnd.choice(begin, end);
+}
+
+template<typename Container>
+typename Container::value_type choice(const Container& container) {
+    return rnd.choice(container);
 }
 
 } // namespace jngen
diff --git a/random.h b/random.h
index eaca4f2..087b2dd 100644
--- a/random.h
+++ b/random.h
@@ -3,6 +3,10 @@
 #include "common.h"
 #include "pattern.h"
 
+#include <algorithm>
+#include <cmath>
+#include <cstdlib>
+#include <iterator>
 #include <limits>
 #include <random>
 #include <string>
@@ -96,11 +100,6 @@ class Random {
         return uniformRandom(n, *this, (uint32_t (Random::*)())&Random::next);
     }
 
-    long next(long n) {
-        ensure(n > 0);
-        return uniformRandom(n, *this, &Random::next64);
-    }
-
     long long next(long long n) {
         ensure(n > 0);
         return uniformRandom(n, *this, &Random::next64);
@@ -112,28 +111,24 @@ class Random {
     }
 
     double next(double n) {
+        ensure(n >= 0);
         return nextf() * n;
     }
 
-    int next(int l, int r) {
-        return l + next(r-l+1);
-    }
-
-    long next(long l, long r) {
-        return l + next(r-l+1);
-    }
-
-    long long next(long long l, long long r) {
-        return l + next(r-l+1);
-    }
+    int next(int l, int r);
+    long long next(long long l, long long r);
+    size_t next(size_t l, size_t r);
+    double next(double l, double r);
 
-    size_t next(size_t l, size_t r) {
-        return l + next(r-l+1);
-    }
+    int wnext(int n, int w);
+    long long wnext(long long n, int w);
+    size_t wnext(size_t n, int w);
+    double wnext(double n, int w);
 
-    double next(double l, double r) {
-        return l + next(r-l);
-    }
+    int wnext(int l, int r, int w);
+    long long wnext(long long l, long long r, int w);
+    size_t wnext(size_t l, size_t r, int w);
+    double wnext(double l, double r, int w);
 
     std::string next(const std::string& pattern) {
         return Pattern(pattern).next([this](int n) { return next(n); });
@@ -149,8 +144,57 @@ class Random {
         return tnext<std::pair<int, int>>(args...);
     }
 
+    template<typename Iterator>
+    typename Iterator::value_type choice(Iterator begin, Iterator end) {
+        auto length = std::distance(begin, end);
+        ensure(length > 0, "Cannot select from a range of negative length");
+        size_t index = tnext<size_t>(length);
+        std::advance(begin, index);
+        return *begin;
+    }
+
+    template<typename Container>
+    typename Container::value_type choice(const Container& container) {
+        ensure(!container.empty(), "Cannot select from an empty container");
+        return choice(container.begin(), container.end());
+    }
+
 private:
+    template<typename T>
+    T baseWnext(T n, int w) {
+        static_assert(std::is_arithmetic<T>::value,
+            "Only numeric types allowed for baseWnext<T>(T n, int w)");
+        if (std::abs(w) <= WNEXT_LIMIT) {
+            T result = next(n);
+            while (w > 0) {
+                result = std::max(result, next(n));
+                --w;
+            }
+            while (w < 0) {
+                result = std::min(result, next(n));
+                ++w;
+            }
+            return result;
+        }
+
+        if (w < 0) {
+            if (std::is_integral<T>::value) {
+                return n - 1 - baseWnext(n, -w);
+            } else {
+                return n - baseWnext(n, -w);
+            }
+        }
+
+        T upperLimit =
+            std::is_integral<T>::value ? n-1 : n;
+
+        double val = std::pow(nextf(), 1.0 / (w + 1));
+        T result = val * n;
+        return std::max(T(0), std::min(result, upperLimit));
+    }
+
     std::mt19937 randomEngine_;
+    constexpr static int WNEXT_LIMIT = 8;
 };
 
 Random rnd;
@@ -247,3 +291,7 @@ void registerGen(int argc, char *argv[], int version = 1) {
     }
     rnd.seed(seed);
 }
+
+#define JNGEN_INCLUDE_RANDOM_INL_H
+#include "random_inl.h"
+#undef JNGEN_INCLUDE_RANDOM_INL_H
diff --git a/sequence_ops.h b/sequence_ops.h
index fb651a6..ded7491 100644
--- a/sequence_ops.h
+++ b/sequence_ops.h
@@ -20,8 +20,12 @@ void shuffle(Iterator begin, Iterator end) {
 
 template<typename Iterator>
 typename Iterator::value_type choice(Iterator begin, Iterator end) {
-    ensure(end > begin, "Cannot select from a range of negative length");
-    return *(begin + rnd.next(end - begin));
+    return rnd.choice(begin, end);
+}
+
+template<typename Container>
+typename Container::value_type choice(const Container& container) {
+    return rnd.choice(container);
 }
 
 } // namespace jngen
diff --git a/tests/random.cpp b/tests/random.cpp
index 1495fb3..9e52909 100644
--- a/tests/random.cpp
+++ b/tests/random.cpp
@@ -111,3 +111,45 @@ BOOST_AUTO_TEST_CASE(several_engines) {
     BOOST_CHECK(etalon == generate(r2));
     BOOST_CHECK(etalon == generate(r1));
 }
+
+BOOST_AUTO_TEST_CASE(test_choice) {
+    BOOST_CHECK(true);
+
+    std::vector<int> a;
+    for (int i = 0; i < 10; ++i) {
+        a.push_back(i);
+    }
+
+    shuffle(a.begin(), a.end());
+    choice(a.begin(), a.end());
+    choice(a);
+
+    rnd.choice(a.begin(), a.end());
+    rnd.choice(a);
+
+    std::set<int> b(a.begin(), a.end());
+    choice(b.begin(), b.end());
+    choice(b);
+
+    rnd.choice(b.begin(), b.end());
+    rnd.choice(b);
+}
+
+BOOST_AUTO_TEST_CASE(wnext) {
+    rnd.seed(987);
+
+    rnd.wnext(0.1, 1);
+    rnd.wnext(1.0, 2.0, 2);
+
+    rnd.wnext(10, 3);
+    rnd.wnext(10, 20, -3);
+
+    rnd.wnext(100ll, 0);
+    rnd.wnext(100ll, 200ll, -40);
+
+    auto a = rnda.randomf(10000, []() { return rnd.wnext(1, 10, 2); });
+    a.sort().unique();
+    ensure(a.size() == 10u);
+    ensure(a[0] == 1);
+    ensure(a[9] == 10);
+}