Hedgehog.hs

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

module Data.Registry.Hedgehog
  ( -- creation / tweaking functions
    GenIO,
    Chooser (..),
    forallS,
    forAllT, -- re-export of forAllT for convenience purpose since we are working in GenIO
    filterGenS,
    genFun,
    genVal,
    genWith,
    modifyGenS,
    setGen,
    setGenIO,
    setGenS,
    specializeGen,
    specializeGenIO,
    specializeGenS,
    tweakGen,
    tweakGenS,
    makeNonEmpty,
    makeNonEmptyS,
    -- combinators to compose different types of generators
    eitherOf,
    hashMapOf,
    listOf,
    listOfMinMax,
    mapOf,
    maybeOf,
    nonEmptyMapOf,
    nonEmptyOf,
    pairOf,
    setOf,
    tripleOf,
    -- cycling values
    choiceChooser,
    chooseOne,
    setCycleChooser,
    setCycleChooserS,
    -- making distinct values
    distinct,
    setDistinct,
    setDistinctFor,
    setDistinctForS,
    setDistinctS,
    -- sampling for GenIO generators
    sampleIO,
  )
where

import Control.Monad.Morph
import Data.HashMap.Strict as HashMap (HashMap, fromList)
import Data.IORef
import Data.List.NonEmpty hiding (cycle, nonEmpty, (!!))
import Data.Map as Map (fromList)
import Data.Maybe as Maybe
import Data.Registry
import Data.Registry.Internal.Hedgehog
import Data.Registry.Internal.Types
import Data.Set as Set (fromList)
import Hedgehog
import Hedgehog.Gen as Gen
import Hedgehog.Internal.Property (forAllT)
import Hedgehog.Range
import Protolude as P
import System.IO.Unsafe

-- * CREATION / TWEAKING OF REGISTRY GENERATORS

-- | Create a GenIO a for a given constructor of type a
genFun :: forall a b. (ApplyVariadic GenIO a b, Typeable a, Typeable b) => a -> Typed b
genFun = funTo @GenIO

-- | Lift a Gen a into GenIO a to be added to a registry
genVal :: forall a. (Typeable a) => Gen a -> Typed (GenIO a)
genVal g = fun (liftGen g)

-- | Extract a generator from a registry
--   We use makeUnsafe assuming that the registry has been checked before
genWith :: forall a ins out. (Typeable a) => Registry ins out -> GenIO a
genWith = makeUnsafe @(GenIO a)

-- | Modify the value of a generator in a given registry
tweakGen :: forall a ins out. (Typeable a) => (a -> a) -> Registry ins out -> Registry ins out
tweakGen f = tweakUnsafe @(GenIO a) (\genA -> f <$> genA)

-- | Modify the registry for a given generator in a State monad
tweakGenS :: forall a m ins out. (Typeable a, MonadState (Registry ins out) m) => (a -> a) -> m ()
tweakGenS f = modify (tweakGen f)

-- | Set a specific generator on the registry the value of a generator in a given registry
setGen :: forall a ins out. (Typeable a) => Gen a -> Registry ins out -> Registry ins out
setGen = setGenIO . liftGen

setGenIO :: forall a ins out. (Typeable a) => GenIO a -> Registry ins out -> Registry ins out
setGenIO genA = tweakUnsafe @(GenIO a) (const genA)

-- | Set a specific generator on the registry the value of a generator in a given registry in a State monad
setGenS :: forall a m ins out. (Typeable a, MonadState (Registry ins out) m) => Gen a -> m ()
setGenS genA = modify (setGen genA)

-- | Specialize a generator in a given context
specializeGen :: forall a b ins out. (Typeable a, Typeable b, Contains (GenIO a) out) => Gen b -> Registry ins out -> Registry ins out
specializeGen g = specializeGenIO @a (liftGen g)

-- | Specialize a generator in a given context
specializeGenIO :: forall a b ins out. (Typeable a, Typeable b, Contains (GenIO a) out) => GenIO b -> Registry ins out -> Registry ins out
specializeGenIO = specialize @(GenIO a)

-- | Specialize a generator in a given context
specializeGenS :: forall a b m ins out. (Typeable a, Typeable b, Contains (GenIO a) out, MonadState (Registry ins out) m) => Gen b -> m ()
specializeGenS g = modify (specializeGen @a @b g)

-- | Modify a generator
modifyGenS :: forall a ins out. (Typeable a) => (GenIO a -> GenIO a) -> PropertyT (StateT (Registry ins out) IO) ()
modifyGenS f = modify (tweakUnsafe @(GenIO a) f)

-- | Filter a generator
filterGenS :: forall a ins out. (Typeable a) => (a -> Bool) -> PropertyT (StateT (Registry ins out) IO) ()
filterGenS = modifyGenS . Gen.filterT

-- | Get a value generated from one of the generators in the registry and modify the registry
--   using a state monad
forallS :: forall a m out. (Typeable a, Show a, MonadIO m) => PropertyT (StateT (Registry _ out) m) a
forallS = do
  r <- P.lift $ get
  withFrozenCallStack $ hoist liftIO $ forAllT (genWith @a r)

-- | Make sure there is always one element of a given type in a list of elements
makeNonEmpty :: forall a ins out. (Typeable a) => Registry ins out -> Registry ins out
makeNonEmpty r =
  -- extract a generator for one element only
  let genA = genWith @a r
   in -- add that element in front of a list of generated elements
      tweakUnsafe @(GenIO [a]) (\genAs -> (:) <$> genA <*> genAs) r

-- | Make sure there is always one element of a given type in a list of elements in a State monad
makeNonEmptyS :: forall a m ins out. (Typeable a, MonadState (Registry ins out) m) => m ()
makeNonEmptyS = modify (makeNonEmpty @a)

