{-# 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) = (Size -> Seed -> TreeT (MaybeT Identity) a)
-> Manual (TreeT (MaybeT Identity) a)
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) = (Size -> Seed -> TreeT (MaybeT Identity) a) -> GenT Identity a
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 = (TreeT (MaybeT Identity) a -> Maybe a)
-> Manual (TreeT (MaybeT Identity) a) -> Manual (Maybe a)
forall a b. (a -> b) -> Manual a -> Manual b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((NodeT (MaybeT Identity) a -> a)
-> Maybe (NodeT (MaybeT Identity) a) -> Maybe a
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap NodeT (MaybeT Identity) a -> a
forall (m :: * -> *) a. NodeT m a -> a
nodeValue (Maybe (NodeT (MaybeT Identity) a) -> Maybe a)
-> (TreeT (MaybeT Identity) a -> Maybe (NodeT (MaybeT Identity) a))
-> TreeT (MaybeT Identity) a
-> Maybe a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TreeT (MaybeT Identity) a -> Maybe (NodeT (MaybeT Identity) a)
forall a b. Coercible a b => a -> b
coerce) (Manual (TreeT (MaybeT Identity) a) -> Manual (Maybe a))
-> (Gen a -> Manual (TreeT (MaybeT Identity) a))
-> Gen a
-> Manual (Maybe a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Gen a -> Manual (TreeT (MaybeT Identity) a)
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 = (a -> b) -> Manual a -> Manual b
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 = (Size -> Seed -> a) -> Manual a
forall a. (Size -> Seed -> a) -> Manual a
Manual ((Size -> Seed -> a) -> Manual a)
-> (Size -> Seed -> a) -> Manual a
forall a b. (a -> b) -> a -> b
$ \Size
_ Seed
_ -> a
x
  <*> :: forall a b. Manual (a -> b) -> Manual a -> Manual 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 = a -> Manual a
forall a. a -> Manual a
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 = (Size -> Seed -> b) -> Manual b
forall a. (Size -> Seed -> a) -> Manual a
Manual ((Size -> Seed -> b) -> Manual b)
-> (Size -> Seed -> b) -> Manual b
forall a b. (a -> b) -> a -> b
$ \Size
size Seed
seed ->
    case Seed -> (Seed, Seed)
Seed.split Seed
seed of
      (Seed
sx, Seed
sf) -> Manual b -> Size -> Seed -> b
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 = (Size -> Seed -> a) -> Manual a
forall a. (Size -> Seed -> a) -> Manual a
Manual ((Size -> Seed -> a) -> Manual a)
-> (Size -> Seed -> a) -> Manual a
forall a b. (a -> b) -> a -> b
$ \Size
size Seed
seed -> Manual a -> Size -> Seed -> a
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) = (Size -> Seed -> [a]) -> Manual [a]
forall a. (Size -> Seed -> a) -> Manual a
Manual ((Size -> Seed -> [a]) -> Manual [a])
-> (Size -> Seed -> [a]) -> Manual [a]
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' a -> [a] -> [a]
forall a. a -> [a] -> [a]
: Int -> Seed -> [a]
go (Int
n' Int -> Int -> Int
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 = ([Maybe (NodeT (MaybeT Identity) a)]
 -> Maybe (NodeT (MaybeT Identity) [a]))
-> [TreeT (MaybeT Identity) a] -> TreeT (MaybeT Identity) [a]
forall a b. Coercible a b => a -> b
coerce (NodeT (MaybeT Identity) [a] -> Maybe (NodeT (MaybeT Identity) [a])
forall a. a -> Maybe a
Just (NodeT (MaybeT Identity) [a]
 -> Maybe (NodeT (MaybeT Identity) [a]))
-> ([Maybe (NodeT (MaybeT Identity) a)]
    -> NodeT (MaybeT Identity) [a])
-> [Maybe (NodeT (MaybeT Identity) a)]
-> Maybe (NodeT (MaybeT Identity) [a])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [NodeT (MaybeT Identity) a] -> NodeT (MaybeT Identity) [a]
forall a.
[NodeT (MaybeT Identity) a] -> NodeT (MaybeT Identity) [a]
interleave' ([NodeT (MaybeT Identity) a] -> NodeT (MaybeT Identity) [a])
-> ([Maybe (NodeT (MaybeT Identity) a)]
    -> [NodeT (MaybeT Identity) a])
-> [Maybe (NodeT (MaybeT Identity) a)]
-> NodeT (MaybeT Identity) [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Maybe (NodeT (MaybeT Identity) a)] -> [NodeT (MaybeT Identity) a]
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 =
  [a] -> [TreeT (MaybeT Identity) [a]] -> NodeT (MaybeT Identity) [a]
forall (m :: * -> *) a. a -> [TreeT m a] -> NodeT m a
NodeT ((NodeT (MaybeT Identity) a -> a)
-> [NodeT (MaybeT Identity) a] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map NodeT (MaybeT Identity) a -> a
forall (m :: * -> *) a. NodeT m a -> a
nodeValue [NodeT (MaybeT Identity) a]
ts) ([TreeT (MaybeT Identity) [a]] -> NodeT (MaybeT Identity) [a])
-> [TreeT (MaybeT Identity) [a]] -> NodeT (MaybeT Identity) [a]
forall a b. (a -> b) -> a -> b
$
    [[TreeT (MaybeT Identity) [a]]] -> [TreeT (MaybeT Identity) [a]]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat
      [ [ Maybe (NodeT (MaybeT Identity) [a]) -> TreeT (MaybeT Identity) [a]
forall a.
Maybe (NodeT (MaybeT Identity) a) -> TreeT (MaybeT Identity) a
wrapTreeT (Maybe (NodeT (MaybeT Identity) [a])
 -> TreeT (MaybeT Identity) [a])
-> (NodeT (MaybeT Identity) [a]
    -> Maybe (NodeT (MaybeT Identity) [a]))
-> NodeT (MaybeT Identity) [a]
-> TreeT (MaybeT Identity) [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NodeT (MaybeT Identity) [a] -> Maybe (NodeT (MaybeT Identity) [a])
forall a. a -> Maybe a
Just (NodeT (MaybeT Identity) [a] -> TreeT (MaybeT Identity) [a])
-> NodeT (MaybeT Identity) [a] -> TreeT (MaybeT Identity) [a]
forall a b. (a -> b) -> a -> b
$ [NodeT (MaybeT Identity) a] -> NodeT (MaybeT Identity) [a]
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' <- Int -> [NodeT (MaybeT Identity) a] -> [[NodeT (MaybeT Identity) a]]
forall a. Int -> [a] -> [[a]]
removes Int
chunkSize [NodeT (MaybeT Identity) a]
ts
        ]
      , [ Maybe (NodeT (MaybeT Identity) [a]) -> TreeT (MaybeT Identity) [a]
forall a.
Maybe (NodeT (MaybeT Identity) a) -> TreeT (MaybeT Identity) a
wrapTreeT (Maybe (NodeT (MaybeT Identity) [a])
 -> TreeT (MaybeT Identity) [a])
-> (NodeT (MaybeT Identity) [a]
    -> Maybe (NodeT (MaybeT Identity) [a]))
-> NodeT (MaybeT Identity) [a]
-> TreeT (MaybeT Identity) [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NodeT (MaybeT Identity) [a] -> Maybe (NodeT (MaybeT Identity) [a])
forall a. a -> Maybe a
Just (NodeT (MaybeT Identity) [a] -> TreeT (MaybeT Identity) [a])
-> NodeT (MaybeT Identity) [a] -> TreeT (MaybeT Identity) [a]
forall a b. (a -> b) -> a -> b
$ [NodeT (MaybeT Identity) a] -> NodeT (MaybeT Identity) [a]
forall a.
[NodeT (MaybeT Identity) a] -> NodeT (MaybeT Identity) [a]
interleave' ([NodeT (MaybeT Identity) a]
xs [NodeT (MaybeT Identity) a]
-> [NodeT (MaybeT Identity) a] -> [NodeT (MaybeT Identity) a]
forall a. [a] -> [a] -> [a]
++ [NodeT (MaybeT Identity) a
y'] [NodeT (MaybeT Identity) a]
-> [NodeT (MaybeT Identity) a] -> [NodeT (MaybeT Identity) a]
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) <- [NodeT (MaybeT Identity) a]
-> [([NodeT (MaybeT Identity) a], NodeT (MaybeT Identity) a,
     [NodeT (MaybeT Identity) a])]
forall a. [a] -> [([a], a, [a])]
splits [NodeT (MaybeT Identity) a]
ts
        , NodeT (MaybeT Identity) a
y' <- (TreeT (MaybeT Identity) a -> Maybe (NodeT (MaybeT Identity) a))
-> [TreeT (MaybeT Identity) a] -> [NodeT (MaybeT Identity) a]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe TreeT (MaybeT Identity) a -> Maybe (NodeT (MaybeT Identity) a)
forall a.
TreeT (MaybeT Identity) a -> Maybe (NodeT (MaybeT Identity) a)
unwrapTreeT (NodeT (MaybeT Identity) a -> [TreeT (MaybeT Identity) a]
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 = (Int -> Bool) -> [Int] -> [Int]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0) ([Int] -> [Int]) -> [Int] -> [Int]
forall a b. (a -> b) -> a -> b
$ (Int -> Int) -> Int -> [Int]
forall a. (a -> a) -> a -> [a]
iterate (Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2) ([NodeT (MaybeT Identity) a] -> Int
forall a. [a] -> Int
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 = Maybe (NodeT (MaybeT Identity) a) -> TreeT (MaybeT Identity) a
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 = TreeT (MaybeT Identity) a -> Maybe (NodeT (MaybeT Identity) a)
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) ([a], a, [a]) -> [([a], a, [a])] -> [([a], a, [a])]
forall a. a -> [a] -> [a]
: (([a], a, [a]) -> ([a], a, [a]))
-> [([a], a, [a])] -> [([a], a, [a])]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\([a]
as, a
b, [a]
cs) -> (a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
as, a
b, [a]
cs)) ([a] -> [([a], a, [a])]
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 [a] -> [[a]] -> [[a]]
forall a. a -> [a] -> [a]
: ([a] -> [a]) -> [[a]] -> [[a]]
forall a b. (a -> b) -> [a] -> [b]
map ([a]
xs1 [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++) ([a] -> [[a]]
go [a]
xs2)
      where
        ([a]
xs1, [a]
xs2) = Int -> [a] -> ([a], [a])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
k [a]
xs