-- |
-- Copyright: © 2022–2025 Jonathan Knowles
-- License: Apache-2.0
--
-- A __lawful__ implementation of 'MultiMap', implemented in terms of 'Map' and
-- 'NESet'.
--
module Data.MonoidMap.Examples.MultiMap.Instances.MultiMap3 where

import Prelude

import Data.Map.Strict
    ( Map )
import Data.Maybe
    ( mapMaybe )
import Data.MonoidMap.Examples.Set.NonEmpty
    ( NESet )

import qualified Data.Map.Merge.Strict as Map
import qualified Data.Map.Strict as Map
import qualified Data.Set as Set
import qualified Data.MonoidMap.Examples.Set.NonEmpty as NESet
import qualified Data.MonoidMap.Examples.MultiMap.Class as Class

newtype MultiMap3 k v = MultiMap (Map k (NESet v))
    deriving stock (MultiMap3 k v -> MultiMap3 k v -> Bool
(MultiMap3 k v -> MultiMap3 k v -> Bool)
-> (MultiMap3 k v -> MultiMap3 k v -> Bool) -> Eq (MultiMap3 k v)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall k v. (Eq k, Eq v) => MultiMap3 k v -> MultiMap3 k v -> Bool
$c== :: forall k v. (Eq k, Eq v) => MultiMap3 k v -> MultiMap3 k v -> Bool
== :: MultiMap3 k v -> MultiMap3 k v -> Bool
$c/= :: forall k v. (Eq k, Eq v) => MultiMap3 k v -> MultiMap3 k v -> Bool
/= :: MultiMap3 k v -> MultiMap3 k v -> Bool
Eq, Int -> MultiMap3 k v -> ShowS
[MultiMap3 k v] -> ShowS
MultiMap3 k v -> String
(Int -> MultiMap3 k v -> ShowS)
-> (MultiMap3 k v -> String)
-> ([MultiMap3 k v] -> ShowS)
-> Show (MultiMap3 k v)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall k v. (Show k, Show v) => Int -> MultiMap3 k v -> ShowS
forall k v. (Show k, Show v) => [MultiMap3 k v] -> ShowS
forall k v. (Show k, Show v) => MultiMap3 k v -> String
$cshowsPrec :: forall k v. (Show k, Show v) => Int -> MultiMap3 k v -> ShowS
showsPrec :: Int -> MultiMap3 k v -> ShowS
$cshow :: forall k v. (Show k, Show v) => MultiMap3 k v -> String
show :: MultiMap3 k v -> String
$cshowList :: forall k v. (Show k, Show v) => [MultiMap3 k v] -> ShowS
showList :: [MultiMap3 k v] -> ShowS
Show)

instance (Ord k, Ord v) => Class.MultiMap MultiMap3 k v where

    fromList :: [(k, Set v)] -> MultiMap3 k v
fromList
        = Map k (NESet v) -> MultiMap3 k v
forall k v. Map k (NESet v) -> MultiMap3 k v
MultiMap
        (Map k (NESet v) -> MultiMap3 k v)
-> ([(k, Set v)] -> Map k (NESet v))
-> [(k, Set v)]
-> MultiMap3 k v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (NESet v -> NESet v -> NESet v)
-> [(k, NESet v)] -> Map k (NESet v)
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
Map.fromListWith NESet v -> NESet v -> NESet v
forall a. Semigroup a => a -> a -> a
(<>)
        ([(k, NESet v)] -> Map k (NESet v))
-> ([(k, Set v)] -> [(k, NESet v)])
-> [(k, Set v)]
-> Map k (NESet v)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((k, Set v) -> Maybe (k, NESet v))
-> [(k, Set v)] -> [(k, NESet v)]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe ((Set v -> Maybe (NESet v)) -> (k, Set v) -> Maybe (k, NESet v)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> (k, a) -> f (k, b)
traverse Set v -> Maybe (NESet v)
forall v. Set v -> Maybe (NESet v)
NESet.nonEmptySet)

    toList :: MultiMap3 k v -> [(k, Set v)]
