{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE EmptyCase #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE QuantifiedConstraints #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE ViewPatterns #-}

module Constrained.GenT where

import Control.Monad
import Data.Foldable
import Data.List.NonEmpty qualified as NE
import GHC.Stack
import System.Random
import Test.QuickCheck hiding (Args, Fun)
import Test.QuickCheck.Gen
import Test.QuickCheck.Random

------------------------------------------------------------------------
-- The Gen Error monad
------------------------------------------------------------------------

-- | A class for different types of errors with a stack of `explain` calls to
-- narrow down problems. The (NE.NonEmpty String) means one cannot cause an
-- Error without at least 1 string to explain it.
class Monad m => MonadGenError m where
  genError :: HasCallStack => NE.NonEmpty String -> m a
  fatalError :: HasCallStack => NE.NonEmpty String -> m a
  explain :: HasCallStack => NE.NonEmpty String -> m a -> m a

-- | genError with one line of explanation
genError1 :: MonadGenError m => String -> m a
genError1 :: forall (m :: * -> *) a. MonadGenError m => String -> m a
genError1 String
s = forall (m :: * -> *) a.
(MonadGenError m, HasCallStack) =>
NonEmpty String -> m a
genError (forall (f :: * -> *) a. Applicative f => a -> f a
pure String
s)

-- | fatalError with one line of explanation
fatalError1 :: MonadGenError m => String -> m a
fatalError1 :: forall (m :: * -> *) a. MonadGenError m => String -> m a
fatalError1 String
s = forall (m :: * -> *) a.
(MonadGenError m, HasCallStack) =>
NonEmpty String -> m a
fatalError (forall (f :: * -> *) a. Applicative f => a -> f a
pure String
s)

-- | explain with one line of explanation
explain1 :: MonadGenError m => String -> m a -> m a
explain1 :: forall (m :: * -> *) a. MonadGenError m => String -> m a -> m a
explain1 String
s = forall (m :: * -> *) a.
(MonadGenError m, HasCallStack) =>
NonEmpty String -> m a -> m a
explain (forall (f :: * -> *) a. Applicative f => a -> f a
pure String
s)

-- | The Gen Error monad, distinguishes between fatal errors
-- and non-fatal errors.
data GE a
  = FatalError [NE.NonEmpty String] (NE.NonEmpty String)
  | GenError [NE.NonEmpty String] (NE.NonEmpty String)
  | Result [NE.NonEmpty String] a
  deriving (GE a -> GE a -> Bool
GE a -> GE a -> Ordering
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
forall {a}. Ord a => Eq (GE a)
forall a. Ord a => GE a -> GE a -> Bool
forall a. Ord a => GE a -> GE a -> Ordering
forall a. Ord a => GE a -> GE a -> GE a
min :: GE a -> GE a -> GE a
$cmin :: forall a. Ord a => GE a -> GE a -> GE a
max :: GE a -> GE a -> GE a
$cmax :: forall a. Ord a => GE a -> GE a -> GE a
>= :: GE a -> GE a -> Bool
$c>= :: forall a. Ord a => GE a -> GE a -> Bool
> :: GE a -> GE a -> Bool
$c> :: forall a. Ord a => GE a -> GE a -> Bool
<= :: GE a -> GE a -> Bool
$c<= :: forall a. Ord a => GE a -> GE a -> Bool
< :: GE a -> GE a -> Bool
$c< :: forall a. Ord a => GE a -> GE a -> Bool
compare :: GE a -> GE a -> Ordering
$ccompare :: forall a. Ord a => GE a -> GE a -> Ordering
Ord, GE a -> GE a -> Bool
forall a. Eq a => GE a -> GE a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: GE a -> GE a -> Bool
$c/= :: forall a. Eq a => GE a -> GE a -> Bool
== :: GE a -> GE a -> Bool
$c== :: forall a. Eq a => GE a -> GE a -> Bool
Eq, Int -> GE a -> ShowS
forall a. Show a => Int -> GE a -> ShowS
forall a. Show a => [GE a] -> ShowS
forall a. Show a => GE a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [GE a] -> ShowS
$cshowList :: forall a. Show a => [GE a] -> ShowS
show :: GE a -> String
$cshow :: forall a. Show a => GE a -> String
showsPrec :: Int -> GE a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> GE a -> ShowS
Show, forall a b. a -> GE b -> GE a
forall a b. (a -> b) -> GE a -> GE b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: forall a b. a -> GE b -> GE a
$c<$ :: forall a b. a -> GE b -> GE a
fmap :: forall a b. (a -> b) -> GE a -> GE b
$cfmap :: forall a b. (a -> b) -> GE a -> GE b
Functor)

instance Applicative GE where
  pure :: forall a. a -> GE a
pure = forall a. [NonEmpty String] -> a -> GE a
Result []
  <*> :: forall a b. GE (a -> b) -> GE a -> GE b
(<*>) = forall (m :: * -> *) a b. Monad m => m (a -> b) -> m a -> m b
ap

instance Monad GE where
  FatalError [NonEmpty String]
es NonEmpty String
err >>= :: forall a b. GE a -> (a -> GE b) -> GE b
>>= a -> GE b
_ = forall a. [NonEmpty String] -> NonEmpty String -> GE a
FatalError [NonEmpty String]
es NonEmpty String
err
  GenError [NonEmpty String]
es NonEmpty String
err >>= a -> GE b
_ = forall a. [NonEmpty String] -> NonEmpty String -> GE a
GenError [NonEmpty String]
es NonEmpty String
err
  Result [NonEmpty String]
_ a
a >>= a -> GE b
k = a -> GE b
k a
a

instance MonadGenError GE where
  genError :: forall a. HasCallStack => NonEmpty String -> GE a
genError NonEmpty String
neStr = forall a. [NonEmpty String] -> NonEmpty String -> GE a
GenError [] NonEmpty String
neStr
  fatalError :: forall a. HasCallStack => NonEmpty String -> GE a
fatalError NonEmpty String
neStr = forall a. [NonEmpty String] -> NonEmpty String -> GE a
FatalError [] NonEmpty String
neStr
  explain :: forall a. HasCallStack => NonEmpty String -> GE a -> GE a
explain NonEmpty String
nes GE a
ge = case GE a
ge of
    GenError [NonEmpty String]
es' NonEmpty String
err -> forall a. [NonEmpty String] -> NonEmpty String -> GE a
GenError (NonEmpty String
nes forall a. a -> [a] -> [a]
: [NonEmpty String]
es') NonEmpty String
err
    FatalError [NonEmpty String]
es' NonEmpty String
err -> forall a. [NonEmpty String] -> NonEmpty String -> GE a
FatalError (NonEmpty String
nes forall a. a -> [a] -> [a]
: [NonEmpty String]
es') NonEmpty String
err
    Result [NonEmpty String]
es' a
a -> forall a. [NonEmpty String] -> a -> GE a
Result (NonEmpty String
nes forall a. a -> [a] -> [a]
: [NonEmpty String]
es') a
a

instance MonadGenError m => MonadGenError (GenT m) where
  genError :: forall a. HasCallStack => NonEmpty String -> GenT m a
genError NonEmpty String
es = forall (m :: * -> *) a. (GenMode -> Gen (m a)) -> GenT m a
GenT forall a b. (a -> b) -> a -> b
$ \GenMode
_ -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
(MonadGenError m, HasCallStack) =>
NonEmpty String -> m a
genError NonEmpty String
es
  fatalError :: forall a. HasCallStack => NonEmpty String -> GenT m a
fatalError NonEmpty String
es = forall (m :: * -> *) a. (GenMode -> Gen (m a)) -> GenT m a
GenT forall a b. (a -> b) -> a -> b
$ \GenMode
_ -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
(MonadGenError m, HasCallStack) =>
NonEmpty String -> m a
fatalError NonEmpty String
es
  explain :: forall a. HasCallStack => NonEmpty String -> GenT m a -> GenT m a
explain NonEmpty String
es GenT m a
gen = forall (m :: * -> *) a. (GenMode -> Gen (m a)) -> GenT m a
GenT forall a b. (a -> b) -> a -> b
$ \GenMode
mode -> forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (m :: * -> *) a.
(MonadGenError m, HasCallStack) =>
NonEmpty String -> m a -> m a
explain NonEmpty String
es) (forall (m :: * -> *) a. GenT m a -> GenMode -> Gen (m a)
runGenT GenT m a
gen GenMode
mode)

instance MonadGenError m => MonadFail (GenT m) where
  fail :: forall a. String -> GenT m a
fail String
s = forall (m :: * -> *) a.
(MonadGenError m, HasCallStack) =>
NonEmpty String -> m a
genError (forall (f :: * -> *) a. Applicative f => a -> f a
pure String
s)

catGEs :: MonadGenError m => [GE a] -> m [a]
catGEs :: forall (m :: * -> *) a. MonadGenError m => [GE a] -> m [a]
catGEs [] = forall (f :: * -> *) a. Applicative f => a -> f a
pure []
catGEs (Result [NonEmpty String]
_ a
a : [GE a]
ges) = (a
a forall a. a -> [a] -> [a]
:) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a. MonadGenError m => [GE a] -> m [a]
catGEs [GE a]
ges
catGEs (GenError {} : [GE a]
ges) = forall (m :: * -> *) a. MonadGenError m => [GE a] -> m [a]
catGEs [GE a]
ges
catGEs (FatalError [NonEmpty String]
es NonEmpty String
e : [GE a]
_) =
  forall (m :: * -> *) r. MonadGenError m => GE r -> m r
runGE forall a b. (a -> b) -> a -> b
$ forall a. [NonEmpty String] -> NonEmpty String -> GE a
FatalError [NonEmpty String]
es NonEmpty String
e

fromGE :: (NE.NonEmpty String -> a) -> GE a -> a
fromGE :: forall a. (NonEmpty String -> a) -> GE a -> a
fromGE NonEmpty String -> a
_ (Result [NonEmpty String]
_ a
a) = a
a
fromGE NonEmpty String -> a
a (GenError [] NonEmpty String
e) = NonEmpty String -> a
a NonEmpty String
e
fromGE NonEmpty String -> a
a (GenError [NonEmpty String]
es NonEmpty String
e) = NonEmpty String -> a
a forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldr1 forall a. Semigroup a => a -> a -> a
(<>) [NonEmpty String]
es forall a. Semigroup a => a -> a -> a
<> NonEmpty String
e
fromGE NonEmpty String -> a
_ (FatalError [NonEmpty String]
es NonEmpty String
e) =
  forall a. HasCallStack => String -> a
error forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String
"\n" forall a. [a] -> [a] -> [a]
++) forall b c a. (b -> c) -> (a -> b) -> a -> c
. [String] -> String
unlines forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (forall a b. (a -> b) -> [a] -> [b]
map forall a. NonEmpty a -> [a]
NE.toList [NonEmpty String]
es) forall a. [a] -> [a] -> [a]
++ (forall a. NonEmpty a -> [a]
NE.toList NonEmpty String
e)

errorGE :: GE a -> a
errorGE :: forall a. GE a -> a
errorGE = forall a. (NonEmpty String -> a) -> GE a -> a
fromGE (forall a. HasCallStack => String -> a
error forall b c a. (b -> c) -> (a -> b) -> a -> c
. [String] -> String
unlines forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. NonEmpty a -> [a]
NE.toList)

isOk :: GE a -> Bool
isOk :: forall a. GE a -> Bool
isOk GenError {} = Bool
False
isOk FatalError {} = Bool
False
isOk Result {} = Bool
True

runGE :: MonadGenError m => GE r -> m r
runGE :: forall (m :: * -> *) r. MonadGenError m => GE r -> m r
runGE (GenError [NonEmpty String]
es NonEmpty String
err) = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr forall (m :: * -> *) a.
(MonadGenError m, HasCallStack) =>
NonEmpty String -> m a -> m a
explain (forall (m :: * -> *) a.
(MonadGenError m, HasCallStack) =>
NonEmpty String -> m a
genError NonEmpty String
err) [NonEmpty String]
es
runGE (FatalError [NonEmpty String]
es NonEmpty String
err) = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr forall (m :: * -> *) a.
(MonadGenError m, HasCallStack) =>
NonEmpty String -> m a -> m a
explain (forall (m :: * -> *) a.
(MonadGenError m, HasCallStack) =>
NonEmpty String -> m a
fatalError NonEmpty String
err) [NonEmpty String]
es
runGE (Result [NonEmpty String]
es r
a) = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr forall (m :: * -> *) a.
(MonadGenError m, HasCallStack) =>
NonEmpty String -> m a -> m a
explain (forall (f :: * -> *) a. Applicative f => a -> f a
pure r
a) [NonEmpty String]
es

fromGEProp :: Testable p => GE p -> Property
fromGEProp :: forall p. Testable p => GE p -> Property
fromGEProp (GenError [NonEmpty String]
es NonEmpty String
err) =
  forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (forall prop. Testable prop => String -> prop -> Property
counterexample forall b c a. (b -> c) -> (a -> b) -> a -> c
. [String] -> String
unlines) (forall prop. Testable prop => String -> prop -> Property
counterexample ([String] -> String
unlines (forall a. NonEmpty a -> [a]
NE.toList NonEmpty String
err)) Bool
False) (forall a b. (a -> b) -> [a] -> [b]
map forall a. NonEmpty a -> [a]
NE.toList [NonEmpty String]
es)
fromGEProp (FatalError [NonEmpty String]
es NonEmpty String
err) =
  forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (forall prop. Testable prop => String -> prop -> Property
counterexample forall b c a. (b -> c) -> (a -> b) -> a -> c
. [String] -> String
unlines) (forall prop. Testable prop => String -> prop -> Property
counterexample ([String] -> String
unlines (forall a. NonEmpty a -> [a]
NE.toList NonEmpty String
err)) Bool
False) (forall a b. (a -> b) -> [a] -> [b]
map forall a. NonEmpty a -> [a]
NE.toList [NonEmpty String]
es)
fromGEProp (Result [NonEmpty String]
es p
p) = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (forall prop. Testable prop => String -> prop -> Property
counterexample forall b c a. (b -> c) -> (a -> b) -> a -> c
. [String] -> String
unlines) (forall prop. Testable prop => prop -> Property
property p
p) (forall a b. (a -> b) -> [a] -> [b]
map forall a. NonEmpty a -> [a]
NE.toList [NonEmpty String]
es)

fromGEDiscard :: Testable p => GE p -> Property
fromGEDiscard :: forall p. Testable p => GE p -> Property
fromGEDiscard (Result [NonEmpty String]
es p
p) = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (forall prop. Testable prop => String -> prop -> Property
counterexample forall b c a. (b -> c) -> (a -> b) -> a -> c
. [String] -> String
unlines forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. NonEmpty a -> [a]
NE.toList) (forall prop. Testable prop => prop -> Property
property p
p) [NonEmpty String]
es
fromGEDiscard GE p
_ = forall a. a
discard

headGE :: Foldable t => t a -> GE a
headGE :: forall (t :: * -> *) a. Foldable t => t a -> GE a
headGE t a
t
  | a
x : [a]
_ <- forall (t :: * -> *) a. Foldable t => t a -> [a]
toList t a
t = forall (f :: * -> *) a. Applicative f => a -> f a
pure a
x
  | Bool
otherwise = forall (m :: * -> *) a.
(MonadGenError m, HasCallStack) =>
NonEmpty String -> m a
fatalError (forall (f :: * -> *) a. Applicative f => a -> f a
pure String
"head of empty structure")

-- | Turn a `GE [a]` to `[a]`, `genError` goes to `[]` and `fatalError` to `error`.
listFromGE :: GE [a] -> [a]
listFromGE :: forall a. GE [a] -> [a]
listFromGE = forall a. (NonEmpty String -> a) -> GE a -> a
fromGE (forall a b. a -> b -> a
const []) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a. MonadGenError m => String -> m a -> m a
explain1 String
"listFromGE"

------------------------------------------------------------------------
-- GenT
------------------------------------------------------------------------

-- | Generation mode - how strict are we about requiring the generator to
-- succeed. This is necessary because sometimes failing to find a value means
-- there is an actual problem (a generator _should_ be satisfiable but for
-- whatever buggy reason it isn't) and sometimes failing to find a value just
-- means there are no values. The latter case is very relevant when you're
-- generating e.g. lists or sets of values that can be empty.
data GenMode
  = Loose
  | Strict
  deriving (Eq GenMode
GenMode -> GenMode -> Bool
GenMode -> GenMode -> Ordering
GenMode -> GenMode -> GenMode
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: GenMode -> GenMode -> GenMode
$cmin :: GenMode -> GenMode -> GenMode
max :: GenMode -> GenMode -> GenMode
$cmax :: GenMode -> GenMode -> GenMode
>= :: GenMode -> GenMode -> Bool
$c>= :: GenMode -> GenMode -> Bool
> :: GenMode -> GenMode -> Bool
$c> :: GenMode -> GenMode -> Bool
<= :: GenMode -> GenMode -> Bool
$c<= :: GenMode -> GenMode -> Bool
< :: GenMode -> GenMode -> Bool
$c< :: GenMode -> GenMode -> Bool
compare :: GenMode -> GenMode -> Ordering
$ccompare :: GenMode -> GenMode -> Ordering
Ord, GenMode -> GenMode -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: GenMode -> GenMode -> Bool
$c/= :: GenMode -> GenMode -> Bool
== :: GenMode -> GenMode -> Bool
$c== :: GenMode -> GenMode -> Bool
Eq, Int -> GenMode -> ShowS
[GenMode] -> ShowS
GenMode -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [GenMode] -> ShowS
$cshowList :: [GenMode] -> ShowS
show :: GenMode -> String
$cshow :: GenMode -> String
showsPrec :: Int -> GenMode -> ShowS
$cshowsPrec :: Int -> GenMode -> ShowS
Show)

-- TODO: put a global suchThat fuel parameter in here? To avoid exponential blowup of nested such
-- thats?
newtype GenT m a = GenT {forall (m :: * -> *) a. GenT m a -> GenMode -> Gen (m a)
runGenT :: GenMode -> Gen (m a)}
  deriving (forall a b. a -> GenT m b -> GenT m a
forall a b. (a -> b) -> GenT m a -> GenT m b
forall (m :: * -> *) a b. Functor m => a -> GenT m b -> GenT m a
forall (m :: * -> *) a b.
Functor m =>
(a -> b) -> GenT m a -> GenT m b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: forall a b. a -> GenT m b -> GenT m a
$c<$ :: forall (m :: * -> *) a b. Functor m => a -> GenT m b -> GenT m a
fmap :: forall a b. (a -> b) -> GenT m a -> GenT m b
$cfmap :: forall (m :: * -> *) a b.
Functor m =>
(a -> b) -> GenT m a -> GenT m b
Functor)

instance Monad m => Applicative (GenT m) where
  pure :: forall a. a -> GenT m a
pure a
x = forall (m :: * -> *) a. (GenMode -> Gen (m a)) -> GenT m a
GenT forall a b. (a -> b) -> a -> b
$ \GenMode
_ -> forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) a. Applicative f => a -> f a
pure a
x)
  <*> :: forall a b. GenT m (a -> b) -> GenT m a -> GenT m b
(<*>) = forall (m :: * -> *) a b. Monad m => m (a -> b) -> m a -> m b
ap

instance Monad m => Monad (GenT m) where
  GenT GenMode -> Gen (m a)
m >>= :: forall a b. GenT m a -> (a -> GenT m b) -> GenT m b
>>= a -> GenT m b
k = forall (m :: * -> *) a. (GenMode -> Gen (m a)) -> GenT m a
GenT forall a b. (a -> b) -> a -> b
$ \GenMode
mode -> forall a. (QCGen -> Int -> a) -> Gen a
MkGen forall a b. (a -> b) -> a -> b
$ \QCGen
r Int
n -> do
    a
a <- forall a. Gen a -> QCGen -> Int -> a
unGen (GenMode -> Gen (m a)
m GenMode
mode) (forall a. Splittable a => a -> a
left QCGen
r) Int
n
    forall a. Gen a -> QCGen -> Int -> a
unGen (forall (m :: * -> *) a. GenT m a -> GenMode -> Gen (m a)
runGenT (a -> GenT m b
k a
a) GenMode
mode) (forall a. Splittable a => a -> a
right QCGen
r) Int
n

strictGen :: GenT m a -> Gen (m a)
strictGen :: forall (m :: * -> *) a. GenT m a -> Gen (m a)
strictGen GenT m a
gen = forall (m :: * -> *) a. GenT m a -> GenMode -> Gen (m a)
runGenT GenT m a
gen GenMode
Strict

genFromGenT :: GenT GE a -> Gen a
genFromGenT :: forall a. GenT GE a -> Gen a
genFromGenT GenT GE a
genT = forall a. GE a -> a
errorGE forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a. GenT m a -> GenMode -> Gen (m a)
runGenT GenT GE a
genT GenMode
Strict

resizeT :: (Int -> Int) -> GenT m a -> GenT m a
resizeT :: forall (m :: * -> *) a. (Int -> Int) -> GenT m a -> GenT m a
resizeT Int -> Int
f (GenT GenMode -> Gen (m a)
gm) = forall (m :: * -> *) a. (GenMode -> Gen (m a)) -> GenT m a
GenT forall a b. (a -> b) -> a -> b
$ \GenMode
mode -> forall a. (Int -> Gen a) -> Gen a
sized forall a b. (a -> b) -> a -> b
$ \Int
sz -> forall a. HasCallStack => Int -> Gen a -> Gen a
resize (Int -> Int
f Int
sz) (GenMode -> Gen (m a)
gm GenMode
mode)

pureGen :: Applicative m => Gen a -> GenT m a
pureGen :: forall (m :: * -> *) a. Applicative m => Gen a -> GenT m a
pureGen Gen a
gen = forall (m :: * -> *) a. (GenMode -> Gen (m a)) -> GenT m a
GenT forall a b. (a -> b) -> a -> b
$ \GenMode
_ -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen a
gen

listOfT :: MonadGenError m => GenT GE a -> GenT m [a]
listOfT :: forall (m :: * -> *) a. MonadGenError m => GenT GE a -> GenT m [a]
listOfT GenT GE a
gen = do
  [GE a]
lst <- forall (m :: * -> *) a. Applicative m => Gen a -> GenT m a
pureGen forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Gen a -> Gen [a]
listOf forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. GenT m a -> GenMode -> Gen (m a)
runGenT GenT GE a
gen GenMode
Loose
  forall (m :: * -> *) a. MonadGenError m => [GE a] -> m [a]
catGEs [GE a]
lst

-- | Generate a list of elements of length at most `goalLen`, but accepting failure
-- to get that many elements so long as `validLen` is true.
-- TODO: possibly one could return "more, fewer, ok" in the `validLen` instead
-- of `Bool`
listOfUntilLenT :: MonadGenError m => GenT GE a -> Int -> (Int -> Bool) -> GenT m [a]
listOfUntilLenT :: forall (m :: * -> *) a.
MonadGenError m =>
GenT GE a -> Int -> (Int -> Bool) -> GenT m [a]
listOfUntilLenT GenT GE a
gen Int
goalLen Int -> Bool
validLen =
  GenT m [a]
genList forall (m :: * -> *) a.
MonadGenError m =>
GenT m a -> (a -> Bool) -> GenT m a
`suchThatT` Int -> Bool
validLen forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) a. Foldable t => t a -> Int
length
  where
    genList :: GenT m [a]
