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

import Prelude ()
import Prelude.Compat

import Data.Traversable (for)

import Test.QuickCheck

import qualified Data.Tree as Tree

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

instance Arbitrary1 Tree.Tree where
    liftArbitrary :: Gen a -> Gen (Tree a)
liftArbitrary arb :: Gen a
arb = Gen (Tree a)
go
      where
        go :: Gen (Tree a)
go = (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
$ \n :: Int
n -> do -- Sized 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
- 1) -- can go negative!
            [Tree a]
forest <- [Int] -> (Int -> Gen (Tree a)) -> Gen [Tree a]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
t a -> (a -> f b) -> f (t b)
for [Int]
pars ((Int -> Gen (Tree a)) -> Gen [Tree a])
-> (Int -> Gen (Tree a)) -> Gen [Tree a]
forall a b. (a -> b) -> a -> b
$ \i :: Int
i -> Int -> Gen (Tree a) -> Gen (Tree a)
forall a. Int -> Gen a -> Gen a
resize Int
i Gen (Tree a)
go
            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 k :: Int
k = case Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Int
k 1 of
            LT -> [Int] -> Gen [Int]
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
            EQ -> [Int] -> Gen [Int]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [1]
            GT -> do
                Int
first <- [Int] -> Gen Int
forall a. [a] -> Gen a
elements [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 (m :: * -> *) a. Monad m => a -> m a
return ([Int] -> Gen [Int]) -> [Int] -> Gen [Int]
forall a b. (a -> b) -> a -> b
$ Int
first Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
rest

    liftShrink :: (a -> [a]) -> Tree a -> [Tree a]
liftShrink shr :: a -> [a]
shr = Tree a -> [Tree a]
go 
      where
        go :: Tree a -> [Tree a]
go (Tree.Node val :: a
val forest :: [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
            | (e :: a
e, fs :: [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 val :: a
val forest :: 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