src/common/rand.c

   1 #include <float.h>
   2 #include <stdint.h>
   3 #include <stdlib.h>
   4
   5 // TODO: use large period prng
   6 static uint64_t seed = -1;
   7 static uint32_t rand32(void)
   8 {
   9         seed = 6364136223846793005ULL*seed + 1;
  10         return seed >> 32;
  11 }
  12 static uint64_t rand64(void)
  13 {
  14         uint64_t u = rand32();
  15         return u<<32 | rand32();
  16 }
  17 static double frand()
  18 {
  19         return rand64() * 0x1p-64;
  20 }
  21 static float frandf()
  22 {
  23         return rand32() * 0x1p-32f;
  24 }
  25 static long double frandl()
  26 {
  27         return rand64() * 0x1p-64L
  28 #if LDBL_MANT_DIG > 64
  29 + rand64() * 0x1p-128L
  30 #endif
  31 ;
  32 }
  33
  34 void t_randseed(uint64_t s)
  35 {
  36         seed = s;
  37 }
  38
  39 /* uniform random in [0,n), n > 0 must hold */
  40 uint64_t t_randn(uint64_t n)
  41 {
  42         uint64_t r, m;
  43
  44         /* m is the largest multiple of n */
  45         m = -1;
  46         m -= m%n;
  47         while ((r = rand64()) >= m);
  48         return r%n;
  49 }
  50
  51 /* uniform on [a,b], a <= b must hold */
  52 uint64_t t_randint(uint64_t a, uint64_t b)
  53 {
  54         uint64_t n = b - a + 1;
  55         if (n)
  56                 return a + t_randn(n);
  57         return rand64();
  58 }
  59
  60 /* shuffle the elements of p and q until the elements in p are well shuffled */
  61 static void shuffle2(uint64_t *p, uint64_t *q, size_t np, size_t nq)
  62 {
  63         size_t r;
  64         uint64_t t;
  65
  66         while (np) {
  67                 r = t_randn(nq+np--);
  68                 t = p[np];
  69                 if (r < nq) {
  70                         p[np] = q[r];
  71                         q[r] = t;
  72                 } else {
  73                         p[np] = p[r-nq];
  74                         p[r-nq] = t;
  75                 }
  76         }
  77 }
  78
  79 /* shuffle the elements of p */
  80 void t_shuffle(uint64_t *p, size_t n)
  81 {
  82         shuffle2(p,0,n,0);
  83 }
  84
  85 void t_randrange(uint64_t *p, size_t n)
  86 {
  87         size_t i;
  88         for (i = 0; i < n; i++)
  89                 p[i] = i;
  90         t_shuffle(p, n);
  91 }
  92
  93 /* hash table insert, 0 means empty, v > 0 must hold, len is power-of-2 */
  94 static int insert(uint64_t *tab, size_t len, uint64_t v)
  95 {
  96         size_t i = v & (len-1);
  97         size_t j = 1;
  98
  99         while (tab[i]) {
 100                 if (tab[i] == v)
 101                         return -1;
 102                 i += j++;
 103                 i &= len-1;
 104         }
 105         tab[i] = v;
 106         return 0;
 107 }
 108
 109 /* choose k unique numbers from [0,n), k <= n */
 110 int t_choose(uint64_t n, size_t k, uint64_t *p)
 111 {
 112         uint64_t *tab;
 113         size_t i, j, len;
 114
 115         if (n < k)
 116                 return -1;
 117
 118         if (n < 16) {
 119                 /* no alloc */
 120                 while (k)
 121                         if (t_randn(n--) < k)
 122                                 p[--k] = n;
 123                 return 0;
 124         }
 125
 126         if (k < 8) {
 127                 /* no alloc, n > 15 > 2*k */
 128                 for (i = 0; i < k;) {
 129                         p[i] = t_randn(n);
 130                         for (j = 0; p[j] != p[i]; j++);
 131                         if (j == i)
 132                                 i++;
 133                 }
 134                 return 0;
 135         }
 136
 137         // TODO: if k < n/k use k*log(k) solution without alloc
 138
 139         if (n < 5*k && (n-k)*sizeof *tab < (size_t)-1) {
 140                 /* allocation is n-k < 4*k */
 141                 tab = malloc((n-k) * sizeof *tab);
 142                 if (!tab)
 143                         return -1;
 144                 for (i = 0; i < k; i++)
 145                         p[i] = i;
 146                 for (; i < n; i++)
 147                         tab[i-k] = i;
 148                 if (k < n-k)
 149                         shuffle2(p, tab, k, n-k);
 150                 else
 151                         shuffle2(tab, p, n-k, k);
 152                 free(tab);
 153                 return 0;
 154         }
 155
 156         /* allocation is 2*k <= len < 4*k */
 157         for (len = 16; len < 2*k; len *= 2);
 158         tab = calloc(len, sizeof *tab);
 159         if (!tab)
 160                 return -1;
 161         for (i = 0; i < k; i++)
 162                 while (insert(tab, len, t_randn(n)+1));
 163         for (i = 0; i < len; i++)
 164                 if (tab[i])
 165                         *p++ = tab[i]-1;
 166         free(tab);
 167         return 0;
 168 }
 169