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;
rand32(void)7 static uint32_t rand32(void)
8 {
9 seed = 6364136223846793005ULL*seed + 1;
10 return seed >> 32;
11 }
rand64(void)12 static uint64_t rand64(void)
13 {
14 uint64_t u = rand32();
15 return u<<32 | rand32();
16 }
frand()17 static double frand()
18 {
19 return rand64() * 0x1p-64;
20 }
frandf()21 static float frandf()
22 {
23 return rand32() * 0x1p-32f;
24 }
frandl()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
t_randseed(uint64_t s)34 void t_randseed(uint64_t s)
35 {
36 seed = s;
37 }
38
39 /* uniform random in [0,n), n > 0 must hold */
t_randn(uint64_t n)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 */
t_randint(uint64_t a,uint64_t b)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 */
shuffle2(uint64_t * p,uint64_t * q,size_t np,size_t nq)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 */
t_shuffle(uint64_t * p,size_t n)80 void t_shuffle(uint64_t *p, size_t n)
81 {
82 shuffle2(p,0,n,0);
83 }
84
t_randrange(uint64_t * p,size_t n)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 */
insert(uint64_t * tab,size_t len,uint64_t v)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 */
t_choose(uint64_t n,size_t k,uint64_t * p)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
170