Fix up frequency

hedgehog/hedgehog.cabal

@@ -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

hedgehog/src/Hedgehog/Internal/Gen.hs

@@ -1,5 +1,6 @@
 {-# OPTIONS_HADDOCK not-home #-}
 {-# LANGUAGE ApplicativeDo #-}
+{-# LANGUAGE BangPatterns #-}
 {-# LANGUAGE CPP #-}
 {-# LANGUAGE DataKinds #-}
 {-# LANGUAGE DeriveFoldable #-}
@@ -12,6 +13,7 @@
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
 {-# LANGUAGE LambdaCase #-}
 {-# LANGUAGE MultiParamTypeClasses #-}
+{-# LANGUAGE PatternGuards #-}
 {-# LANGUAGE RankNTypes #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE StandaloneDeriving #-}
@@ -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.Lazy as IM
+import qualified Data.List as List
 import           Data.List.NonEmpty (NonEmpty)
 import qualified Data.List.NonEmpty as NonEmpty
 import           Data.Map (Map)
@@ -1170,28 +1174,44 @@ 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.
 --
+--   If all frequencies are zero, then nothing is generated.
 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)
-
-    n <- integral $ Range.constant 1 total
-    pick n 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
+  | Just ((total, final), sum_map) <- IM.maxViewWithKey sum_map'
+  = do
+      n <- integral $ Range.constant 1 total
+      case IM.lookupGE n sum_map of
+        Just (_, a) -> a
+        Nothing -> final
+  | otherwise
+  = discard -- All frequencies were 0, so we discard.
+            -- Alternatively, we could throw an error here.
+    where
+      --[(1, x), (7, y), (10, z)]  In
+      --[(1, x), (8, y), (18, z)]  Out
+      --
+      -- We convert the whole list to an IntMap, but then we
+      -- chop off the last element and use it as a "fall-through".
+      -- This avoids an "impossible" error and may make generation
+      -- ever so slightly faster when the list is short.
+      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
 
 -- | Modifies combinators which choose from a list of generators, like 'choice'
 --   or 'frequency', so that they can be used in recursive scenarios.