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

module Data.GenValidity.Map
    ( genStructurallyValidMapOf
    , genStructurallyValidMapOfInvalidValues
#if MIN_VERSION_containers(0,5,9)
    , genStructurallyInvalidMap
#endif
    ) where
#if !MIN_VERSION_base(4,8,0)
import Control.Applicative (pure, (<$>), (<*>))
#endif
import Data.GenValidity
import Data.Validity.Map ()
import Test.QuickCheck

import Data.Map (Map)
import qualified Data.Map as M
#if MIN_VERSION_containers(0,5,9)
import qualified Data.Map.Internal as Internal
#endif

#if MIN_VERSION_containers(0,5,9)
instance (Ord k, GenUnchecked k, GenUnchecked v) => GenUnchecked (Map k v) where
    genUnchecked :: Gen (Map k v)
genUnchecked =
        (Int -> Gen (Map k v)) -> Gen (Map k v)
forall a. (Int -> Gen a) -> Gen a
sized ((Int -> Gen (Map k v)) -> Gen (Map k v))
-> (Int -> Gen (Map k v)) -> Gen (Map k v)
forall a b. (a -> b) -> a -> b
$ \n :: Int
n ->
            case Int
n of
                0 -> Map k v -> Gen (Map k v)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Map k v
forall k a. Map k a
Internal.Tip
                _ -> do
                    (a :: Int
a, b :: Int
b, c :: Int
c, d :: Int
d, e :: Int
e) <- Int -> Gen (Int, Int, Int, Int, Int)
genSplit5 Int
n
                    Int -> k -> v -> Map k v -> Map k v -> Map k v
