Small optimization in 'project' from math.random

When computing the Mersenne number, instead of spreading 1's a fixed
number of times (with shifts of 1, 2, 4, 8, 16, and 32), spread only
until the number becomes a Mersenne number.

1 files changed, 9 insertions, 21 deletions

diff --git a/lmathlib.c b/lmathlib.c
index a0981772..bd34c888 100644
--- a/lmathlib.c
+++ b/lmathlib.c
@@ -533,7 +533,7 @@ typedef struct {
 ** Project the random integer 'ran' into the interval [0, n].
 ** Because 'ran' has 2^B possible values, the projection can only be
 ** uniform when the size of the interval is a power of 2 (exact
-** division). Otherwise, to get a uniform projection into [0, n], we
+** division). So, to get a uniform projection into [0, n], we
 ** first compute 'lim', the smallest Mersenne number not smaller than
 ** 'n'. We then project 'ran' into the interval [0, lim].  If the result
 ** is inside [0, n], we are done. Otherwise, we try with another 'ran',
@@ -541,26 +541,14 @@ typedef struct {
 */
 static lua_Unsigned project (lua_Unsigned ran, lua_Unsigned n,
                              RanState *state) {
-  if ((n & (n + 1)) == 0)  /* is 'n + 1' a power of 2? */
-    return ran & n;  /* no bias */
-  else {
-    lua_Unsigned lim = n;
-    /* compute the smallest (2^b - 1) not smaller than 'n' */
-    lim |= (lim >> 1);
-    lim |= (lim >> 2);
-    lim |= (lim >> 4);
-    lim |= (lim >> 8);
-    lim |= (lim >> 16);
-#if (LUA_MAXUNSIGNED >> 31) >= 3
-    lim |= (lim >> 32);  /* integer type has more than 32 bits */
-#endif
-    lua_assert((lim & (lim + 1)) == 0  /* 'lim + 1' is a power of 2, */
-      && lim >= n  /* not smaller than 'n', */
-      && (lim >> 1) < n);  /* and it is the smallest one */
-    while ((ran &= lim) > n)  /* project 'ran' into [0..lim] */
-      ran = I2UInt(nextrand(state->s));  /* not inside [0..n]? try again */
-    return ran;
-  }
+  lua_Unsigned lim = n;  /* to compute the Mersenne number */
+  int sh;  /* how much to spread bits to the right in 'lim' */
+  /* spread '1' bits in 'lim' until it becomes a Mersenne number */
+  for (sh = 1; (lim & (lim + 1)) != 0; sh *= 2)
+    lim |= (lim >> sh);  /* spread '1's to the right */
+  while ((ran &= lim) > n)  /* project 'ran' into [0..lim] and test */
+    ran = I2UInt(nextrand(state->s));  /* not inside [0..n]? try again */
+  return ran;
 }