@@ -1,12 +1,13 @@
11module Control.ST.Random
22
33import Control.ST
4+import Data.Vect
5+
6+%access public export
47
5−public export
68Random : Type
79Random = State Integer
810
9−export
1011getRandom : (rnd : Var) -> ST m Integer [rnd ::: Random]
1112getRandom rnd =
1213 do
@@ -14,3 +15,28 @@ getRandom rnd =
1415 let st' = assert_total ((1664525 * st + 1013904223) `prim__sremBigInt` (pow 2 32))
1516 write rnd st'
1617 pure st'
18+
19+||| Generates a random Integer in a given range
20+rndInt : (rnd : Var) -> Integer -> Integer -> ST m Integer [rnd ::: Random]
21+rndInt rnd lower upper = do v <- getRandom rnd
22+ pure $ assert_total ((v `prim__sremBigInt` (upper - lower)) + lower)
23+
24+||| Generate a random number in `Fin (S k)`
25+|||
26+||| Note that rndFin k takes values 0, 1, ..., k.
27+rndFin : (rnd : Var) -> (k : Nat) -> ST m (Fin (S k)) [rnd ::: Random]
28+rndFin rnd k = do let v = assert_total $ !(getRandom rnd) `prim__sremBigInt` (cast (S k))
29+ pure (toFin v)
30+ where toFin : Integer -> Fin (S k)
31+ toFin x = case integerToFin x (S k) of
32+ Just v => v
33+ Nothing => toFin (assert_smaller x (x - cast (S k)))
34+
35+||| Select a random element from a vector
36+rndSelect' : (rnd : Var) -> Vect (S k) a -> ST m a [rnd ::: Random]
37+rndSelect' {k} rnd xs = pure (Vect.index !(rndFin rnd k) xs)
38+
39+||| Select a random element from a non-empty list
40+rndSelect : (rnd : Var) -> (xs : List a) -> {auto ok : NonEmpty xs} -> ST m a [rnd ::: Random]
41+rndSelect rnd (x::xs) {ok=NonEmpty} = pure ( !(rndSelect' rnd (x::(fromList xs))))
42+rndSelect rnd [] {ok=NonEmpty} impossible