{-# LANGUAGE CPP              #-}
{-# LANGUAGE FlexibleContexts #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
module Test.QuickCheck.Instances.Containers () where

import Prelude ()
import Test.QuickCheck.Instances.CustomPrelude

import Test.QuickCheck
       (Arbitrary (..), Arbitrary1 (..), CoArbitrary (..), Function (..), Gen,
       arbitrary1, chooseInt, functionMap, liftShrink2, shrink1, shuffle,
       sized)

import qualified Data.Tree as Tree

-------------------------------------------------------------------------------
-- containers
-------------------------------------------------------------------------------

instance Arbitrary1 Tree.Tree where
    liftArbitrary :: Gen a -> Gen (Tree a)
liftArbitrary Gen a
arb = (Int -> Gen (Tree a)) -> Gen (Tree a)
forall a. (Int -> Gen a) -> Gen a
sized ((Int -> Gen (Tree a)) -> Gen (Tree a))
-> (Int -> Gen (Tree a)) -> Gen (Tree a)
forall a b. (a -> b) -> a -> b
$ \Int
n -> do
        Int
k <- (Int, Int) -> Gen Int
chooseInt (Int
0, Int
n)
        Int -> Gen (Tree a)
go Int
k
      where
        go :: Int -> Gen (Tree a)
go Int
n = do -- n is the size of the trees.
            a
value <- Gen a
arb
            [Int]
pars <- Int -> Gen [Int]
arbPartition (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) -- can go negative!
            [Tree a]
forest <- (Int -> Gen (Tree a)) -> [Int] -> Gen [Tree a]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse Int -> Gen (Tree a)
go [Int]
pars
            Tree a -> Gen (Tree a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Tree a -> Gen (Tree a)) -> Tree a -> Gen (Tree a)
forall a b. (a -> b) -> a -> b
$ a -> [Tree a] -> Tree a
forall a. a -> Forest a -> Tree a
Tree.Node a
value [Tree a]
forest

        arbPartition :: Int -> Gen [Int]
        arbPartition :: Int -> Gen [Int]
arbPartition Int
k = case Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Int
k Int
1 of
            Ordering
LT -> [Int] -> Gen [Int]
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
            Ordering
EQ -> [Int] -> Gen [Int]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Int
1]
            Ordering
GT -> do
                Int
first <- (Int, Int) -> Gen Int
chooseInt (Int
1, Int
k)
                [Int]
rest <- Int -> Gen [Int]
arbPartition (Int -> Gen [Int]) -> Int -> Gen [Int]
forall a b. (a -> b) -> a -> b
$ Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
first
                [Int] -> Gen [Int]
forall a. [a] -> Gen [a]
shuffle (Int
first Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
rest)

    liftShrink :: (a -> [a]) -> Tree a -> [Tree a]
liftShrink a -> [a]
shr = Tree a -> [Tree a]
go
      where
        go :: Tree a -> [Tree a]
go (Tree.Node a
val [Tree a]
forest) = [Tree a]
forest [Tree a] -> [Tree a] -> [Tree a]
forall a. [a] -> [a] -> [a]
++
            [ a -> [Tree a] -> Tree a
forall a. a -> Forest a -> Tree a
Tree.Node a
e [Tree a]
fs
            | (a
e, [Tree a]
fs) <- (a -> [a])
-> ([Tree a] -> [[Tree a]]) -> (a, [Tree a]) -> [(a, [Tree a])]
forall (f :: * -> * -> *) a b.
Arbitrary2 f =>
(a -> [a]) -> (b -> [b]) -> f a b -> [f a b]
liftShrink2 a -> [a]
shr ((Tree a -> [Tree a]) -> [Tree a] -> [[Tree a]]
forall (f :: * -> *) a. Arbitrary1 f => (a -> [a]) -> f a -> [f a]
liftShrink Tree a -> [Tree a]
go) (a
val, [Tree a]
forest)
            ]

instance Arbitrary a => Arbitrary (Tree.Tree a) where
    arbitrary :: Gen (Tree a)
arbitrary = Gen (Tree a)
forall (f :: * -> *) a. (Arbitrary1 f, Arbitrary a) => Gen (f a)
arbitrary1
    shrink :: Tree a -> [Tree a]
shrink = Tree a -> [Tree a]
forall (f :: * -> *) a. (Arbitrary1 f, Arbitrary a) => f a -> [f a]
shrink1

instance CoArbitrary a => CoArbitrary (Tree.Tree a) where
    coarbitrary :: Tree a -> Gen b -> Gen b
coarbitrary (Tree.Node a
val Forest a
forest) =
        a -> Gen b -> Gen b
forall a b. CoArbitrary a => a -> Gen b -> Gen b
coarbitrary a
val (Gen b -> Gen b) -> (Gen b -> Gen b) -> Gen b -> Gen b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Forest a -> Gen b -> Gen b
forall a b. CoArbitrary a => a -> Gen b -> Gen b
coarbitrary Forest a
forest

instance Function a => Function (Tree.Tree a) where
    function :: (Tree a -> b) -> Tree a :-> b
function = (Tree a -> (a, Forest a))
-> ((a, Forest a) -> Tree a) -> (Tree a -> b) -> Tree a :-> b
forall b a c.
Function b =>
(a -> b) -> (b -> a) -> (a -> c) -> a :-> c
functionMap (\(Tree.Node a
x Forest a
xs) -> (a
x,Forest a
xs)) ((a -> Forest a -> Tree a) -> (a, Forest a) -> Tree a
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> Forest a -> Tree a
forall a. a -> Forest a -> Tree a
Tree.Node)