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

module Data.GenValidity.Set
  ( genSetOf,
    genSeperate,
    genSeperateFor,
    genSeperateForNE,
    genValidSeperateFor,
    genValidSeperateForNE,
  )
where

import Data.Containers.ListUtils
import Data.GenValidity
import Data.List.NonEmpty (NonEmpty (..))
import qualified Data.List.NonEmpty as NE
import Data.Set (Set)
import qualified Data.Set as S
import Data.Validity.Set ()
import Test.QuickCheck

instance (Ord v, GenValid v) => GenValid (Set v) where
  genValid :: Gen (Set v)
genValid = [v] -> Set v
forall a. Ord a => [a] -> Set a
S.fromList ([v] -> Set v) -> Gen [v] -> Gen (Set v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen [v]
forall a. GenValid a => Gen a
genValid
  shrinkValid :: Set v -> [Set v]
shrinkValid = ([v] -> Set v) -> [[v]] -> [Set v]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [v] -> Set v
forall a. Ord a => [a] -> Set a
S.fromList ([[v]] -> [Set v]) -> (Set v -> [[v]]) -> Set v -> [Set v]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [v] -> [[v]]
forall a. GenValid a => a -> [a]
shrinkValid ([v] -> [[v]]) -> (Set v -> [v]) -> Set v -> [[v]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Set v -> [v]
forall a. Set a -> [a]
S.toList

genSetOf :: Ord v => Gen v -> Gen (Set v)
genSetOf :: Gen v -> Gen (Set v)
genSetOf Gen v
g = [v] -> Set v
forall a. Ord a => [a] -> Set a
S.fromList ([v] -> Set v) -> Gen [v] -> Gen (Set v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen v -> Gen [v]
forall a. Gen a -> Gen [a]
genListOf Gen v
g

genValidSeperateFor :: (GenValid b, Eq b) => [a] -> Gen [(b, a)]
genValidSeperateFor :: [a] -> Gen [(b, a)]
genValidSeperateFor = Gen b -> [a] -> Gen [(b, a)]
forall b a. Eq b => Gen b -> [a] -> Gen [(b, a)]
genSeperateFor Gen b
forall a. GenValid a => Gen a
genValid

genValidSeperateForNE :: (GenValid b, Eq b) => NonEmpty a -> Gen (NonEmpty (b, a))
genValidSeperateForNE :: NonEmpty a -> Gen (NonEmpty (b, a))
genValidSeperateForNE = Gen b -> NonEmpty a -> Gen (NonEmpty (b, a))
forall b a. Eq b => Gen b -> NonEmpty a -> Gen (NonEmpty (b, a))
genSeperateForNE Gen b
forall a. GenValid a => Gen a
genValid

genSeperate :: Ord a => Gen a -> Gen [a]
genSeperate :: Gen a -> Gen [a]
genSeperate Gen a
g = [a] -> [a]
forall a. Ord a => [a] -> [a]
nubOrd ([a] -> [a]) -> Gen [a] -> Gen [a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen a -> Gen [a]
forall a. Gen a -> Gen [a]
genListOf Gen a
g

-- TODO these two can likely be optimised
genSeperateFor :: Eq b => Gen b -> [a] -> Gen [(b, a)]
genSeperateFor :: Gen b -> [a] -> Gen [(b, a)]
genSeperateFor Gen b
_ [] = [(b, a)] -> Gen [(b, a)]
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
genSeperateFor Gen b
g (a
a : [a]
as) = NonEmpty (b, a) -> [(b, a)]
forall a. NonEmpty a -> [a]
NE.toList (NonEmpty (b, a) -> [(b, a)])
-> Gen (NonEmpty (b, a)) -> Gen [(b, a)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen b -> NonEmpty a -> Gen (NonEmpty (b, a))
forall b a. Eq b => Gen b -> NonEmpty a -> Gen (NonEmpty (b, a))
genSeperateForNE Gen b
g (a
a a -> [a] -> NonEmpty a
forall a. a -> [a] -> NonEmpty a
:| [a]
as)

genSeperateForNE :: Eq b => Gen b -> NonEmpty a -> Gen (NonEmpty (b, a))
genSeperateForNE :: Gen b -> NonEmpty a -> Gen (NonEmpty (b, a))
genSeperateForNE Gen b
g (a
a :| [a]
as) = do
  [(b, a)]
restTups <- Gen b -> [a] -> Gen [(b, a)]
forall b a. Eq b => Gen b -> [a] -> Gen [(b, a)]
genSeperateFor Gen b
g [a]
as
  b
b <- Gen b
g Gen b -> (b -> Bool) -> Gen b
forall a. Gen a -> (a -> Bool) -> Gen a
`suchThat` (b -> [b] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` ((b, a) -> b) -> [(b, a)] -> [b]
forall a b. (a -> b) -> [a] -> [b]
map (b, a) -> b
forall a b. (a, b) -> a
fst [(b, a)]
restTups)
  NonEmpty (b, a) -> Gen (NonEmpty (b, a))
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((b
b, a
a) (b, a) -> [(b, a)] -> NonEmpty (b, a)
forall a. a -> [a] -> NonEmpty a
:| [(b, a)]
restTups)