{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE PartialTypeSignatures #-}

module Data.Registry.Internal.Hedgehog
  ( GenIO,
    Chooser (..),
    -- cycling values
    cycleWith,
    chooseOne,
    choiceChooser,
    cycleChooser,
    -- making distinct values
    distinct,
    distinctWith,
    -- utilities
    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 Protolude as P
import Prelude (show, (!!))

-- | All the generators we use are lifted into GenIO to allow some generators to be stateful
type GenIO = GenT IO

-- | Lift a pure generator into another monad like IO
liftGen :: (Monad m) => Gen a -> GenT m a
liftGen :: Gen a -> GenT m a
liftGen = (forall a. Identity a -> m a) -> Gen a -> GenT m a
forall k (t :: (* -> *) -> k -> *) (m :: * -> *) (n :: * -> *)
       (b :: k).
(MFunctor t, Monad m) =>
(forall a. m a -> n a) -> t m b -> t n b
hoist (a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> m a) -> (Identity a -> a) -> Identity a -> m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Identity a -> a
forall a. Identity a -> a
runIdentity)

-- * CHOOSING VALUES DETERMINISTICALLY

-- | Given a choosing strategy pick a generator
--   This is possibly a stateful operation
chooseOne :: GenIO Chooser -> [GenIO a] -> GenIO a
chooseOne :: GenIO Chooser -> [GenIO a] -> GenIO a
chooseOne GenIO Chooser
chooser [GenIO a]
gs = do
  Chooser
c <- GenIO Chooser
chooser
  GenT IO (GenIO a) -> GenIO a
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join (GenT IO (GenIO a) -> GenIO a) -> GenT IO (GenIO a) -> GenIO a
forall a b. (a -> b) -> a -> b
$ IO (GenIO a) -> GenT IO (GenIO a)
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
P.lift (IO (GenIO a) -> GenT IO (GenIO a))
-> IO (GenIO a) -> GenT IO (GenIO a)
forall a b. (a -> b) -> a -> b
$ Chooser -> [GenIO a] -> IO (GenIO a)
Chooser -> forall a. [GenIO a] -> IO (GenIO a)
pickOne Chooser
c [GenIO a]
gs

-- | Chooser for randomly selecting a generator
choiceChooser :: Chooser
choiceChooser :: Chooser
choiceChooser = Chooser :: Text -> (forall a. [GenIO a] -> IO (GenIO a)) -> Chooser
Chooser {chooserType :: Text
chooserType = Text
"choice", pickOne :: forall a. [GenIO a] -> IO (GenIO a)
pickOne = GenT IO a -> IO (GenT IO a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (GenT IO a -> IO (GenT IO a))
-> ([GenT IO a] -> GenT IO a) -> [GenT IO a] -> IO (GenT IO a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [GenT IO a] -> GenT IO a
forall (m :: * -> *) a. MonadGen m => [m a] -> m a
Gen.choice}

-- | Chooser for deterministically choosing elements in a list
--   by cycling over them, which requires to maintain some state about the last position
cycleChooser :: IO Chooser
cycleChooser :: IO Chooser
cycleChooser = do
  IORef Int
ref <- Int -> IO (IORef Int)
forall a. a -> IO (IORef a)
newIORef Int
0
  Chooser -> IO Chooser
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Chooser -> IO Chooser) -> Chooser -> IO Chooser
forall a b. (a -> b) -> a -> b
$ Chooser :: Text -> (forall a. [GenIO a] -> IO (GenIO a)) -> Chooser
Chooser {chooserType :: Text
chooserType = Text
"cycle", pickOne :: forall a. [GenIO a] -> IO (GenIO a)
pickOne = IORef Int -> [GenT IO a] -> IO (GenT IO a)
forall (m :: * -> *) a.
MonadIO m =>
IORef Int -> [GenT m a] -> IO (GenT m a)
cycleWith IORef Int
ref}

-- | A "chooser" strategy
--   The type can be used to debug specializations
data Chooser = Chooser
  { Chooser -> Text
chooserType :: Text,
    Chooser -> forall a. [GenIO a] -> IO (GenIO a)
pickOne :: forall a. [GenIO a] -> IO (GenIO a)
  }

instance Show Chooser where
  show :: Chooser -> String
show Chooser
c = Text -> String
forall a b. ConvertText a b => a -> b
toS (Chooser -> Text
chooserType Chooser
c)

-- | Pick a generator in a list based on the previous position selected
cycleWith :: (MonadIO m) => IORef Int -> [GenT m a] -> IO (GenT m a)
cycleWith :: IORef Int -> [GenT m a] -> IO (GenT m a)
cycleWith IORef Int
ref [GenT m a]
gs = do
  Int
n <- IORef Int -> IO Int
forall a. IORef a -> IO a
readIORef IORef Int
ref
  IORef Int -> (Int -> Int) -> IO ()
forall a. IORef a -> (a -> a) -> IO ()
modifyIORef IORef Int
ref Int -> Int
increment
  GenT m a -> IO (GenT m a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([GenT m a]
gs [GenT m a] -> Int -> GenT m a
forall a. [a] -> Int -> a
!! Int
n)
  where
    increment :: Int -> Int
increment Int
i = if Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [GenT m a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
P.length [GenT m a]
gs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 then Int
0 else Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1

-- * MAKING DISTINCT VALUES

-- | Create a generator for distinct values
--   This is a stateful operation
distinct :: (MonadIO m, Eq a) => GenT m a -> IO (GenT m a)
distinct :: GenT m a -> IO (GenT m a)
distinct GenT m a
g = do
  IORef [a]
ref <- [a] -> IO (IORef [a])
forall a. a -> IO (IORef a)
newIORef []
  GenT m a -> IO (GenT m a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (GenT m a -> IO (GenT m a)) -> GenT m a -> IO (GenT m a)
forall a b. (a -> b) -> a -> b
$ IORef [a] -> GenT m a -> GenT m a
forall (m :: * -> *) a.
(MonadIO m, Eq a) =>
IORef [a] -> GenT m a -> GenT m a
distinctWith IORef [a]
ref GenT m a
g

-- | Generate distinct values based on the values already generated
distinctWith :: (MonadIO m, Eq a) => IORef [a] -> GenT m a -> GenT m a
distinctWith :: IORef [a] -> GenT m a -> GenT m a
distinctWith IORef [a]
ref GenT m a
g = (Size -> Seed -> TreeT (MaybeT m) a) -> GenT m a
forall (m :: * -> *) a.
(Size -> Seed -> TreeT (MaybeT m) a) -> GenT m a
GenT ((Size -> Seed -> TreeT (MaybeT m) a) -> GenT m a)
-> (Size -> Seed -> TreeT (MaybeT m) a) -> GenT m a
forall a b. (a -> b) -> a -> b
$ \Size
size Seed
seed -> do
  [a]
as <- IO [a] -> TreeT (MaybeT m) [a]
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO [a] -> TreeT (MaybeT m) [a]) -> IO [a] -> TreeT (MaybeT m) [a]
forall a b. (a -> b) -> a -> b
$ IORef [a] -> IO [a]
forall a. IORef a -> IO a
readIORef IORef [a]
ref
  a
a <- Size -> Seed -> GenT m a -> TreeT (MaybeT m) a
forall (m :: * -> *) a.
Size -> Seed -> GenT m a -> TreeT (MaybeT m) a
runGenT Size
size Seed
seed (GenT m a -> TreeT (MaybeT m) a) -> GenT m a -> TreeT (MaybeT m) a
forall a b. (a -> b) -> a -> b
$ ((a -> Bool) -> GenT m a -> GenT m a
forall (m :: * -> *) a. MonadGen m => (a -> Bool) -> m a -> m a
Gen.filterT (Bool -> Bool
not (Bool -> Bool) -> (a -> Bool) -> a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> [a] -> Bool) -> [a] -> a -> Bool
forall a b c. (a -> b -> c) -> b -> a -> c
flip a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
elem [a]
as)) GenT m a
g
  IO () -> TreeT (MaybeT m) ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> TreeT (MaybeT m) ()) -> IO () -> TreeT (MaybeT m) ()
forall a b. (a -> b) -> a -> b
$ IORef [a] -> [a] -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef [a]
ref (a
a a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
as)
  a -> TreeT (MaybeT m) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
a

-- * UTILITIES

-- | Sample GenIO values
sampleIO :: GenIO a -> IO a
sampleIO :: GenIO a -> IO a
sampleIO GenIO a
gen =
  let loop :: Int -> IO a
loop Int
n =
        if Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0
          then Text -> IO a
forall a. HasCallStack => Text -> a
panic Text
"Hedgehog.Gen.sample: too many discards, could not generate a sample"
          else do
            Seed
seed <- IO Seed
forall (m :: * -> *). MonadIO m => m Seed
Seed.random
            NodeT Maybe a
r [TreeT IO (Maybe a)]
_ <- TreeT IO (Maybe a) -> IO (NodeT IO (Maybe a))
forall (m :: * -> *) a. TreeT m a -> m (NodeT m a)
runTreeT (TreeT IO (Maybe a) -> IO (NodeT IO (Maybe a)))
-> TreeT IO (Maybe a) -> IO (NodeT IO (Maybe a))
forall a b. (a -> b) -> a -> b
$ Size -> Seed -> GenIO a -> TreeT IO (Maybe a)
forall (m :: * -> *) a.
Monad m =>
Size -> Seed -> GenT m a -> TreeT m (Maybe a)
evalGenT Size
30 Seed
seed GenIO a
gen
            case Maybe a
r of
              Maybe a
Nothing ->
                Int -> IO a
loop (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
              Just a
a ->
                a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
a
   in Int -> IO a
loop (Int
100 :: Int)