genList = do
      [GE a]
res <- forall (m :: * -> *) a. Applicative m => Gen a -> GenT m a
pureGen forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Int -> Gen a -> Gen [a]
vectorOf Int
goalLen forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. GenT m a -> GenMode -> Gen (m a)
runGenT GenT GE a
gen GenMode
Loose
      forall (m :: * -> *) a. MonadGenError m => [GE a] -> m [a]
catGEs [GE a]
res

vectorOfT :: MonadGenError m => Int -> GenT GE a -> GenT m [a]
vectorOfT :: forall (m :: * -> *) a.
MonadGenError m =>
Int -> GenT GE a -> GenT m [a]
vectorOfT Int
i GenT GE a
gen = forall (m :: * -> *) a. (GenMode -> Gen (m a)) -> GenT m a
GenT forall a b. (a -> b) -> a -> b
$ \GenMode
mode -> do
  GE [a]
res <- forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Int -> Gen a -> Gen [a]
vectorOf Int
i forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. GenT m a -> GenMode -> Gen (m a)
runGenT GenT GE a
gen GenMode
Strict
  case GenMode
mode of
    GenMode
Strict -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) r. MonadGenError m => GE r -> m r
runGE GE [a]
res
    GenMode
Loose -> case GE [a]
res of
      FatalError [NonEmpty String]
