code / tomo-random

   1 # Random Number Generator (RNG) implementation based on ChaCha
   2 
   3 use <assert.h>
   4 use ./sysrandom.h
   5 use ./chacha.h
   6 
   7 struct chacha_ctx(j0,j1,j2,j3,j4,j5,j6,j7,j8,j9,j10,j11,j12,j13,j14,j15:Int32; secret)
   8     func from_seed(seed:[Byte]=[] -> chacha_ctx)
   9         ctx := chacha_ctx(0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0)
  10         C_code`
  11             uint8_t seed_bytes[KEYSZ + IVSZ] = {};
  12             if (@seed.length <= (int64_t)sizeof(seed_bytes)) {
  13                 for (int64_t i = 0; i < (int64_t)sizeof(seed_bytes); i++)
  14                     seed_bytes[i] = i < @seed.length ? *(uint8_t*)(@seed.data + i*@seed.stride) : 0;
  15             } else {
  16                 // If the seed is too big, we use as many bytes from the start of the seed as we can
  17                 // fit and then hash the rest of the seed and use the hash bytes at the end:
  18                 for (int64_t i = 0; i < KEYSZ + IVSZ - (int64_t)sizeof(uint64_t); i++)
  19                     seed_bytes[i] = i < @seed.length ? *(uint8_t*)(@seed.data + i*@seed.stride) : 0;
  20                 List_t rest = List$from(@seed, I((int64_t)sizeof(seed_bytes)));
  21                 uint64_t hash = generic_hash(&rest, List$info(&Byte$info));
  22                 memcpy(seed_bytes + KEYSZ + IVSZ - sizeof(uint64_t), &hash, sizeof(hash));
  23             }
  24             chacha_keysetup((void*)&@ctx, seed_bytes);
  25             chacha_ivsetup((void*)&@ctx, seed_bytes + KEYSZ);
  26         `
  27         return ctx
  28 
  29 random := RandomNumberGenerator.new()
  30 
  31 func _os_random_bytes(count:Int64 -> [Byte])
  32     return C_code:[Byte]`
  33         uint8_t *random_bytes = GC_MALLOC_ATOMIC(@count);
  34         assert(getrandom(random_bytes, (size_t)@count, 0) == (size_t)@count);
  35         (List_t){.length=@count, .data=random_bytes, .stride=1, .atomic=1}
  36     `
  37 struct RandomNumberGenerator(_chacha:chacha_ctx, _random_bytes:[Byte]=[]; secret)
  38     func new(seed:[Byte]?=none -> RandomNumberGenerator)
  39         ctx := chacha_ctx.from_seed(seed or _os_random_bytes(40))
  40         return RandomNumberGenerator(ctx, [])
  41 
  42     func _rekey(rng:&RandomNumberGenerator)
  43         new_bytes : [Byte]
  44         C_code `
  45             Byte_t new_keystream[KEYSZ + IVSZ] = {};
  46             // Fill the buffer with the keystream
  47             chacha_encrypt_bytes((void*)&@rng->_chacha, new_keystream, new_keystream, sizeof(new_keystream));
  48             // Immediately reinitialize for backtracking resistance
  49             chacha_keysetup((void*)&@rng->_chacha, new_keystream);
  50             chacha_ivsetup((void*)&@rng->_chacha, new_keystream + KEYSZ);
  51             @new_bytes = (List_t){.data=GC_MALLOC_ATOMIC(1024), .length=1024, .stride=1, .atomic=1};
  52             memset(@new_bytes.data, 0, @new_bytes.length);
  53             chacha_encrypt_bytes((void*)&@rng->_chacha, @new_bytes.data, @new_bytes.data, @new_bytes.length);
  54         `
  55         rng._random_bytes = new_bytes
  56 
  57     func _fill_bytes(rng:&RandomNumberGenerator, dest:&Memory, needed:Int64)
  58         C_code `
  59             while (@needed > 0) {
  60                 if (@rng->_random_bytes.length == 0)
  61                     @(rng._rekey());
  62 
  63                 assert(@rng->_random_bytes.stride == 1);
  64                 if (@rng->_random_bytes.data_refcount > 0) {
  65                     List$compact(&@rng->_random_bytes, sizeof(Byte_t));
  66                 }
  67 
  68                 int64_t batch_size = MIN(@needed, @rng->_random_bytes.length);
  69                 uint8_t *batch_src = @rng->_random_bytes.data;
  70                 memcpy(@dest, batch_src, batch_size);
  71                 memset(batch_src, 0, batch_size);
  72                 @rng->_random_bytes.data += batch_size;
  73                 @rng->_random_bytes.length -= batch_size;
  74                 @dest += batch_size;
  75                 @needed -= batch_size;
  76             }
  77         `
  78 
  79     func bytes(rng:&RandomNumberGenerator, count:Int -> [Byte])
  80         count64 := Int64(count)
  81         buf := C_code:@Memory`GC_MALLOC_ATOMIC(@count64)`
  82         rng._fill_bytes(buf, count64)
  83         return C_code:[Byte]`(List_t){.data=@buf, .stride=1, .atomic=1, .length=@count64}`
  84 
  85     func byte(rng:&RandomNumberGenerator, min:Byte=0, max:Byte=Byte.max -> Byte)
  86         fail("Random minimum value $min is larger than the maximum value $max") if min > max
  87         return min if min == max
  88         random_byte : &Byte
  89         rng._fill_bytes(random_byte, 1)
  90         if min == 0 and max == Byte.max
  91             return random_byte[]
  92 
  93         return C_code:Byte`
  94             Byte_t range = (Byte_t)@max - (Byte_t)@min + 1;
  95             Byte_t min_r = -range % range;
  96             for (;;) {
  97                 @(rng._fill_bytes(random_byte, 1));
  98                 if (*@random_byte >= min_r) break;
  99             }
 100             @min + (*@random_byte % range)
 101         `
 102 
 103     func bool(rng:&RandomNumberGenerator, probability=0.5 -> Bool)
 104         if probability == 0.5
 105             return rng.byte() < 0x80
 106         else
 107             return rng.num(0., 1.) < 0.5
 108 
 109     func int64(rng:&RandomNumberGenerator, min=Int64.min, max=Int64.max -> Int64)
 110         fail("Random minimum value $min is larger than the maximum value $max") if min > max
 111         return min if min == max
 112         random_int64 : &Int64
 113         rng._fill_bytes(random_int64, 8)
 114         if min == Int64.min and max == Int64.max
 115             return random_int64
 116 
 117         return C_code:Int64`
 118             uint64_t range = (uint64_t)@max - (uint64_t)@min + 1;
 119             uint64_t min_r = -range % range;
 120             uint64_t r;
 121             @random_int64 = (int64_t*)&r;
 122             for (;;) {
 123                 @(rng._fill_bytes(random_int64, 8));
 124                 if (r >= min_r) break;
 125             }
 126             (int64_t)((uint64_t)@min + (r % range))
 127         `
 128 
 129     func int32(rng:&RandomNumberGenerator, min=Int32.min, max=Int32.max -> Int32)
 130         fail("Random minimum value $min is larger than the maximum value $max") if min > max
 131         return min if min == max
 132         random_int32 : &Int32
 133         rng._fill_bytes(random_int32, 8)
 134         if min == Int32.min and max == Int32.max
 135             return random_int32
 136 
 137         return C_code:Int32`
 138             uint32_t range = (uint32_t)@max - (uint32_t)@min + 1;
 139             uint32_t min_r = -range % range;
 140             uint32_t r;
 141             @random_int32 = (int32_t*)&r;
 142             for (;;) {
 143                 @(rng._fill_bytes(random_int32, 4));
 144                 if (r >= min_r) break;
 145             }
 146             (int32_t)((uint32_t)@min + (r % range))
 147         `
 148 
 149     func int16(rng:&RandomNumberGenerator, min=Int16.min, max=Int16.max -> Int16)
 150         fail("Random minimum value $min is larger than the maximum value $max") if min > max
 151         return min if min == max
 152         random_int16 : &Int16
 153         rng._fill_bytes(random_int16, 8)
 154         if min == Int16.min and max == Int16.max
 155             return random_int16
 156 
 157         return C_code:Int16`
 158             uint16_t range = (uint16_t)@max - (uint16_t)@min + 1;
 159             uint16_t min_r = -range % range;
 160             uint16_t r;
 161             @random_int16 = (int16_t*)&r;
 162             for (;;) {
 163                 @(rng._fill_bytes(random_int16, 2));
 164                 if (r >= min_r) break;
 165             }
 166             (int16_t)((uint16_t)@min + (r % range))
 167         `
 168 
 169     func int8(rng:&RandomNumberGenerator, min=Int8.min, max=Int8.max -> Int8)
 170         fail("Random minimum value $min is larger than the maximum value $max") if min > max
 171         return min if min == max
 172         random_int8 : &Int8
 173         rng._fill_bytes(random_int8, 1)
 174         if min == Int8.min and max == Int8.max
 175             return random_int8[]
 176 
 177         return C_code:Int8`
 178             uint8_t range = (uint8_t)@max - (uint8_t)@min + 1;
 179             uint8_t min_r = -range % range;
 180             uint8_t r;
 181             @random_int8 = (int8_t*)&r;
 182             for (;;) {
 183                 @(rng._fill_bytes(random_int8, 1));
 184                 if (r >= min_r) break;
 185             }
 186             (int8_t)((uint8_t)@min + (r % range))
 187         `
 188 
 189     func num(rng:&RandomNumberGenerator, min=0., max=1. -> Num)
 190         return C_code:Num`
 191             if (@min > @max) fail("Random minimum value (", @min, ") is larger than the maximum value (", @max, ")");
 192             if (@min == @max) return @min;
 193 
 194             union {
 195                 Num_t num;
 196                 uint64_t bits;
 197             } r = {.bits=0}, one = {.num=1.0};
 198             @(rng._fill_bytes(C_code:&Num`&r.num`, 8));
 199 
 200             // Set r.num to 1.<random-bits>
 201             r.bits &= ~(0xFFFULL << 52);
 202             r.bits |= (one.bits & (0xFFFULL << 52));
 203 
 204             r.num -= 1.0;
 205 
 206             (@min == 0.0 && @max == 1.0) ? r.num : ((1.0-r.num)*@min + r.num*@max)
 207         `
 208 
 209     func num32(rng:&RandomNumberGenerator, min=Num32(0.), max=Num32(1.) -> Num32)
 210         return Num32(rng.num(Num(min), Num(max)))
 211 
 212     func int(rng:&RandomNumberGenerator, min:Int, max:Int -> Int)
 213         return C_code:Int`
 214             if (likely(((@min.small & @max.small) & 1) != 0)) {
 215                 int32_t r = @(rng.int32(Int32(min), Int32(max)));
 216                 return I_small(r);
 217             }
 218 
 219             int32_t cmp = @(min <> max);
 220             if (cmp > 0)
 221                 fail("Random minimum value (", @min, ") is larger than the maximum value (", @max, ")");
 222             if (cmp == 0) return @min;
 223 
 224             mpz_t range_size;
 225             mpz_init_set_int(range_size, @max);
 226             if (@min.small & 1) {
 227                 mpz_t min_mpz;
 228                 mpz_init_set_si(min_mpz, @min.small >> 2);
 229                 mpz_sub(range_size, range_size, min_mpz);
 230             } else {
 231                 mpz_sub(range_size, range_size, @min.big);
 232             }
 233 
 234             gmp_randstate_t gmp_rng;
 235             gmp_randinit_default(gmp_rng);
 236             int64_t seed = @(rng.int64());
 237             gmp_randseed_ui(gmp_rng, (unsigned long)seed);
 238 
 239             mpz_t r;
 240             mpz_init(r);
 241             mpz_urandomm(r, gmp_rng, range_size);
 242 
 243             gmp_randclear(gmp_rng);
 244             Int$plus(@min, Int$from_mpz(r))
 245         `
 246 
 247 
 248 func main()
 249     >> bytes := "asdf".utf8()
 250     rng := RandomNumberGenerator.new(bytes)
 251     >> rng.num()
 252     >> rng.num()
 253     >> rng.num()
 254     >> rng.num(0, 100)
 255     >> rng.byte()
 256     >> rng.bytes(20)
 257 
 258     cached := rng
 259 
 260     >> assert cached == rng
 261     >> assert rng.int64(1, 1000000) == cached.int64(1, 1000000)
 262 
 263     seed1 := [Byte(i) for i in 255]
 264     rng1 := RandomNumberGenerator.new(seed1)
 265 
 266     # Similar at the start, but different at the end
 267     seed2 := [(if i == 255 then Byte(0) else Byte(i)) for i in 255]
 268     rng2 := RandomNumberGenerator.new(seed2)
 269 
 270     assert rng1 != rng2
 271 
 272 
 273     >> random.bool()
 274     >> random.bool(0.7)
 275     >> random.byte(1, 10)
 276     >> random.int8(0xB, 0xF)
 277     >> random.int16(1, 10)
 278     >> random.int32(1, 10)
 279     >> random.int64(1, 10)
 280     >> random.num(1, 10)
 281     >> random.num32(1, 10)