forall k a. Int -> k -> a -> Map k a -> Map k a -> Map k a
Internal.Bin (Int -> k -> v -> Map k v -> Map k v -> Map k v)
-> Gen Int -> Gen (k -> v -> Map k v -> Map k v -> Map k v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Gen Int -> Gen Int
forall a. Int -> Gen a -> Gen a
resize Int
a Gen Int
forall a. GenUnchecked a => Gen a
genUnchecked Gen (k -> v -> Map k v -> Map k v -> Map k v)
-> Gen k -> Gen (v -> Map k v -> Map k v -> Map k v)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*>
                        Int -> Gen k -> Gen k
forall a. Int -> Gen a -> Gen a
resize Int
b Gen k
forall a. GenUnchecked a => Gen a
genUnchecked Gen (v -> Map k v -> Map k v -> Map k v)
-> Gen v -> Gen (Map k v -> Map k v -> Map k v)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*>
                        Int -> Gen v -> Gen v
forall a. Int -> Gen a -> Gen a
resize Int
c Gen v
forall a. GenUnchecked a => Gen a
genUnchecked Gen (Map k v -> Map k v -> Map k v)
-> Gen (Map k v) -> Gen (Map k v -> Map k v)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*>
                        Int -> Gen (Map k v) -> Gen (Map k v)
forall a. Int -> Gen a -> Gen a
resize Int
d Gen (Map k v)
forall a. GenUnchecked a => Gen a
genUnchecked Gen (Map k v -> Map k v) -> Gen (Map k v) -> Gen (Map k v)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*>
                        Int -> Gen (Map k v) -> Gen (Map k v)
forall a. Int -> Gen a -> Gen a
resize Int
e Gen (Map k v)
forall a. GenUnchecked a => Gen a
genUnchecked
    shrinkUnchecked :: Map k v -> [Map k v]
shrinkUnchecked Internal.Tip = []
    shrinkUnchecked (Internal.Bin s :: Int
s k :: k
k a :: v
a m1 :: Map k v
m1 m2 :: Map k v
m2) =
        Map k v
forall k a. Map k a
Internal.Tip Map k v -> [Map k v] -> [Map k v]
forall a. a -> [a] -> [a]
:
        [Map k v
m1, Map k v
m2] [Map k v] -> [Map k v] -> [Map k v]
forall a. [a] -> [a] -> [a]
++
        [ Int -> k -> v -> Map k v -> Map k v -> Map k v
forall k a. Int -> k -> a -> Map k a -> Map k a -> Map k a
Internal.Bin Int
s' k
k' v
a' Map k v
m1' Map k v
m2'
        | (s' :: Int
s', k' :: k
k', a' :: v
a', m1' :: Map k v
m1', m2' :: Map k v
m2') <- (Int, k, v, Map k v, Map k v) -> [(Int, k, v, Map k v, Map k v)]
forall a. GenUnchecked a => a -> [a]
shrinkUnchecked (Int
s, k
k, v
a, Map k v
m1, Map k v
m2)
        ]
#else
instance (Ord k, GenUnchecked k, GenUnchecked v) => GenUnchecked (Map k v) where
    genUnchecked = M.fromList <$> genUnchecked
    shrinkUnchecked = fmap M.fromList . shrinkUnchecked . M.toList
#endif
instance (Show k, Ord k, GenValid k, GenValid v) => GenValid (Map k v) where
    genValid :: Gen (Map k v)
genValid = [(k, v)] -> Map k v
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(k, v)] -> Map k v) -> Gen [(k, v)] -> Gen (Map k v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen [(k, v)]
forall a. GenValid a => Gen a
genValid
    shrinkValid :: Map k v -> [Map k v]
shrinkValid = ([(k, v)] -> Map k v) -> [[(k, v)]] -> [Map k v]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [(k, v)] -> Map k v
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([[(k, v)]] -> [Map k v])
-> (Map k v -> [[(k, v)]]) -> Map k v -> [Map k v]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(k, v)] -> [[(k, v)]]
forall a. GenValid a => a -> [a]
shrinkValid ([(k, v)] -> [[(k, v)]])
-> (Map k v -> [(k, v)]) -> Map k v -> [[(k, v)]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map k v -> [(k, v)]
forall k a. Map k a -> [(k, a)]
M.toList
#if MIN_VERSION_containers(0,5,9)
instance (Show k, Ord k, GenUnchecked k, GenInvalid k, GenUnchecked v, GenInvalid v) => GenInvalid (Map k v) where
    genInvalid :: Gen (Map k v)
genInvalid =
        [Gen (Map k v)] -> Gen (Map k v)
forall a. [Gen a] -> Gen a
oneof
            [Gen (Map k v)
forall k v.
(Ord k, GenUnchecked k, GenInvalid k, GenUnchecked v,
 GenInvalid v) =>
Gen (Map k v)
genStructurallyValidMapOfInvalidValues, Gen (Map k v)
forall k v.
(Show k, Ord k, GenUnchecked k, GenUnchecked v) =>
Gen (Map k v)
genStructurallyInvalidMap]
#else
instance (Show k, Ord k, GenUnchecked k, GenInvalid k, GenUnchecked v, GenInvalid v) => GenInvalid (Map k v) where
    genInvalid = genStructurallyValidMapOfInvalidValues
#endif
genStructurallyValidMapOf :: Ord k => Gen (k, v) -> Gen (Map k v)
genStructurallyValidMapOf :: Gen (k, v) -> Gen (Map k v)
genStructurallyValidMapOf g :: Gen (k, v)
g =
    (Int -> Gen (Map k v)) -> Gen (Map k v)
forall a. (Int -> Gen a) -> Gen a
sized ((Int -> Gen (Map k v)) -> Gen (Map k v))
-> (Int -> Gen (Map k v)) -> Gen (Map k v)
forall a b. (a -> b) -> a -> b
$ \n :: Int
n ->
        case Int
n of
            0 -> Map k v -> Gen (Map k v)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Map k v
forall k a. Map k a
M.empty
            _ -> do
                (kv :: Int
kv, m :: Int
m) <- Int -> Gen (Int, Int)
genSplit Int
n
                (key :: k
key, val :: v
val) <- Int -> Gen (k, v) -> Gen (k, v)
forall a. Int -> Gen a -> Gen a
resize Int
kv Gen (k, v)
g
                Map k v
rest <- Int -> Gen (Map k v) -> Gen (Map k v)
forall a. Int -> Gen a -> Gen a
resize Int
m (Gen (Map k v) -> Gen (Map k v)) -> Gen (Map k v) -> Gen (Map k v)
forall a b. (a -> b) -> a -> b
$ Gen (k, v) -> Gen (Map k v)
forall k v. Ord k => Gen (k, v) -> Gen (Map k v)
genStructurallyValidMapOf Gen (k, v)
g
                Map k v -> Gen (Map k v)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Map k v -> Gen (Map k v)) -> Map k v -> Gen (Map k v)
forall a b. (a -> b) -> a -> b
$ k -> v -> Map k v -> Map k v
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
key v
val Map k v
rest

