{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ScopedTypeVariables #-}

-- | Manual generators.
--
-- This module provides functions to convert hedgehog 'Gen's to and from a
-- 'Manual' generators, and functions to manipulate these manual generators.
module Test.Hedgehog.Extra.Manual (
  Manual (Manual),
  unManual,
  toManual,
  fromManual,
  dontShrink,

  -- * Combinators
  sized,
  replicate,
  interleave,

  -- * Auxiliary
  wrapTreeT,
  unwrapTreeT,
)
where

import Control.Monad (ap, liftM)
import Control.Monad.Trans.Maybe (MaybeT (MaybeT))
import Data.Coerce (coerce)
import Data.Functor.Identity (Identity (Identity))
import Data.Maybe (catMaybes, mapMaybe)
import Hedgehog (Gen, Seed, Size)
import Hedgehog.Internal.Gen (GenT (GenT))
import qualified Hedgehog.Internal.Seed as Seed
import Hedgehog.Internal.Tree (NodeT (NodeT), TreeT (TreeT), nodeChildren, nodeValue)
import Prelude hiding (replicate)

newtype Manual a = Manual {forall a. Manual a -> Size -> Seed -> a
unManual :: Size -> Seed -> a}

toManual :: Gen a -> Manual (TreeT (MaybeT Identity) a)
toManual :: forall a. Gen a -> Manual (TreeT (MaybeT Identity) a)
toManual (GenT Size -> Seed -> TreeT (MaybeT Identity) a
f) = forall a. (Size -> Seed -> a) -> Manual a
Manual Size -> Seed -> TreeT (MaybeT Identity) a
f

fromManual :: Manual (TreeT (MaybeT Identity) a) -> Gen a
fromManual :: forall a. Manual (TreeT (MaybeT Identity) a) -> Gen a
fromManual (Manual Size -> Seed -> TreeT (MaybeT Identity) a
f) = forall (m :: * -> *) a.
(Size -> Seed -> TreeT (MaybeT m) a) -> GenT m a
GenT Size -> Seed -> TreeT (MaybeT Identity) a
f

dontShrink :: Gen a -> Manual (Maybe a)
dontShrink :: forall a. Gen a -> Manual (Maybe a)
dontShrink = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall (m :: * -> *) a. NodeT m a -> a
nodeValue forall b c a. (b -> c) -> (a -> b) -> a -> c
. coerce :: forall a b. Coercible a b => a -> b
coerce) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Gen a -> Manual (TreeT (MaybeT Identity) a)
toManual

instance Functor Manual where
  fmap :: forall a b. (a -> b) -> Manual a -> Manual b
fmap = forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM

instance Applicative Manual where
  pure :: forall a. a -> Manual a
pure a
x = forall a. (Size -> Seed -> a) -> Manual a
Manual forall a b. (a -> b) -> a -> b
$ \Size
_ Seed
_ -> a
x
  <*> :: forall a b. Manual (a -> b) -> Manual a -> Manual b
(<*>) = forall (m :: * -> *) a b. Monad m => m (a -> b) -> m a -> m b
ap

instance Monad Manual where
  return :: forall a. a -> Manual a
return = forall (f :: * -> *) a. Applicative f => a -> f a
pure
  Manual Size -> Seed -> a
x >>= :: forall a b. Manual a -> (a -> Manual b) -> Manual b
>>= a -> Manual b
f = forall a. (Size -> Seed -> a) -> Manual a
Manual forall a b. (a -> b) -> a -> b
$ \Size
size Seed
seed ->
    case Seed -> (Seed, Seed)
Seed.split Seed
seed of
      (Seed
sx, Seed
sf) -> forall a. Manual a -> Size -> Seed -> a
unManual (a -> Manual b
f (Size -> Seed -> a
x Size
size Seed
sx)) Size
size Seed
sf

{-------------------------------------------------------------------------------
  Combinators
-------------------------------------------------------------------------------}

sized :: (Size -> Manual a) -> Manual a
sized :: forall a. (Size -> Manual a) -> Manual a
sized Size -> Manual a
f = forall a. (Size -> Seed -> a) -> Manual a
Manual forall a b. (a -> b) -> a -> b
$ \Size
size Seed
seed -> forall a. Manual a -> Size -> Seed -> a
unManual (Size -> Manual a
f Size
size) Size
size Seed
seed

-- | A version of 'Control.Monad.replicateM' specific to 'Manual'.
replicate :: forall a. Int -> Manual a -> Manual [a]
replicate :: forall a. Int -> Manual a -> Manual [a]
replicate Int
n (Manual Size -> Seed -> a
f) = forall a. (Size -> Seed -> a) -> Manual a
Manual forall a b. (a -> b) -> a -> b
$ \Size
size Seed
seed ->
  let go :: Int -> Seed -> [a]
      go :: Int -> Seed -> [a]
go Int
0 Seed
_ = []
      go !Int
n' Seed
s = case Seed -> (Seed, Seed)
Seed.split Seed
s of
        (Seed
s', Seed
s'') -> Size -> Seed -> a
f Size
size Seed
s' forall a. a -> [a] -> [a]
: Int -> Seed -> [a]
go (Int
n' forall a. Num a => a -> a -> a
- Int
1) Seed
s''
   in Int -> Seed -> [a]
go Int
n Seed
seed

interleave :: [TreeT (MaybeT Identity) a] -> TreeT (MaybeT Identity) [a]
interleave :: forall a.
[TreeT (MaybeT Identity) a] -> TreeT (MaybeT Identity) [a]
interleave = coerce :: forall a b. Coercible a b => a -> b
coerce (forall a. a -> Maybe a
Just forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a.
[NodeT (MaybeT Identity) a] -> NodeT (MaybeT Identity) [a]
interleave' forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. [Maybe a] -> [a]
catMaybes)

