{-# 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 Hedgehog.Extra.Manual
  ( Manual(Manual)
  , unManual
  , toManual
  , fromManual
  , dontShrink
    -- * Combinators
  , sized
  , replicate
  , interleave
    -- * Auxiliary
  , wrapTreeT
  , unwrapTreeT
  )
where

import           Prelude hiding (replicate)

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           Hedgehog.Internal.Tree (NodeT (NodeT), TreeT (TreeT), nodeChildren, nodeValue)

import qualified Hedgehog.Internal.Seed as Seed


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


toManual :: Gen a -> Manual (TreeT (MaybeT Identity) a)
toManual :: 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 :: Manual (TreeT (MaybeT Identity) a) -> Gen a
fromManual (Manual Size -> Seed -> TreeT (MaybeT Identity) a
f) = (Size -> Seed -> TreeT (MaybeT Identity) a) -> Gen 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 :: Gen a -> Manual (Maybe a)
dontShrink = (TreeT (MaybeT Identity) a -> Maybe a)
-> Manual (TreeT (MaybeT Identity) a) -> Manual (Maybe a)
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 (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)
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 :: (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 :: 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
  <*> :: 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 :: a -> Manual a
return         = a -> Manual a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
  Manual Size -> Seed -> a
x >>= :: 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 :: (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 :: 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 :: [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]
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' :: [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 (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 :: Maybe (NodeT (MaybeT Identity) a) -> TreeT (MaybeT Identity) a
wrapTreeT = Maybe (NodeT (MaybeT Identity) a) -> TreeT (MaybeT Identity) a
coerce


unwrapTreeT :: TreeT (MaybeT Identity) a -> Maybe (NodeT (MaybeT Identity) a)
unwrapTreeT :: TreeT (MaybeT Identity) a -> Maybe (NodeT (MaybeT Identity) a)
unwrapTreeT = TreeT (MaybeT Identity) a -> Maybe (NodeT (MaybeT Identity) a)
coerce


splits :: [a] -> [([a], a, [a])]
splits :: [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 (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\([a]
as,a
b,[a]
cs) -> (a
xa -> [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 :: 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