{-# OPTIONS_GHC -fno-warn-orphans #-}
{-# LANGUAGE CPP #-}

module Data.GenValidity.Tree where
#if !MIN_VERSION_base(4,8,0)
import Control.Applicative ((<*>))
import Data.Functor ((<$>))
#endif
import Data.GenValidity
import Data.List.NonEmpty (NonEmpty(..))
import qualified Data.List.NonEmpty as NE
import Data.Validity.Tree ()

import Test.QuickCheck

import Data.Tree

instance GenUnchecked a => GenUnchecked (Tree a) where
  genUnchecked :: Gen (Tree a)
genUnchecked = Gen a -> Gen (Tree a)
forall a. Gen a -> Gen (Tree a)
genTreeOf Gen a
forall a. GenUnchecked a => Gen a
genUnchecked
  shrinkUnchecked :: Tree a -> [Tree a]
shrinkUnchecked (Node v :: a
v ts :: [Tree a]
ts) = [a -> [Tree a] -> Tree a
forall a. a -> Forest a -> Tree a
Node a
v' [Tree a]
ts' | (v' :: a
v', ts' :: [Tree a]
ts') <- (a, [Tree a]) -> [(a, [Tree a])]
forall a. GenUnchecked a => a -> [a]
shrinkUnchecked (a
v, [Tree a]
ts)]

instance GenValid a => GenValid (Tree a) where
  genValid :: Gen (Tree a)
genValid = Gen a -> Gen (Tree a)
forall a. Gen a -> Gen (Tree a)
genTreeOf Gen a
forall a. GenValid a => Gen a
genValid
  shrinkValid :: Tree a -> [Tree a]
shrinkValid (Node v :: a
v ts :: [Tree a]
ts) = [a -> [Tree a] -> Tree a
forall a. a -> Forest a -> Tree a
Node a
v' [Tree a]
ts' | (v' :: a
v', ts' :: [Tree a]
ts') <- (a, [Tree a]) -> [(a, [Tree a])]
forall a. GenValid a => a -> [a]
shrinkValid (a
v, [Tree a]
ts)]

-- | There should be at least one invalid element, either it's here or it's
-- further down the tree.
instance (GenUnchecked a, GenInvalid a) => GenInvalid (Tree a) where
  genInvalid :: Gen (Tree a)
genInvalid =
    (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
      Int
size <- Int -> Gen Int
upTo Int
n
      (a :: Int
a, b :: Int
b) <- Int -> Gen (Int, Int)
genSplit Int
size
      [Gen (Tree a)] -> Gen (Tree a)
forall a. [Gen a] -> Gen a
oneof
        [ a -> [Tree a] -> Tree a
forall a. a -> Forest a -> Tree a
Node (a -> [Tree a] -> Tree a) -> Gen a -> Gen ([Tree a] -> Tree a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Gen a -> Gen a
forall a. Int -> Gen a -> Gen a
resize Int
a Gen a
forall a. GenInvalid a => Gen a
genInvalid Gen ([Tree a] -> Tree a) -> Gen [Tree a] -> Gen (Tree a)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Int -> Gen [Tree a] -> Gen [Tree a]
forall a. Int -> Gen a -> Gen a
resize Int
b Gen [Tree a]
forall a. GenUnchecked a => Gen a
genUnchecked
        , a -> [Tree a] -> Tree a
forall a. a -> Forest a -> Tree a
Node (a -> [Tree a] -> Tree a) -> Gen a -> Gen ([Tree a] -> Tree a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Gen a -> Gen a
forall a. Int -> Gen a -> Gen a
resize Int
a Gen a
forall a. GenUnchecked a => Gen a
genUnchecked Gen ([Tree a] -> Tree a) -> Gen [Tree a] -> Gen (Tree a)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Int -> Gen [Tree a] -> Gen [Tree a]
forall a. Int -> Gen a -> Gen a
resize Int
b Gen [Tree a]
forall a. GenInvalid a => Gen a
genInvalid
        ]
  shrinkInvalid :: Tree a -> [Tree a]
shrinkInvalid (Node v :: a
v ts :: [Tree a]
ts) =
    if a -> Bool
forall a. Validity a => a -> Bool
isInvalid a
v
      then a -> [Tree a] -> Tree a
forall a. a -> Forest a -> Tree a
Node (a -> [Tree a] -> Tree a) -> [a] -> [[Tree a] -> Tree a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> [a]
forall a. GenInvalid a => a -> [a]
shrinkInvalid a
v [[Tree a] -> Tree a] -> [[Tree a]] -> [Tree a]
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [Tree a] -> [[Tree a]]
forall a. GenUnchecked a => a -> [a]
shrinkUnchecked [Tree a]
ts
      else a -> [Tree a] -> Tree a
forall a. a -> Forest a -> Tree a
Node (a -> [Tree a] -> Tree a) -> [a] -> [[Tree a] -> Tree a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> [a]
forall a. GenUnchecked a => a -> [a]
shrinkUnchecked a
v [[Tree a] -> Tree a] -> [[Tree a]] -> [Tree a]
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [Tree a] -> [[Tree a]]
forall a. GenInvalid a => a -> [a]
shrinkInvalid [Tree a]
ts

-- | Generate a tree of values that are generated as specified.
--
-- This takes the size parameter much better into account
genTreeOf :: Gen a -> Gen (Tree a)
genTreeOf :: Gen a -> Gen (Tree a)
genTreeOf func :: Gen a
func = do
    NonEmpty a
ne <- Gen a -> Gen (NonEmpty a)
forall a. Gen a -> Gen (NonEmpty a)
genNonEmptyOf Gen a
func
    NonEmpty a -> Gen (Tree a)
forall a. NonEmpty a -> Gen (Tree a)
turnIntoTree NonEmpty a
ne
 where
        turnIntoTree :: NonEmpty a -> Gen (Tree a)
        turnIntoTree :: NonEmpty a -> Gen (Tree a)
turnIntoTree (e :: a
e :| es :: [a]
es) = do
          [NonEmpty a]
groups <- [a] -> Gen [NonEmpty a]
forall a. [a] -> Gen [NonEmpty a]
turnIntoGroups [a]
es
          [Tree a]
subtrees <- (NonEmpty a -> Gen (Tree a)) -> [NonEmpty a] -> Gen [Tree a]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM NonEmpty a -> Gen (Tree a)
forall a. NonEmpty a -> Gen (Tree a)
turnIntoTree [NonEmpty a]
groups
          Tree a -> Gen (Tree a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> [Tree a] -> Tree a
forall a. a -> Forest a -> Tree a
Node a
e [Tree a]
subtrees)

        turnIntoGroups :: [a] -> Gen [NonEmpty a]
        turnIntoGroups :: [a] -> Gen [NonEmpty a]
turnIntoGroups = [a] -> [a] -> Gen [NonEmpty a]
forall a. [a] -> [a] -> Gen [NonEmpty a]
go []
          where
            go :: [a] -> [a] -> Gen [NonEmpty a]
            go :: [a] -> [a] -> Gen [NonEmpty a]
go acc :: [a]
acc [] =
              case [a] -> Maybe (NonEmpty a)
forall a. [a] -> Maybe (NonEmpty a)
NE.nonEmpty [a]
acc of
                Nothing -> [NonEmpty a] -> Gen [NonEmpty a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
                Just ne :: NonEmpty a
ne -> [NonEmpty a] -> Gen [NonEmpty a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [NonEmpty a
ne]
            go acc :: [a]
acc (e :: a
e:es :: [a]
es) =
              [(Int, Gen [NonEmpty a])] -> Gen [NonEmpty a]
forall a. [(Int, Gen a)] -> Gen a
frequency
                [ ( 1
                  , do [NonEmpty a]
rest <- [a] -> [a] -> Gen [NonEmpty a]
forall a. [a] -> [a] -> Gen [NonEmpty a]
go [] [a]
es
                       [NonEmpty a] -> Gen [NonEmpty a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((a
e a -> [a] -> NonEmpty a
forall a. a -> [a] -> NonEmpty a
:| [a]
acc) NonEmpty a -> [NonEmpty a] -> [NonEmpty a]
forall a. a -> [a] -> [a]
: [NonEmpty a]
rest))
                , (4, [a] -> [a] -> Gen [NonEmpty a]
forall a. [a] -> [a] -> Gen [NonEmpty a]
go (a
e a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
acc) [a]
es)
                ]