about summary refs log tree commit diff
path: root/users/glittershark/xanthous/src/Xanthous/Random.hs
diff options
context:
space:
mode:
Diffstat (limited to 'users/glittershark/xanthous/src/Xanthous/Random.hs')
-rw-r--r--users/glittershark/xanthous/src/Xanthous/Random.hs16
1 files changed, 16 insertions, 0 deletions
diff --git a/users/glittershark/xanthous/src/Xanthous/Random.hs b/users/glittershark/xanthous/src/Xanthous/Random.hs
index 41c80ab73c4c..6d34109df7f8 100644
--- a/users/glittershark/xanthous/src/Xanthous/Random.hs
+++ b/users/glittershark/xanthous/src/Xanthous/Random.hs
@@ -10,6 +10,7 @@ module Xanthous.Random
   , weightedBy
   , subRand
   , chance
+  , chooseSubset
   ) where
 --------------------------------------------------------------------------------
 import Xanthous.Prelude
@@ -17,6 +18,7 @@ import Xanthous.Prelude
 import           Data.List.NonEmpty (NonEmpty(..))
 import           Control.Monad.Random.Class (MonadRandom(getRandomR, getRandom))
 import           Control.Monad.Random (Rand, evalRand, mkStdGen, StdGen)
+import           Data.Functor.Compose
 import           Data.Random.Shuffle.Weighted
 import           Data.Random.Distribution
 import           Data.Random.Distribution.Uniform
@@ -66,10 +68,16 @@ instance Choose (a, a) where
   choose (x, y) = choose (x :| [y])
 
 newtype Weighted w t a = Weighted (t (w, a))
+  deriving (Functor, Foldable) via (t `Compose` (,) w)
+
+instance Traversable t => Traversable (Weighted w t) where
+  traverse f (Weighted twa) = Weighted <$> (traverse . traverse) f twa
 
 evenlyWeighted :: [a] -> Weighted Int [] a
 evenlyWeighted = Weighted . itoList
 
+-- | Weight the elements of some functor by a function. Larger values of 'w' per
+-- its 'Ord' instance will be more likely to be generated
 weightedBy :: Functor t => (a -> w) -> t a -> Weighted w t a
 weightedBy weighting xs = Weighted $ (weighting &&& id) <$> xs
 
@@ -96,6 +104,14 @@ chance
   -> m Bool
 chance n = choose $ weightedBy (bool 1 (n * 2)) bools
 
+-- | Choose a random subset of *about* @w@ of the elements of the given
+-- 'Witherable' structure
+chooseSubset :: ( Num w, Ord w, Distribution Uniform w, Excludable w
+               , Witherable t
+               , MonadRandom m
+               ) => w -> t a -> m (t a)
+chooseSubset = filterA . const . chance
+
 --------------------------------------------------------------------------------
 
 bools :: NonEmpty Bool