es NonEmpty String
e -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) r. MonadGenError m => GE r -> m r
runGE (forall a. [NonEmpty String] -> NonEmpty String -> GE a
GenError [NonEmpty String]
es NonEmpty String
e)
      GE [a]
_ -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) r. MonadGenError m => GE r -> m r
runGE GE [a]
res

infixl 2 `suchThatT`
suchThatT :: MonadGenError m => GenT m a -> (a -> Bool) -> GenT m a
suchThatT :: forall (m :: * -> *) a.
MonadGenError m =>
GenT m a -> (a -> Bool) -> GenT m a
suchThatT GenT m a
g a -> Bool
p = forall (m :: * -> *) a.
MonadGenError m =>
Int -> GenT m a -> (a -> Bool) -> GenT m a
suchThatWithTryT Int
100 GenT m a
g a -> Bool
p

suchThatWithTryT :: MonadGenError m => Int -> GenT m a -> (a -> Bool) -> GenT m a
suchThatWithTryT :: forall (m :: * -> *) a.
MonadGenError m =>
Int -> GenT m a -> (a -> Bool) -> GenT m a
suchThatWithTryT Int
tries GenT m a
g a -> Bool
p = do
  GenMode
mode <- forall (m :: * -> *). Applicative m => GenT m GenMode
getMode
  let (Int
n, NonEmpty String -> GenT m a
cont) = case GenMode
mode of
        GenMode
