Skip to content

Commit

Permalink
Fix up frequency
Browse files Browse the repository at this point in the history
* Make `frequency` detect negative frequencies, frequency
  total overflow, and zero total frequency.

* Rewrite `frequency` to build an `IntMap` to represent the
  frequency list, which greatly improves efficiency for long
  lists.

* Bump `containers` lower bound to 0.5.11.

Fixes #337.
  • Loading branch information
treeowl committed Nov 11, 2019
1 parent 420bd95 commit 748aa70
Show file tree
Hide file tree
Showing 2 changed files with 34 additions and 20 deletions.
2 changes: 1 addition & 1 deletion hedgehog/hedgehog.cabal
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,7 @@ library
, async >= 2.0 && < 2.3
, bytestring >= 0.10 && < 0.11
, concurrent-output >= 1.7 && < 1.11
, containers >= 0.4 && < 0.7
, containers >= 0.5.11 && < 0.7
, directory >= 1.2 && < 1.4
, erf >= 2.0 && < 2.1
, exceptions >= 0.7 && < 0.11
Expand Down
52 changes: 33 additions & 19 deletions hedgehog/src/Hedgehog/Internal/Gen.hs
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
{-# OPTIONS_HADDOCK not-home #-}
{-# LANGUAGE ApplicativeDo #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DeriveFoldable #-}
Expand All @@ -12,6 +13,7 @@
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PatternGuards #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
Expand Down Expand Up @@ -202,6 +204,8 @@ import Data.Coerce (coerce)
import Data.Foldable (for_, toList)
import Data.Functor.Identity (Identity(..))
import Data.Int (Int8, Int16, Int32, Int64)
import qualified Data.IntMap.Strict as IM
import qualified Data.List as List
import Data.List.NonEmpty (NonEmpty)
import qualified Data.List.NonEmpty as NonEmpty
import Data.Map (Map)
Expand Down Expand Up @@ -1170,28 +1174,38 @@ choice = \case
--
-- This generator shrinks towards the first generator in the list.
--
-- /The input list must be non-empty./
-- /The sum of the frequencies must be at least @1@ and at most @'maxBound' :: 'Int'@.
-- No frequency may be negative./
--
frequency :: MonadGen m => [(Int, m a)] -> m a
frequency = \case
[] ->
error "Hedgehog.Gen.frequency: used with empty list"
xs0 -> do
let
pick n = \case
[] ->
error "Hedgehog.Gen.frequency/pick: used with empty list"
(k, x) : xs ->
if n <= k then
x
else
pick (n - k) xs

total =
sum (fmap fst xs0)

-- We calculate a running sum of the individual frequencies and build
-- an IntMap mapping the results to the generators. This makes the
-- resulting generator much faster than a naive list-based one when
-- the input list is long, and not much slower when it's short.
frequency xs0 =
do
n <- integral $ Range.constant 1 total
pick n xs0
case IM.lookupGE n sum_map of
Just (_, a) -> a
Nothing -> error "Hedgehog.Gen.frequency: Something went wrong."
where
--[(1, x), (7, y), (10, z)] In
--[(1, x), (8, y), (18, z)] Out
sum_map = IM.fromDistinctAscList $ List.unfoldr go (0, xs0)
where
go (_, []) = Nothing
go (n, (k, x) : xs)
| k < 0 = error "Hedgehog.Gen.frequency: Negative frequency."
-- nk < 0 means the sum overflowed.
| nk < 0 = error "Hedgehog.Gen.frequency: Frequency sum above maxBound :: Int"
| k > 0 = Just ((nk, x), (nk, xs))
| otherwise = go (n, xs)
where !nk = n + fromIntegral k
total
| Just (mx, _) <- IM.lookupMax sum_map
= mx
| otherwise
= error "Hedgehog.Gen.frequency: frequencies sum to zero"

-- | Modifies combinators which choose from a list of generators, like 'choice'
-- or 'frequency', so that they can be used in recursive scenarios.
Expand Down

0 comments on commit 748aa70

Please sign in to comment.