about summary refs log tree commit diff
path: root/src/Xanthous/Random.hs
diff options
context:
space:
mode:
Diffstat (limited to 'src/Xanthous/Random.hs')
-rw-r--r--src/Xanthous/Random.hs40
1 files changed, 40 insertions, 0 deletions
diff --git a/src/Xanthous/Random.hs b/src/Xanthous/Random.hs
new file mode 100644
index 000000000000..a3a1124f2780
--- /dev/null
+++ b/src/Xanthous/Random.hs
@@ -0,0 +1,40 @@
+{-# LANGUAGE TupleSections #-}
+{-# LANGUAGE UndecidableInstances #-}
+
+module Xanthous.Random
+  ( Choose(..)
+  , ChooseElement(..)
+  ) where
+
+import Xanthous.Prelude
+import Data.List.NonEmpty (NonEmpty)
+import System.Random
+import Control.Monad.Random.Class (MonadRandom(getRandomR))
+
+class Choose a where
+  type RandomResult a
+  choose :: MonadRandom m => a -> m (RandomResult a)
+
+newtype ChooseElement a = ChooseElement a
+
+instance MonoFoldable a => Choose (ChooseElement a) where
+  type RandomResult (ChooseElement a) = Maybe (Element a)
+  choose (ChooseElement xs) = do
+    chosenIdx <- getRandomR (0, olength xs - 1)
+    let pick _ (Just x) = Just x
+        pick (x, i) Nothing
+          | i == chosenIdx = Just x
+          | otherwise = Nothing
+    pure $ ofoldr pick Nothing $ zip (toList xs) [0..]
+
+instance MonoFoldable a => Choose (NonNull a) where
+  type RandomResult (NonNull a) = Element a
+  choose
+    = fmap (fromMaybe (error "unreachable")) -- why not lol
+    . choose
+    . ChooseElement
+    . toNullable
+
+instance Choose (NonEmpty a) where
+  type RandomResult (NonEmpty a) = a
+  choose = choose . fromNonEmpty @[_]