REF: Refactor uint generators

Refactor uint generators so that they can be used more generally

Author: bashtage

randomstate/distributions.c

@@ -1401,32 +1401,82 @@ static inline uint64_t gen_mask(uint64_t max)
  * Fills an array with cnt random npy_uint64 between off and off + rng
  * inclusive. The numbers wrap if rng is sufficiently large.
  */
-void random_bounded_uint64_fill(aug_state *state, uint64_t off, uint64_t rng, npy_intp cnt, uint64_t *out)
-{
-    uint64_t val, mask;
-    npy_intp i;
 
-    if (rng == 0) {
-        for (i = 0; i < cnt; i++) {
-            out[i] = off;
-        }
-        return;
-    }
 
-    /* Smallest bit mask >= max */
-    mask = gen_mask(rng);
+inline uint64_t random_bounded_uint64(aug_state *state, uint64_t off, uint64_t rng, uint64_t mask)
+{
+    uint64_t val;
+    if (rng == 0)
+        return off;
 
     if (rng <= 0xffffffffUL) {
-        for (i = 0; i < cnt; i++) {
-            while ((val = (random_uint32(state) & mask)) > rng);
-            out[i] =  off + val;
-        }
+        while ((val = (random_uint32(state) & mask)) > rng);
     }
     else {
-        for (i = 0; i < cnt; i++) {
-            while ((val = (random_uint64(state) & mask)) > rng);
-            out[i] =  off + val;
+        while ((val = (random_uint64(state) & mask)) > rng);
+    }
+    return  off + val;
+}
+
+
+inline uint32_t random_bounded_uint32(aug_state *state, uint32_t off, uint32_t rng, uint32_t mask)
+{
+    uint32_t val;
+    if (rng == 0)
+        return off;
+
+    while ((val = (random_uint32(state) & mask)) > rng);
+    return off + val;
+}
+
+inline uint16_t random_buffered_bounded_uint16(aug_state *state, uint16_t off, uint16_t rng, uint16_t mask, int *bcnt, uint32_t *buf)
+{
+    uint16_t val;
+    if (rng == 0)
+        return off;
+
+    do {
+        if (!(bcnt[0])) {
+            buf[0] = random_uint32(state);
+            bcnt[0] = 1;
+        }
+        else {
+            buf[0] >>= 16;
+            bcnt[0] -= 1;
         }
+        val = (uint16_t)buf[0] & mask;
+    } while (val > rng);
+    return  off + val;
+}
+
+inline uint8_t random_buffered_bounded_uint8(aug_state *state, uint8_t off, uint8_t rng, uint8_t mask, int *bcnt, uint32_t *buf)
+{
+    uint8_t val;
+    if (rng == 0)
+        return off;
+    do {
+        if (!(bcnt[0])) {
+            buf[0] = random_uint32(state);
+            bcnt[0] = 3;
+        }
+        else {
+            buf[0] >>= 8;
+            bcnt[0] -= 1;
+        }
+        val = (uint8_t)buf[0] & mask;
+    } while (val > rng);
+    return  off + val;
+}
+
+void random_bounded_uint64_fill(aug_state *state, uint64_t off, uint64_t rng, npy_intp cnt, uint64_t *out)
+{
+    uint64_t mask;
+    npy_intp i;
+
+    /* Smallest bit mask >= max */
+    mask = gen_mask(rng);
+    for (i = 0; i < cnt; i++) {
+        out[i] = random_bounded_uint64(state, off, rng, mask);
     }
 }
 
@@ -1437,22 +1487,13 @@ void random_bounded_uint64_fill(aug_state *state, uint64_t off, uint64_t rng, np
  */
 void random_bounded_uint32_fill(aug_state *state, uint32_t off, uint32_t rng, npy_intp cnt, uint32_t *out)
 {
-    uint32_t val, mask = rng;
+    uint32_t val, mask;
     npy_intp i;
 
-    if (rng == 0) {
-        for (i = 0; i < cnt; i++) {
-            out[i] = off;
-        }
-        return;
-    }
-
     /* Smallest bit mask >= max */
     mask = (uint32_t)gen_mask(rng);
-
     for (i = 0; i < cnt; i++) {
-        while ((val = (random_uint32(state) & mask)) > rng);
-        out[i] =  off + val;
+        out[i] =  random_bounded_uint32(state, off, rng, mask);
     }
 }
 
@@ -1463,34 +1504,15 @@ void random_bounded_uint32_fill(aug_state *state, uint32_t off, uint32_t rng, np
  */
 void random_bounded_uint16_fill(aug_state *state, uint16_t off, uint16_t rng, npy_intp cnt, uint16_t *out)
 {
-    uint16_t val, mask;
+    uint16_t mask;
     npy_intp i;
     uint32_t buf = 0;
     int bcnt = 0;
 
-    if (rng == 0) {
-        for (i = 0; i < cnt; i++) {
-            out[i] = off;
-        }
-        return;
-    }
-
     /* Smallest bit mask >= max */
     mask = (uint16_t)gen_mask(rng);
-
     for (i = 0; i < cnt; i++) {
-        do {
-            if (!bcnt) {
-                buf = random_uint32(state);
-                bcnt = 1;
-            }
-            else {
-                buf >>= 16;
-                bcnt--;
-            }
-            val = (uint16_t)buf & mask;
-        } while (val > rng);
-        out[i] =  off + val;
+        out[i] = random_buffered_bounded_uint16(state, off, rng, mask, &bcnt, &buf);
     }
 }
 
@@ -1500,35 +1522,17 @@ void random_bounded_uint16_fill(aug_state *state, uint16_t off, uint16_t rng, np
  */
 void random_bounded_uint8_fill(aug_state *state, uint8_t off, uint8_t rng, npy_intp cnt, uint8_t *out)
 {
-    uint8_t val, mask = rng;
+    uint8_t mask;
     npy_intp i;
     uint32_t buf = 0;
     int bcnt = 0;
 
-    if (rng == 0) {
-        for (i = 0; i < cnt; i++) {
-            out[i] = off;
-        }
-        return;
-    }
-
     /* Smallest bit mask >= max */
     mask = (uint8_t)gen_mask(rng);
-
     for (i = 0; i < cnt; i++) {
-        do {
-            if (!bcnt) {
-                buf = random_uint32(state);
-                bcnt = 3;
-            }
-            else {
-                buf >>= 8;
-                bcnt--;
-            }
-            val = (uint8_t)buf & mask;
-        } while (val > rng);
-        out[i] =  off + val;
+        out[i] = random_buffered_bounded_uint8(state, off, rng, mask, &bcnt, &buf);
     }
+
 }