Strict -> (Int
tries, forall (m :: * -> *) a.
(MonadGenError m, HasCallStack) =>
NonEmpty String -> m a
fatalError)
        GenMode
Loose -> (Int
1 :: Int, forall (m :: * -> *) a.
(MonadGenError m, HasCallStack) =>
NonEmpty String -> m a
genError) -- TODO: Maybe 1 is not the right number here!
  Int -> (NonEmpty String -> GenT m a) -> GenT m a
go Int
n NonEmpty String -> GenT m a
cont
  where
    go :: Int -> (NonEmpty String -> GenT m a) -> GenT m a
go Int
0 NonEmpty String -> GenT m a
cont = NonEmpty String -> GenT m a
cont (forall (f :: * -> *) a. Applicative f => a -> f a
pure (String
"Ran out of tries (" forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Int
tries forall a. [a] -> [a] -> [a]
++ String
") on suchThatWithTryT"))
    go Int
n NonEmpty String -> GenT m a
cont = do
      a
a <- GenT m a
g
      if a -> Bool
p a
a then forall (f :: * -> *) a. Applicative f => a -> f a
pure a
a else forall (m :: * -> *) a. (Int -> Int) -> GenT m a -> GenT m a
scaleT (forall a. Num a => a -> a -> a
+ Int
1) forall a b. (a -> b) -> a -> b
$ Int -> (NonEmpty String -> GenT m a) -> GenT m a
go (Int
n forall a. Num a => a -> a -> a
- Int
1) NonEmpty String -> GenT m a
cont