-- * CONTAINERS COMBINATORS

-- | Create a generator for a pair
pairOf :: forall a b. GenIO a -> GenIO b -> GenIO (a, b)
pairOf ga gb = (,) <$> ga <*> gb

-- | Create a generator for a triple
tripleOf :: forall a b c. GenIO a -> GenIO b -> GenIO c -> GenIO (a, b, c)
tripleOf ga gb gc = (,,) <$> ga <*> gb <*> gc

-- | Create a default generator for a small list of elements
listOf :: forall a. GenIO a -> GenIO [a]
listOf = Gen.list (linear 0 10)

-- | Create a default generator for a list of elements of min elements and max elements
listOfMinMax :: forall a. Int -> Int -> GenIO a -> GenIO [a]
listOfMinMax min' max' = Gen.list (linear min' max')

-- | Create a default generator for a small non-empty list of elements
nonEmptyOf :: GenIO a -> GenIO (NonEmpty a)
nonEmptyOf = Gen.nonEmpty (linear 1 10)

-- | Create a default generator for a Maybe, choosing evenly between Nothing and Just
maybeOf :: forall a. GenIO a -> GenIO (Maybe a)
maybeOf genA = choice [pure Nothing, Just <$> genA]

-- | Create a default generator for a Either, choosing evenly between Left and Right
eitherOf :: forall a b. GenIO a -> GenIO b -> GenIO (Either a b)
eitherOf genA genB = choice [Left <$> genA, Right <$> genB]

-- | Create a default generator for a small set of elements
setOf :: forall a. (Ord a) => GenIO a -> GenIO (Set a)
setOf = fmap Set.fromList . listOf

-- | Create a default generator for map of key/values
mapOf :: forall k v. (Ord k) => GenIO k -> GenIO v -> GenIO (Map k v)
mapOf gk gv = Map.fromList <$> listOf (pairOf gk gv)

-- | Create a default generator for HashMap of key/values
hashMapOf :: forall k v. (Ord k, Hashable k) => GenIO k -> GenIO v -> GenIO (HashMap k v)
hashMapOf gk gv = HashMap.fromList <$> listOf (pairOf gk gv)

-- | Create a default generator for a small non-empty map of elements
nonEmptyMapOf :: forall k v. (Ord k) => GenIO k -> GenIO v -> GenIO (Map k v)
nonEmptyMapOf gk gv = do
  h <- pairOf gk gv
  t <- listOf (pairOf gk gv)
  pure (Map.fromList (h : t))

-- * STATEFUL GENERATORS

-- * CHOOSING VALUES DETERMINISTICALLY

-- | Set a cycling chooser for a specific data type
{-# NOINLINE setCycleChooser #-}
setCycleChooser :: forall a ins out. (Typeable a, Contains (GenIO a) out) => Registry ins out -> Registry ins out
setCycleChooser r = unsafePerformIO $ do
  c <- cycleChooser
  pure $ specializeValTo @GenIO @(GenIO a) c r

-- | Set a cycling chooser for a specific data type
{-# NOINLINE setCycleChooserS #-}
setCycleChooserS :: forall a m ins out. (Typeable a, Contains (GenIO a) out, MonadState (Registry ins out) m, MonadIO m) => m ()
setCycleChooserS =
  let c = unsafePerformIO cycleChooser
   in do
        r <- get
        let r' = specializeValTo @GenIO @(GenIO a) c r
        put r'

-- * MAKING DISTINCT VALUES

-- | Generate distinct values for a specific data type
{-# NOINLINE setDistinct #-}
setDistinct :: forall a ins out. (Eq a, Typeable a, Contains (GenIO a) out) => Registry ins out -> Registry ins out
setDistinct = setDistinctWithRef @a (unsafePerformIO $ newIORef [])

setDistinctWithRef :: forall a ins out. (Eq a, Typeable a, Contains (GenIO a) out) => IORef [a] -> Registry ins out -> Registry ins out
setDistinctWithRef ref r = setGenIO (distinctWith ref (makeFast @(GenIO a) r)) r

-- | Generate distinct values for a specific data type
{-# NOINLINE setDistinctS #-}
setDistinctS :: forall a m ins out. (Eq a, Typeable a, Contains (GenIO a) out, MonadState (Registry ins out) m, MonadIO m) => m ()
setDistinctS =
  let ref = unsafePerformIO $ newIORef []
   in modify (setDistinctWithRef @a ref)

-- | Generate distinct values for a specific data type, when used inside another data type
{-# NOINLINE setDistinctFor #-}
setDistinctFor :: forall a b ins out. (Typeable a, Contains (GenIO a) out, Eq b, Typeable b, Contains (GenIO b) out) => Registry ins out -> Registry ins out
setDistinctFor = setDistinctForWithRef @a @b (unsafePerformIO $ newIORef [])

setDistinctForWithRef :: forall a b ins out. (Typeable a, Contains (GenIO a) out, Eq b, Typeable b, Contains (GenIO b) out) => IORef [b] -> Registry ins out -> Registry ins out
setDistinctForWithRef ref r = specializeGenIO @a (distinctWith ref (makeFast @(GenIO b) r)) r

-- | Generate distinct values for a specific data type, when used inside another data type
{-# NOINLINE setDistinctForS #-}
setDistinctForS :: forall a b m ins out. (Typeable a, Contains (GenIO a) out, Eq b, Typeable b, Contains (GenIO b) out, MonadState (Registry ins out) m, MonadIO m) => m ()
setDistinctForS =
  let ref = unsafePerformIO $ newIORef []
   in modify (setDistinctForWithRef @a @b ref)