toList (MultiMap Map k (NESet v)
m) = (NESet v -> Set v) -> (k, NESet v) -> (k, Set v)
forall a b. (a -> b) -> (k, a) -> (k, b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap NESet v -> Set v
forall v. NESet v -> Set v
NESet.toSet ((k, NESet v) -> (k, Set v)) -> [(k, NESet v)] -> [(k, Set v)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map k (NESet v) -> [(k, NESet v)]
forall k a. Map k a -> [(k, a)]
Map.toList Map k (NESet v)
m

    empty :: MultiMap3 k v
empty = Map k (NESet v) -> MultiMap3 k v
forall k v. Map k (NESet v) -> MultiMap3 k v
MultiMap Map k (NESet v)
forall k a. Map k a
Map.empty

    lookup :: k -> MultiMap3 k v -> Set v
lookup k
k (MultiMap Map k (NESet v)
m) = Set v -> (NESet v -> Set v) -> Maybe (NESet v) -> Set v
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Set v
forall a. Set a
Set.empty NESet v -> Set v
forall v. NESet v -> Set v
NESet.toSet (k -> Map k (NESet v) -> Maybe (NESet v)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup k
k Map k (NESet v)
m)

    null :: MultiMap3 k v -> Bool
null (MultiMap Map k (NESet v)
m) = Map k (NESet v) -> Bool
forall k a. Map k a -> Bool
Map.null Map k (NESet v)
m

    nonNull :: MultiMap3 k v -> Bool
nonNull (MultiMap Map k (NESet v)
m) = Bool -> Bool
not (Map k (NESet v) -> Bool
forall k a. Map k a -> Bool
Map.null Map k (NESet v)
m)

    nonNullKey :: k -> MultiMap3 k v -> Bool
nonNullKey k
k (MultiMap Map k (NESet v)
m) = k -> Map k (NESet v) -> Bool
forall k a. Ord k => k -> Map k a -> Bool
Map.member k
k Map k (NESet v)
m

    nonNullKeys :: MultiMap3 k v -> Set k
nonNullKeys (MultiMap Map k (NESet v)
m) = Map k (NESet v) -> Set k
forall k a. Map k a -> Set k
Map.keysSet Map k (NESet v)
m

    nonNullCount :: MultiMap3 k v -> Int
nonNullCount (MultiMap Map k (NESet v)
m) = Map k (NESet v) -> Int
forall k a. Map k a -> Int
Map.size Map k (NESet v)
m

    isSubmapOf :: MultiMap3 k v -> MultiMap3 k v -> Bool
isSubmapOf (MultiMap Map k (NESet v)
m1) (MultiMap Map k (NESet v)
m2) =
        (NESet v -> NESet v -> Bool)
-> Map k (NESet v) -> Map k (NESet v) -> Bool
forall k a b.
Ord k =>
(a -> b -> Bool) -> Map k a -> Map k b -> Bool
Map.isSubmapOfBy NESet v -> NESet v -> Bool
forall v. Ord v => NESet v -> NESet v -> Bool
NESet.isSubsetOf Map k (NESet v)
m1 Map k (NESet v)
m2

    update :: k -> Set v -> MultiMap3 k v -> MultiMap3 k v
update k
k Set v
vs (MultiMap Map k (NESet v)
m) =
        case Set v -> Maybe (NESet v)
forall v. Set v -> Maybe (NESet v)
NESet.nonEmptySet Set v
vs of
            Maybe (NESet v)
Nothing -> Map k (NESet v) -> MultiMap3 k v
forall k v. Map k (NESet v) -> MultiMap3 k v
MultiMap (k -> Map k (NESet v) -> Map k (NESet v)
forall k a. Ord k => k -> Map k a -> Map k a
Map.delete k
k    Map k (NESet v)
m)
            Just NESet v
ys -> Map k (NESet v) -> MultiMap3 k v
forall k v. Map k (NESet v) -> MultiMap3 k v
MultiMap (k -> NESet v -> Map k (NESet v) -> Map k (NESet v)
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert k
k NESet v
ys Map k (NESet v)
m)

    insert :: k -> Set v -> MultiMap3 k v -> MultiMap3 k v
insert k
k Set v
vs (MultiMap Map k (NESet v)
m) =
        case Set v -> Maybe (NESet v)
forall v. Set v -> Maybe (NESet v)
NESet.nonEmptySet Set v
xs of
            Maybe (NESet v)
Nothing -> Map k (NESet v) -> MultiMap3 k v
forall k v. Map k (NESet v) -> MultiMap3 k v
MultiMap (k -> Map k (NESet v) -> Map k (NESet v)
forall k a. Ord k => k -> Map k a -> Map k a
Map.delete k
k    Map k (NESet v)
m)
            Just NESet v
ys -> Map k (NESet v) -> MultiMap3 k v
forall k v. Map k (NESet v) -> MultiMap3 k v
MultiMap (k -> NESet v -> Map k (NESet v) -> Map k (NESet v)
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert k
k NESet v
ys Map k (NESet v)
m)
      where
        xs :: Set v
xs = Set v -> (NESet v -> Set v) -> Maybe (NESet v) -> Set v
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Set v
forall a. Set a
Set.empty NESet v -> Set v
forall v. NESet v -> Set v
NESet.toSet (k -> Map k (NESet v) -> Maybe (NESet v)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup k
k Map k (NESet v)
m) Set v -> Set v -> Set v
forall a. Ord a => Set a -> Set a -> Set a
`Set.union` Set v
vs

    remove :: k -> Set v -> MultiMap3 k v -> MultiMap3 k v
remove k
k Set v
vs (MultiMap Map k (NESet v)
m) =
        case Set v -> Maybe (NESet v)
forall v. Set v -> Maybe (NESet v)
NESet.nonEmptySet Set v
xs of
            Maybe (NESet v)
Nothing -> Map k (NESet v) -> MultiMap3 k v
forall k v. Map k (NESet v) -> MultiMap3 k v
MultiMap (k -> Map k (NESet v) -> Map k (NESet v)
forall k a. Ord k => k -> Map k a -> Map k a
Map.delete k
k    Map k (NESet v)
m)
            Just NESet v
ys -> Map k (NESet v) -> MultiMap3 k v
forall k v. Map k (NESet v) -> MultiMap3 k v
MultiMap (k -> NESet v -> Map k (NESet v) -> Map k (NESet v)
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert k
k NESet v
ys Map k (NESet v)
m)
      where
        xs :: Set v
xs = Set v -> (NESet v -> Set v) -> Maybe (NESet v) -> Set v
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Set v
forall a. Set a
Set.empty NESet v -> Set v
forall v. NESet v -> Set v
NESet.toSet (k -> Map k (NESet v) -> Maybe (NESet v)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup k
k Map k (NESet v)
m) Set v -> Set v -> Set v
forall a. Ord a => Set a -> Set a -> Set a
`Set.difference` Set v
vs

    union :: MultiMap3 k v -> MultiMap3 k v -> MultiMap3 k v
union (MultiMap Map k (NESet v)
m1) (MultiMap Map k (NESet v)
m2) = Map k (NESet v) -> MultiMap3 k v
forall k v. Map k (NESet v) -> MultiMap3 k v
MultiMap (Map k (NESet v) -> MultiMap3 k v)
-> Map k (NESet v) -> MultiMap3 k v
forall a b. (a -> b) -> a -> b
$
        (NESet v -> NESet v -> NESet v)
-> Map k (NESet v) -> Map k (NESet v) -> Map k (NESet v)
forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
Map.unionWith NESet v -> NESet v -> NESet v
forall v. Ord v => NESet v -> NESet v -> NESet v
NESet.union Map k (NESet v)
m1 Map k (NESet v)
m2

    intersection :: MultiMap3 k v -> MultiMap3 k v -> MultiMap3 k v
intersection (MultiMap Map k (NESet v)
m1) (MultiMap Map k (NESet v)
m2) = Map k (NESet v) -> MultiMap3 k v
forall k v. Map k (NESet v) -> MultiMap3 k v
MultiMap (Map k (NESet v) -> MultiMap3 k v)
-> Map k (NESet v) -> MultiMap3 k v
forall a b. (a -> b) -> a -> b
$
        SimpleWhenMissing k (NESet v) (NESet v)
-> SimpleWhenMissing k (NESet v) (NESet v)
-> SimpleWhenMatched k (NESet v) (NESet v) (NESet v)
-> Map k (NESet v)
-> Map k (NESet v)
-> Map k (NESet v)
forall k a c b.
Ord k =>
SimpleWhenMissing k a c
-> SimpleWhenMissing k b c
-> SimpleWhenMatched k a b c
-> Map k a
-> Map k b
-> Map k c
Map.merge
            SimpleWhenMissing k (NESet v) (NESet v)
forall (f :: * -> *) k x y. Applicative f => WhenMissing f k x y
Map.dropMissing
            SimpleWhenMissing k (NESet v) (NESet v)
forall (f :: * -> *) k x y. Applicative f => WhenMissing f k x y
Map.dropMissing
            ((k -> NESet v -> NESet v -> Maybe (NESet v))
-> SimpleWhenMatched k (NESet v) (NESet v) (NESet v)
forall (f :: * -> *) k x y z.
Applicative f =>
(k -> x -> y -> Maybe z) -> WhenMatched f k x y z
Map.zipWithMaybeMatched k -> NESet v -> NESet v -> Maybe (NESet v)
Ord v => k -> NESet v -> NESet v -> Maybe (NESet v)
mergeValues)
            Map k (NESet v)
m1
            Map k (NESet v)
m2
      where
        mergeValues :: Ord v => k -> NESet v -> NESet v -> Maybe (NESet v)
        mergeValues :: Ord v => k -> NESet v -> NESet v -> Maybe (NESet v)
mergeValues k
_k NESet v
s1 NESet v
s2 = Set v -> Maybe (NESet v)
forall v. Set v -> Maybe (NESet v)
NESet.nonEmptySet (NESet v -> NESet v -> Set v
forall v. Ord v => NESet v -> NESet v -> Set v
NESet.intersection NESet v
s1 NESet v
s2)