scaleT :: (Int -> Int) -> GenT m a -> GenT m a
scaleT :: forall (m :: * -> *) a. (Int -> Int) -> GenT m a -> GenT m a
scaleT Int -> Int
sc (GenT GenMode -> Gen (m a)
gen) = forall (m :: * -> *) a. (GenMode -> Gen (m a)) -> GenT m a
GenT forall a b. (a -> b) -> a -> b
$ \GenMode
mode -> forall a. (Int -> Int) -> Gen a -> Gen a
scale Int -> Int
sc forall a b. (a -> b) -> a -> b
$ GenMode -> Gen (m a)
gen GenMode
mode

getMode :: Applicative m => GenT m GenMode
getMode :: forall (m :: * -> *). Applicative m => GenT m GenMode
getMode = forall (m :: * -> *) a. (GenMode -> Gen (m a)) -> GenT m a
GenT forall a b. (a -> b) -> a -> b
$ \GenMode
mode -> forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall (f :: * -> *) a. Applicative f => a -> f a
pure GenMode
mode)

withMode :: GenMode -> GenT m a -> GenT m a
withMode :: forall (m :: * -> *) a. GenMode -> GenT m a -> GenT m a
withMode GenMode
mode GenT m a
gen = forall (m :: * -> *) a. (GenMode -> Gen (m a)) -> GenT m a
GenT forall a b. (a -> b) -> a -> b
$ \GenMode
_ -> forall (m :: * -> *) a. GenT m a -> GenMode -> Gen (m a)
runGenT GenT m a
gen GenMode
mode

