{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE PartialTypeSignatures #-}
module Data.Registry.Internal.Hedgehog (
GenIO
, Chooser (..)
, cycleWith
, chooseOne
, choiceChooser
, cycleChooser
, distinct
, distinctWith
, sampleIO
) where
import Control.Monad.Trans.Maybe
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 (Node (..), Tree (..))
import Prelude (show, (!!))
import Protolude as P
type GenIO = GenT IO
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.filter (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
r <- evalGenIO 30 seed gen
case r of
Nothing ->
loop (n - 1)
Just a ->
pure a
in
loop (100 :: Int)
evalGenIO :: Size -> Seed -> GenIO a -> IO (Maybe a)
evalGenIO size seed g = do
r <- runMaybeT . runTree $ runGenT size seed g
pure $ case r of
Nothing -> Nothing
Just (Node a _) -> Just a