mirror of
https://github.com/logos-storage/logos-storage-nim.git
synced 2026-01-09 17:03:09 +00:00
90 lines
2.4 KiB
Nim
90 lines
2.4 KiB
Nim
## Nim-Codex
|
|
## Copyright (c) 2021 Status Research & Development GmbH
|
|
## Licensed under either of
|
|
## * Apache License, version 2.0, ([LICENSE-APACHE](LICENSE-APACHE))
|
|
## * MIT license ([LICENSE-MIT](LICENSE-MIT))
|
|
## at your option.
|
|
## This file may not be copied, modified, or distributed except according to
|
|
## those terms.
|
|
|
|
import std/sugar
|
|
import pkg/libp2p/crypto/crypto
|
|
import pkg/bearssl/rand
|
|
|
|
type
|
|
RngSampleError = object of CatchableError
|
|
Rng* = ref HmacDrbgContext
|
|
|
|
var rng {.threadvar.}: Rng
|
|
|
|
proc instance*(t: type Rng): Rng =
|
|
if rng.isNil:
|
|
rng = newRng()
|
|
rng
|
|
|
|
# Random helpers: similar as in stdlib, but with HmacDrbgContext rng
|
|
# TODO: Move these somewhere else?
|
|
const randMax = 18_446_744_073_709_551_615'u64
|
|
|
|
proc rand*(rng: Rng, max: Natural): int =
|
|
if max == 0:
|
|
return 0
|
|
|
|
while true:
|
|
let x = rng[].generate(uint64)
|
|
if x < randMax - (randMax mod (uint64(max) + 1'u64)): # against modulo bias
|
|
return int(x mod (uint64(max) + 1'u64))
|
|
|
|
proc sampleNoReplacement[T](a: seq[T], n: int): seq[T] =
|
|
if n > a.len:
|
|
raise newException(
|
|
RngSampleError,
|
|
"Cannot sample " & $n & " elements from a set of size " & $a.len &
|
|
" without replacement.",
|
|
)
|
|
|
|
if n == a.len:
|
|
return a
|
|
|
|
var x = a
|
|
collect:
|
|
for i in 0 ..< n:
|
|
swap(x[i], x[i + rng.rand(x.len - i - 1)])
|
|
x[i]
|
|
|
|
proc sampleWithReplacement[T](a: seq[T], n: int): seq[T] =
|
|
collect:
|
|
for i in 0 ..< n:
|
|
a[rng.rand(a.high)]
|
|
|
|
proc sample*[T](rng: Rng, a: openArray[T]): T =
|
|
result = a[rng.rand(a.high)]
|
|
|
|
proc sample*[T](rng: Rng, a: seq[T], n: int, replace: bool = false): seq[T] =
|
|
## Sample `n` elements from a set `a` with or without replacement. In case of
|
|
## sampling with replacement, `n` must not be greater than the size of `a`.
|
|
## In case of sampling without replacement, `n` must not be greater than the
|
|
## size of `a` minus 1.
|
|
if replace:
|
|
sampleWithReplacement(a, n)
|
|
else:
|
|
sampleNoReplacement(a, n)
|
|
|
|
proc sample*[T](
|
|
rng: Rng, sample, exclude: openArray[T]
|
|
): T {.raises: [Defect, RngSampleError].} =
|
|
if sample == exclude:
|
|
raise newException(RngSampleError, "Sample and exclude arrays are the same!")
|
|
|
|
while true:
|
|
result = rng.sample(sample)
|
|
if exclude.find(result) != -1:
|
|
continue
|
|
|
|
break
|
|
|
|
proc shuffle*[T](rng: Rng, a: var openArray[T]) =
|
|
for i in countdown(a.high, 1):
|
|
let j = rng.rand(i)
|
|
swap(a[i], a[j])
|