oneofT :: MonadGenError m => [GenT GE a] -> GenT m a
oneofT :: forall (m :: * -> *) a. MonadGenError m => [GenT GE a] -> GenT m a
oneofT [GenT GE a]
gs = do
  GenMode
mode <- forall (m :: * -> *). Applicative m => GenT m GenMode
getMode
  GE a
r <-
    forall (m :: * -> *) a.
(MonadGenError m, HasCallStack) =>
NonEmpty String -> m a -> m a
explain (forall (f :: * -> *) a. Applicative f => a -> f a
pure String
"suchThatT in oneofT") forall a b. (a -> b) -> a -> b
$
      forall (m :: * -> *) a. Applicative m => Gen a -> GenT m a
pureGen (forall a. HasCallStack => [Gen a] -> Gen a
oneof [forall (m :: * -> *) a. GenT m a -> GenMode -> Gen (m a)
runGenT GenT GE a
g GenMode
mode | GenT GE a
g <- [GenT GE a]
gs]) forall (m :: * -> *) a.
MonadGenError m =>
GenT m a -> (a -> Bool) -> GenT m a
`suchThatT` forall a. GE a -> Bool
isOk
  forall (m :: * -> *) r. MonadGenError m => GE r -> m r
runGE GE a
r

frequencyT :: MonadGenError m => [(Int, GenT GE a)] -> GenT m a
frequencyT :: forall (m :: * -> *) a.
MonadGenError m =>
[(Int, GenT GE a)] -> GenT m a
frequencyT [(Int, GenT GE a)]
gs = do
  GenMode
mode <- forall (m :: * -> *). Applicative m => GenT m GenMode
getMode
  GE a
r <-
    forall (m :: * -> *) a.
(MonadGenError m, HasCallStack) =>
NonEmpty String -> m a -> m a
explain (forall (f :: * -> *) a. Applicative f => a -> f a
pure String
"suchThatT in oneofT") forall a b. (a -> b) -> a -> b
$
      forall (m :: * -> *) a. Applicative m => Gen a -> GenT m a
pureGen (forall a. HasCallStack => [(Int, Gen a)] -> Gen a
frequency [(Int
f, forall (m :: * -> *) a. GenT m a -> GenMode -> Gen (m a)
runGenT GenT GE a
g GenMode
mode) | (Int
f, GenT GE a
g) <- [(Int, GenT GE a)]
gs]) forall (m :: * -> *) a.
MonadGenError m =>
GenT m a -> (a -> Bool) -> GenT m a
`suchThatT` forall a. GE a -> Bool
isOk
  forall (m :: * -> *) r. MonadGenError m => GE r -> m r
runGE GE a
r

chooseT :: (Random a, Ord a, Show a, MonadGenError m) => (a, a) -> GenT m a
chooseT :: forall a (m :: * -> *).
(Random a, Ord a, Show a, MonadGenError m) =>
(a, a) -> GenT m a
chooseT (a
a, a
b)
  | a
b forall a. Ord a => a -> a -> Bool
< a
a = forall (m :: * -> *) a.
(MonadGenError m, HasCallStack) =>
NonEmpty String -> m a
genError (forall (f :: * -> *) a. Applicative f => a -> f a
pure (String
"chooseT (" forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show a
a forall a. [a] -> [a] -> [a]
++ String
", " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show a
b forall a. [a] -> [a] -> [a]
++ String
")"))
  | Bool
otherwise = forall (m :: * -> *) a. Applicative m => Gen a -> GenT m a
pureGen forall a b. (a -> b) -> a -> b
$ forall a. Random a => (a, a) -> Gen a
choose (a
a, a
b)

sizeT :: Monad m => GenT m Int
sizeT :: forall (m :: * -> *). Monad m => GenT m Int
sizeT = forall (m :: * -> *) a. (GenMode -> Gen (m a)) -> GenT m a
GenT forall a b. (a -> b) -> a -> b
$ \GenMode
mode -> forall a. (Int -> Gen a) -> Gen a
sized forall a b. (a -> b) -> a -> b
$ \Int
n -> forall (m :: * -> *) a. GenT m a -> GenMode -> Gen (m a)
runGenT (forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
n) GenMode
mode

-- ==================================================================
-- Reflective analysis of the internal GE structure of (GenT GE x)
-- This allows "catching" internal FatalError and GenError, and allowing
-- the program to control what happens in those cases.

-- | Always succeeds, but returns the internal GE structure for analysis
inspect :: forall m x. MonadGenError m => GenT GE x -> GenT m (GE x)
inspect :: forall (m :: * -> *) x.
MonadGenError m =>
GenT GE x -> GenT m (GE x)
inspect (GenT GenMode -> Gen (GE x)
f) = forall (m :: * -> *) a. (GenMode -> Gen (m a)) -> GenT m a
GenT GenMode -> Gen (m (GE x))
g
  where
    g :: GenMode -> Gen (m (GE x))
g GenMode
mode = do GE x
result <- GenMode -> Gen (GE x)
f GenMode
mode; forall (f :: * -> *) a. Applicative f => a -> f a
pure @Gen (forall (f :: * -> *) a. Applicative f => a -> f a
pure @m GE x
result)

