@@ -44,12 +44,10 @@ import Data.Array ((:), length, zip, sortBy)
4444import Data.Array.NonEmpty (NonEmptyArray )
4545import Data.Array.NonEmpty as NEA
4646import Data.Enum (class BoundedEnum , fromEnum , toEnum )
47- import Data.Foldable (fold )
47+ import Data.Semigroup. Foldable (foldMap1 )
4848import Data.Int (toNumber , floor )
4949import Data.List (List (..), toUnfoldable )
50- import Data.List.NonEmpty (NonEmptyList )
51- import Data.List.NonEmpty as NEL
52- import Data.Maybe (fromJust )
50+ import Data.Maybe (Maybe (..), fromJust )
5351import Data.Monoid.Additive (Additive (..))
5452import Data.Newtype (unwrap )
5553import Data.Tuple (Tuple (..), fst , snd )
@@ -162,15 +160,19 @@ oneOf xs = do
162160
163161-- | Create a random generator which selects and executes a random generator from
164162-- | a non-empty, weighted list of random generators.
165- frequency :: forall a . NonEmptyList (Tuple Number (Gen a )) -> Gen a
166- frequency = NEL .uncons >>> \{ head: x, tail: xs } -> let
167- xxs = Cons x xs
168- total = unwrap $ fold (map (Additive <<< fst) xxs :: List (Additive Number ))
169- pick _ d Nil = d
170- pick n d (Cons (Tuple k x') xs') = if n <= k then x' else pick (n - k) d xs'
163+ frequency :: forall a . NonEmptyArray (Tuple Number (Gen a )) -> Gen a
164+ frequency xxs =
165+ let
166+ default = snd $ NEA .head xxs
167+ total = unwrap $ foldMap1 (Additive <<< fst) xxs
168+ pick i n = case NEA .index xxs i of
169+ Nothing -> default
170+ Just (Tuple k x')
171+ | n <= k -> x'
172+ | otherwise -> pick (i + 1 ) (n - k)
171173 in do
172174 n <- choose zero total
173- pick n (snd x) xxs
175+ pick 0 n
174176
175177-- | Create a random generator which generates an array of random values.
176178arrayOf :: forall a . Gen a -> Gen (Array a )
0 commit comments