{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE PartialTypeSignatures #-}
module Data.Registry.Internal.Hedgehog (
GenIO
, Chooser (..)
, cycleWith
, chooseOne
, choiceChooser
, cycleChooser
, distinct
, distinctWith
, liftGen
, sampleIO
) where
import Control.Monad.Morph
import Data.IORef
import Data.Maybe as Maybe
import Hedgehog
import Hedgehog.Gen as Gen
import Hedgehog.Internal.Gen as Gen
import Hedgehog.Internal.Seed as Seed (random)
import Hedgehog.Internal.Tree as Tree (NodeT (..), runTreeT)
import Prelude (show, (!!))
import Protolude as P
type GenIO = GenT IO
liftGen :: (Monad m) => Gen a -> GenT m a
liftGen = hoist (pure . runIdentity)
chooseOne :: GenIO Chooser -> [GenIO a] -> GenIO a
chooseOne chooser gs = do
c <- chooser
join $ P.lift $ pickOne c gs
choiceChooser :: Chooser
choiceChooser = Chooser { chooserType = "choice", pickOne = pure . Gen.choice }
cycleChooser :: IO Chooser
cycleChooser = do
ref <- newIORef 0
pure $ Chooser { chooserType = "cycle", pickOne = cycleWith ref }
data Chooser = Chooser {
chooserType :: Text
, pickOne :: forall a . [GenIO a] -> IO (GenIO a)
}
instance Show Chooser where
show c = toS (chooserType c)
cycleWith :: (MonadIO m) => IORef Int -> [GenT m a] -> IO (GenT m a)
cycleWith ref gs = do
n <- readIORef ref
modifyIORef ref increment
pure (gs !! n)
where increment i = if i == P.length gs - 1 then 0 else i + 1
distinct :: (MonadIO m, Eq a) => GenT m a -> IO (GenT m a)
distinct g = do
ref <- newIORef []
pure $ distinctWith ref g
distinctWith :: (MonadIO m, Eq a) => IORef [a] -> GenT m a -> GenT m a
distinctWith ref g = GenT $ \size seed -> do
as <- liftIO $ readIORef ref
a <- runGenT size seed $ (Gen.filterT (not . flip elem as)) g
liftIO $ writeIORef ref (a:as)
pure a
sampleIO :: GenIO a -> IO a
sampleIO gen =
let
loop n =
if n <= 0 then
panic "Hedgehog.Gen.sample: too many discards, could not generate a sample"
else do
seed <- Seed.random
NodeT r _ <- runTreeT $ evalGenT 30 seed gen
case r of
Nothing ->
loop (n - 1)
Just a ->
pure a
in
loop (100 :: Int)