-- | Ignore all kinds of Errors, by squashing them into Nothing
tryGenT :: MonadGenError m => GenT GE a -> GenT m (Maybe a)
tryGenT :: forall (m :: * -> *) a.
MonadGenError m =>
GenT GE a -> GenT m (Maybe a)
tryGenT GenT GE a
g = do
  GE a
r <- forall (m :: * -> *) x.
MonadGenError m =>
GenT GE x -> GenT m (GE x)
inspect GenT GE a
g
  case GE a
r of
    FatalError [NonEmpty String]
_ NonEmpty String
_ -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Maybe a
Nothing
    GenError [NonEmpty String]
_ NonEmpty String
_ -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Maybe a
Nothing
    Result [NonEmpty String]
es a
a -> forall (m :: * -> *) a.
(MonadGenError m, HasCallStack) =>
NonEmpty String -> m a -> m a
explain (forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldr1 forall a. Semigroup a => a -> a -> a
(<>) [NonEmpty String]
es) (forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just a
a)

-- Pass on the error messages of both kinds of Errors, by squashing and combining both of them into Left constructor
catchGenT :: MonadGenError m => GenT GE a -> GenT m (Either (NE.NonEmpty String) a)
catchGenT :: forall (m :: * -> *) a.
MonadGenError m =>
GenT GE a -> GenT m (Either (NonEmpty String) a)
catchGenT GenT GE a
g = do
  GE a
r <- forall (m :: * -> *) x.
MonadGenError m =>
GenT GE x -> GenT m (GE x)
inspect GenT GE a
g
  case GE a
r of
    FatalError [NonEmpty String]
es NonEmpty String
e -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a b. a -> Either a b
Left (forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldr1 forall a. Semigroup a => a -> a -> a
(<>) [NonEmpty String]
es forall a. Semigroup a => a -> a -> a
<> NonEmpty String
e)
    GenError [NonEmpty String]
es NonEmpty String
e -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a b. a -> Either a b
Left (forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldr1 forall a. Semigroup a => a -> a -> a
(<>) [NonEmpty String]
es forall a. Semigroup a => a -> a -> a
<> NonEmpty String
e)
    Result [NonEmpty String]
es a
a -> forall (m :: * -> *) a.
(MonadGenError m, HasCallStack) =>
NonEmpty String -> m a -> m a
explain (forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldr1 forall a. Semigroup a => a -> a -> a
(<>) [NonEmpty String]
es) (forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a b. b -> Either a b
Right a
a)

-- Pass on the error messages of both kinds of Errors in the Gen (not the GenT) monad
catchGen :: GenT GE a -> Gen (Either (NE.NonEmpty String) a)
catchGen :: forall a. GenT GE a -> Gen (Either (NonEmpty String) a)
catchGen GenT GE a
g = forall a. GenT GE a -> Gen a
genFromGenT (forall (m :: * -> *) a.
MonadGenError m =>
GenT GE a -> GenT m (Either (NonEmpty String) a)
catchGenT GenT GE a
g)

-- | Return the first successfull result from a list of computations, if they all fail
--   return a list of the error messages from each one.
firstGenT :: forall m a. MonadGenError m => [GenT GE a] -> GenT m (Either [NE.NonEmpty String] a)
firstGenT :: forall (m :: * -> *) a.
MonadGenError m =>
[GenT GE a] -> GenT m (Either [NonEmpty String] a)
firstGenT [GenT GE a]
gs = [GenT GE a]
-> [NonEmpty String] -> GenT m (Either [NonEmpty String] a)
loop [GenT GE a]
gs []
  where
    loop :: [GenT GE a] -> [NE.NonEmpty String] -> GenT m (Either [NE.NonEmpty String] a)
    loop :: [GenT GE a]
-> [NonEmpty String] -> GenT m (Either [NonEmpty String] a)
loop [] [NonEmpty String]
ys = forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall a b. a -> Either a b
Left (forall a. [a] -> [a]
reverse [NonEmpty String]
ys))
    loop (GenT GE a
x : [GenT GE a]
xs) [NonEmpty String]
ys = do
      Either (NonEmpty String) a
this <- forall (m :: * -> *) a.
MonadGenError m =>
GenT GE a -> GenT m (Either (NonEmpty String) a)
catchGenT GenT GE a
x
      case Either (NonEmpty String) a
this of
        Left NonEmpty String
zs -> [GenT GE a]
-> [NonEmpty String] -> GenT m (Either [NonEmpty String] a)
loop [GenT GE a]
xs (NonEmpty String
zs forall a. a -> [a] -> [a]
: [NonEmpty String]
ys)
        Right a
a -> forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall a b. b -> Either a b
Right a
a)

liftGen :: forall x. (forall m. MonadGenError m => GenT m x) -> GenT GE x
liftGen :: forall x.
(forall (m :: * -> *). MonadGenError m => GenT m x) -> GenT GE x
liftGen forall (m :: * -> *). MonadGenError m => GenT m x
x = forall (m :: * -> *). MonadGenError m => GenT m x
x

-- Drop a (GenT GE) computation into a (GenT m) computation. Some error information
-- is lost, as the two components of FatalError and GenError are folded into one component.
dropGen :: MonadGenError m => GenT GE a -> GenT m a
dropGen :: forall (m :: * -> *) a. MonadGenError m => GenT GE a -> GenT m a
dropGen GenT GE a
y = do
  GE a
r <- forall (m :: * -> *) x.
MonadGenError m =>
GenT GE x -> GenT m (GE x)
inspect GenT GE a
y
  case GE a
r of
    FatalError [NonEmpty String]
