τau / math/rand /

rand.tau

source
/Users/niconex/Documents/tau/stdlib/math/rand/rand.tau
1# rand - pseudo random numbers.2#3# New returns a source seeded with a number, so a program that wants the same4# sequence every time can have it. The functions of the module itself use one5# source seeded from the clock at startup:6#7#	rand = import("math/rand")8#	rand.Intn(6) + 19#10# The generator is xorshift64*, which is small and good enough for shuffles,11# sampling and tests. It is not for keys or tokens: Crypto reads those from12# the system.1314syscall = import("syscall")1516mask63 = 0x7fffffffffffffff1718# ushr shifts right without carrying the sign along, which >> does.19ushr = fn(x, n) { (x >> n) & ((1 << (64 - n)) - 1) }2021# mix is splitmix64, used to spread a seed over the whole word: xorshift22# starting from a small number takes a while to look random.23mix = fn(x) {24	x = x + -704602925438635313125	z = x ^ ushr(x, 30)26	z = z * -465889528055300768727	z = z ^ ushr(z, 27)28	z = z * -772359229311070568529	return z ^ ushr(z, 31)30}3132# New returns a source of random numbers seeded with seed. Two sources with33# the same seed give the same sequence.34New = fn(seed) {35	r = new()36	r.state = mix(int(seed, 64))37	if r.state == 0 {38		r.state = 1 # xorshift stays at zero forever39	}4041	# next advances the state and returns the whole 64 bit word, sign and all.42	r.next = fn() {43		s = r.state44		s = s ^ ushr(s, 12)45		s = s ^ (s << 25)46		s = s ^ ushr(s, 27)47		r.state = s48		return s * -704602925438635313149	}5051	# Int returns a number between 0 and 2^63 - 1.52	r.Int = fn() { r.next() & mask63 }5354	# Intn returns a number between 0 and n - 1. n must be positive.55	r.Intn = fn(n) {56		if n <= 0 {57			return error("rand: Intn needs a positive n, got {n}")58		}5960		# The numbers above the last whole multiple of n are thrown away,61		# otherwise the low ones would come up more often than the high ones.62		limit = mask63 - mask63 % n63		for {64			v = r.Int()65			if v < limit {66				return v % n67			}68		}69	}7071	# Float returns a number between 0 (included) and 1 (excluded), with the72	# 53 bits a float can hold.73	r.Float = fn() { float(ushr(r.next() & mask63, 10)) / float(1 << 53) }7475	# Bool returns true about half the time.76	r.Bool = fn() { (r.next() & 1) == 1 }7778	# Bytes returns n random bytes.79	r.Bytes = fn(n) {80		out = []81		for i = 0; i < n; ++i {82			out = append(out, r.Int() & 0xff)83		}84		return bytes(out)85	}8687	# Perm returns the numbers from 0 to n - 1 in random order.88	r.Perm = fn(n) {89		out = []90		for i = 0; i < n; ++i {91			out = append(out, i)92		}93		return r.Shuffle(out)94	}9596	# Shuffle returns the elements of xs in random order, leaving xs alone.97	r.Shuffle = fn(xs) {98		out = []99		for i = 0; i < len(xs); ++i {100			out = append(out, xs[i])101		}102103		# Fisher-Yates, from the end down.104		for i = len(out) - 1; i > 0; --i {105			j = r.Intn(i + 1)106			tmp = out[i]107			out[i] = out[j]108			out[j] = tmp109		}110		return out111	}112113	# Choice returns one element of xs, or null when there are none.114	r.Choice = fn(xs) {115		if len(xs) == 0 {116			return null117		}118		return xs[r.Intn(len(xs))]119	}120121	return r122}123124# The source behind the functions of the module, seeded with the clock. The125# monotonic reading goes in too, so two programs started in the same126# millisecond don't run the same sequence.127global = New(syscall.TimeMillis() * 1000003 + syscall.TimeMono())128129# Seed restarts the sequence of the module from seed, for a program that wants130# the same numbers on every run. The state is written into the source that is131# already there, since a function assigning to global would only be writing to132# a name of its own.133Seed = fn(seed) {134	global.state = New(seed).state135	return null136}137138Int = fn() { global.Int() }139Intn = fn(n) { global.Intn(n) }140Float = fn() { global.Float() }141Bool = fn() { global.Bool() }142Bytes = fn(n) { global.Bytes(n) }143Perm = fn(n) { global.Perm(n) }144Shuffle = fn(xs) { global.Shuffle(xs) }145Choice = fn(xs) { global.Choice(xs) }146147# Crypto returns n bytes from the system generator, the ones to use for keys,148# tokens and passwords. Unlike the rest of the module it can fail, when149# /dev/urandom isn't there.150#151# ponytail: /dev/urandom rather than getrandom(2), so no C is needed. On152# Windows it fails, which is where a sys_getrandom would have to go.153Crypto = fn(n) {154	if n < 0 {155		return error("rand: Crypto needs a length, got {n}")156	}157	if n == 0 {158		return bytes(0)159	}160161	if failed(fd = syscall.Open("/dev/urandom", syscall.O_RDONLY, 0)) {162		return error("rand: no system generator: {fd}")163	}164165	buf = bytes(n)166	got = 0167	for got < n {168		if failed(k = syscall.Read(fd, slice(buf, got, n), n - got)) {169			syscall.Close(fd)170			return k171		}172		if k == 0 {173			syscall.Close(fd)174			return error("rand: /dev/urandom ended early")175		}176		got = got + k177	}178179	syscall.Close(fd)180	return buf181}