-- Note: M.fromList <$> genInvalid does not work because of this line in the Data.Map documentation:
-- ' If the list contains more than one value for the same key, the last value for the key is retained.'
genStructurallyValidMapOfInvalidValues ::
       (Ord k, GenUnchecked k, GenInvalid k, GenUnchecked v, GenInvalid v) => Gen (Map k v)
genStructurallyValidMapOfInvalidValues :: Gen (Map k v)
genStructurallyValidMapOfInvalidValues =
    (Int -> Gen (Map k v)) -> Gen (Map k v)
forall a. (Int -> Gen a) -> Gen a
sized ((Int -> Gen (Map k v)) -> Gen (Map k v))
-> (Int -> Gen (Map k v)) -> Gen (Map k v)
forall a b. (a -> b) -> a -> b
$ \n :: Int
n -> do
        (k :: Int
k, v :: Int
v, m :: Int
m) <- Int -> Gen (Int, Int, Int)
genSplit3 Int
n
        let go :: Gen k -> Gen a -> Gen (Map k a)
go g1 :: Gen k
g1 g2 :: Gen a
g2 = do
                k
key <- Int -> Gen k -> Gen k
forall a. Int -> Gen a -> Gen a
resize Int
k Gen k
g1
                a
val <- Int -> Gen a -> Gen a
forall a. Int -> Gen a -> Gen a
resize Int
v Gen a
g2
                Map k a
rest <-
                    Int -> Gen (Map k a) -> Gen (Map k a)
forall a. Int -> Gen a -> Gen a
resize Int
m (Gen (Map k a) -> Gen (Map k a)) -> Gen (Map k a) -> Gen (Map k a)
forall a b. (a -> b) -> a -> b
$
                    Gen (k, a) -> Gen (Map k a)
forall k v. Ord k => Gen (k, v) -> Gen (Map k v)
genStructurallyValidMapOf (Gen (k, a) -> Gen (Map k a)) -> Gen (k, a) -> Gen (Map k a)
forall a b. (a -> b) -> a -> b
$
                    (,) (k -> a -> (k, a)) -> Gen k -> Gen (a -> (k, a))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen k
forall a. GenUnchecked a => Gen a
genUnchecked Gen (a -> (k, a)) -> Gen a -> Gen (k, a)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Gen a
forall a. GenUnchecked a => Gen a
genUnchecked
                Map k a -> Gen (Map k a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Map k a -> Gen (Map k a)) -> Map k a -> Gen (Map k a)
forall a b. (a -> b) -> a -> b
$ k -> a -> Map k a -> Map k a
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
key a
val Map k a
rest
        [Gen (Map k v)] -> Gen (Map k v)
forall a. [Gen a] -> Gen a
oneof [Gen k -> Gen v -> Gen (Map k v)
forall k a.
(Ord k, GenUnchecked k, GenUnchecked a) =>
Gen k -> Gen a -> Gen (Map k a)
go Gen k
forall a. GenInvalid a => Gen a
genInvalid Gen v
forall a. GenUnchecked a => Gen a
genUnchecked, Gen k -> Gen v -> Gen (Map k v)
forall k a.
(Ord k, GenUnchecked k, GenUnchecked a) =>
Gen k -> Gen a -> Gen (Map k a)
go Gen k
forall a. GenUnchecked a => Gen a
genUnchecked Gen v
forall a. GenInvalid a => Gen a
genInvalid]
#if MIN_VERSION_containers(0,5,9)
genStructurallyInvalidMap ::
       (Show k, Ord k, GenUnchecked k, GenUnchecked v) => Gen (Map k v)
genStructurallyInvalidMap :: Gen (Map k v)
genStructurallyInvalidMap = do
    Map k v
v <- Gen (Map k v)
forall a. GenUnchecked a => Gen a
genUnchecked
    if Map k v -> Bool
forall k a. Ord k => Map k a -> Bool
M.valid Map k v
v
        then (Int -> Int) -> Gen (Map k v) -> Gen (Map k v)
forall a. (Int -> Int) -> Gen a -> Gen a
scale (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ 1) Gen (Map k v)
forall k v.
(Show k, Ord k, GenUnchecked k, GenUnchecked v) =>
Gen (Map k v)
genStructurallyInvalidMap
        else Map k v -> Gen (Map k v)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Map k v
v
#endif