es NonEmpty String
e -> forall (m :: * -> *) a.
(MonadGenError m, HasCallStack) =>
NonEmpty String -> m a
fatalError (forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldr1 forall a. Semigroup a => a -> a -> a
(<>) [NonEmpty String]
es forall a. Semigroup a => a -> a -> a
<> NonEmpty String
e)
    GenError [NonEmpty String]
es NonEmpty String
e -> forall (m :: * -> *) a.
(MonadGenError m, HasCallStack) =>
NonEmpty String -> m a
genError (forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldr1 forall a. Semigroup a => a -> a -> a
(<>) [NonEmpty String]
es forall a. Semigroup a => a -> a -> a
<> NonEmpty String
e)
    Result [NonEmpty String]
es a
a -> forall (m :: * -> *) a.
(MonadGenError m, HasCallStack) =>
NonEmpty String -> m a -> m a
explain (forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldr1 forall a. Semigroup a => a -> a -> a
(<>) [NonEmpty String]
es) (forall (f :: * -> *) a. Applicative f => a -> f a
pure a
a)

-- | Run one of the actions with frequency proportional to the count. If it fails, run the other.
frequency2 :: forall m a. MonadGenError m => (Int, GenT GE a) -> (Int, GenT GE a) -> GenT m a
frequency2 :: forall (m :: * -> *) a.
MonadGenError m =>
(Int, GenT GE a) -> (Int, GenT GE a) -> GenT m a
frequency2 (Int
n, GenT GE a
g1) (Int
m, GenT GE a
g2)
  | Int
n forall a. Ord a => a -> a -> Bool
<= Int
0 Bool -> Bool -> Bool
&& Int
m forall a. Ord a => a -> a -> Bool
<= Int
0 = forall (m :: * -> *) a.
(MonadGenError m, HasCallStack) =>
NonEmpty String -> m a
fatalError (forall (f :: * -> *) a. Applicative f => a -> f a
pure (String
"Non positive frequencies in frequency2 " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show (Int
n, Int
m)))
  | Int
n forall a. Ord a => a -> a -> Bool
<= Int
0 = forall (m :: * -> *) a. MonadGenError m => GenT GE a -> GenT m a
dropGen GenT GE a
g2
  | Int
m forall a. Ord a => a -> a -> Bool
<= Int
0 = forall (m :: * -> *) a. MonadGenError m => GenT GE a -> GenT m a
dropGen GenT GE a
g1
  | Bool
True = do
      Int
i <- forall (m :: * -> *) a. Applicative m => Gen a -> GenT m a
pureGen forall a b. (a -> b) -> a -> b
$ forall a. Random a => (a, a) -> Gen a
choose (Int
1, Int
n forall a. Num a => a -> a -> a
+ Int
m)
      Either [NonEmpty String] a
ans <- if Int
i forall a. Ord a => a -> a -> Bool
<= Int
n then forall (m :: * -> *) a.
MonadGenError m =>
[GenT GE a] -> GenT m (Either [NonEmpty String] a)
firstGenT [GenT GE a
g1, GenT GE a
g2] else forall (m :: * -> *) a.
MonadGenError m =>
[GenT GE a] -> GenT m (Either [NonEmpty String] a)
firstGenT [GenT GE a
g2, GenT GE a
g1]
      case Either [NonEmpty String] a
ans of
        Left [NonEmpty String]
_ -> forall (m :: * -> *) a.
(MonadGenError m, HasCallStack) =>
NonEmpty String -> m a
fatalError (forall (f :: * -> *) a. Applicative f => a -> f a
pure String
"Both branches of frequency2 fail")
        Right a
x -> forall (f :: * -> *) a. Applicative f => a -> f a
pure a
x

-- ======================================

-- | Temporarily extend the stack while executing 'm', and revert to the old stack if successful
push :: forall m a. MonadGenError m => [String] -> GenT GE a -> GenT m a
push :: forall (m :: * -> *) a.
MonadGenError m =>
[String] -> GenT GE a -> GenT m a
push [] GenT GE a
m = forall (m :: * -> *) a. MonadGenError m => GenT GE a -> GenT m a
dropGen GenT GE a
m
push (String
x : [String]
xs) GenT GE a
m =
  case forall (m :: * -> *) a.
(MonadGenError m, HasCallStack) =>
NonEmpty String -> m a -> m a
explain (String
x forall a. a -> [a] -> NonEmpty a
NE.:| [String]
xs) GenT GE a
m of
    GenT GenMode -> Gen (GE a)
f -> (forall (m :: * -> *) a. (GenMode -> Gen (m a)) -> GenT m a
GenT GenMode -> Gen (m a)
g)
      where
        g :: GenMode -> Gen (m a)
        g :: GenMode -> Gen (m a)
g GenMode
mode = do
          GE a
result <- GenMode -> Gen (GE a)
f GenMode
mode
          case GE a
result of
            Result (NonEmpty String
_ : [NonEmpty String]
ys) a
a -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) r. MonadGenError m => GE r -> m r
runGE (forall a. [NonEmpty String] -> a -> GE a
Result [NonEmpty String]
ys a
a)
            GE a
other -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) r. MonadGenError m => GE r -> m r
runGE GE a
other

pushGE :: forall a. [String] -> GE a -> GE a
pushGE :: forall a. [String] -> GE a -> GE a
pushGE [] GE a
x = GE a
x
pushGE (String
x : [String]
xs) GE a
m = do
  case forall (m :: * -> *) a.
(MonadGenError m, HasCallStack) =>
NonEmpty String -> m a -> m a
explain (String
x forall a. a -> [a] -> NonEmpty a
NE.:| [String]
xs) GE a
m of
    Result (NonEmpty String
_ : [NonEmpty String]
ys) a
a -> forall a. [NonEmpty String] -> a -> GE a
Result [NonEmpty String]
ys a
a
    GE a
other -> GE a
other