interleave' :: [NodeT (MaybeT Identity) a] -> NodeT (MaybeT Identity) [a]
interleave' :: forall a.
[NodeT (MaybeT Identity) a] -> NodeT (MaybeT Identity) [a]
interleave' [NodeT (MaybeT Identity) a]
ts =
  forall (m :: * -> *) a. a -> [TreeT m a] -> NodeT m a
NodeT (forall a b. (a -> b) -> [a] -> [b]
map forall (m :: * -> *) a. NodeT m a -> a
nodeValue [NodeT (MaybeT Identity) a]
ts) forall a b. (a -> b) -> a -> b
$
    forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat
      [ [ forall a.
Maybe (NodeT (MaybeT Identity) a) -> TreeT (MaybeT Identity) a
wrapTreeT forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall a.
[NodeT (MaybeT Identity) a] -> NodeT (MaybeT Identity) [a]
interleave' [NodeT (MaybeT Identity) a]
ts'
        | Int
chunkSize <- [Int]
chunkSizes
        , [NodeT (MaybeT Identity) a]
ts' <- forall a. Int -> [a] -> [[a]]
removes Int
chunkSize [NodeT (MaybeT Identity) a]
ts
        ]
      , [ forall a.
Maybe (NodeT (MaybeT Identity) a) -> TreeT (MaybeT Identity) a
wrapTreeT forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall a.
[NodeT (MaybeT Identity) a] -> NodeT (MaybeT Identity) [a]
interleave' ([NodeT (MaybeT Identity) a]
xs forall a. [a] -> [a] -> [a]
++ [NodeT (MaybeT Identity) a
y'] forall a. [a] -> [a] -> [a]
++ [NodeT (MaybeT Identity) a]
zs)
        | ([NodeT (MaybeT Identity) a]
xs, NodeT (MaybeT Identity) a
y, [NodeT (MaybeT Identity) a]
zs) <- forall a. [a] -> [([a], a, [a])]
splits [NodeT (MaybeT Identity) a]
ts
        , NodeT (MaybeT Identity) a
y' <- forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe forall a.
TreeT (MaybeT Identity) a -> Maybe (NodeT (MaybeT Identity) a)
unwrapTreeT (forall (m :: * -> *) a. NodeT m a -> [TreeT m a]
nodeChildren NodeT (MaybeT Identity) a
y)
        ]
      ]
  where
    -- Chunks we try to remove from the list
    --
    -- For example, if the list has length 10, @chunkSizes = [10,5,2,1]@
    chunkSizes :: [Int]
    chunkSizes :: [Int]
chunkSizes = forall a. (a -> Bool) -> [a] -> [a]
takeWhile (forall a. Ord a => a -> a -> Bool
> Int
0) forall a b. (a -> b) -> a -> b
$ forall a. (a -> a) -> a -> [a]
iterate (forall a. Integral a => a -> a -> a
`div` Int
2) (forall (t :: * -> *) a. Foldable t => t a -> Int
length [NodeT (MaybeT Identity) a]
ts)

{-------------------------------------------------------------------------------
  Auxiliary
-------------------------------------------------------------------------------}

wrapTreeT :: Maybe (NodeT (MaybeT Identity) a) -> TreeT (MaybeT Identity) a
wrapTreeT :: forall a.
Maybe (NodeT (MaybeT Identity) a) -> TreeT (MaybeT Identity) a
wrapTreeT = coerce :: forall a b. Coercible a b => a -> b
coerce

unwrapTreeT :: TreeT (MaybeT Identity) a -> Maybe (NodeT (MaybeT Identity) a)
unwrapTreeT :: forall a.
TreeT (MaybeT Identity) a -> Maybe (NodeT (MaybeT Identity) a)
unwrapTreeT = coerce :: forall a b. Coercible a b => a -> b
coerce

splits :: [a] -> [([a], a, [a])]
splits :: forall a. [a] -> [([a], a, [a])]
splits [] = []
splits (a
x : [a]
xs) = ([], a
x, [a]
xs) forall a. a -> [a] -> [a]
: forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\([a]
as, a
b, [a]
cs) -> (a
x forall a. a -> [a] -> [a]
: [a]
as, a
b, [a]
cs)) (forall a. [a] -> [([a], a, [a])]
splits [a]
xs)

-- | @removes n@ splits a list into chunks of size n and returns all possible
-- lists where one of these chunks has been removed.
--
-- Examples:
--
-- > removes 1 [1..3] == [[2,3],[1,3],[1,2]]
-- > removes 2 [1..4] == [[3,4],[1,2]]
-- > removes 2 [1..5] == [[3,4,5],[1,2,5],[1,2,3,4]]
-- > removes 3 [1..5] == [[4,5],[1,2,3]]
--
-- Note that the last chunk we delete might have fewer elements than @n@.
removes :: forall a. Int -> [a] -> [[a]]
removes :: forall a. Int -> [a] -> [[a]]
removes Int
k = [a] -> [[a]]
go
  where
    go :: [a] -> [[a]]
    go :: [a] -> [[a]]
go [] = []
    go [a]
xs = [a]
xs2 forall a. a -> [a] -> [a]
: forall a b. (a -> b) -> [a] -> [b]
map ([a]
xs1 forall a. [a] -> [a] -> [a]
++) ([a] -> [[a]]
go [a]
xs2)
      where
        ([a]
xs1, [a]
xs2) = forall a. Int -> [a] -> ([a], [a])
splitAt Int
k [a]
xs