module IntLike.MultiMap
  ( IntLikeMultiMap
  , empty
  , size
  , toList
  , insert
  , member
  , invertDisjoint
  , unsafeInvertDisjoint
  , fromInvertedMap
  )
where

import Control.Monad (foldM)
import Data.Coerce (Coercible)
import Data.Foldable (foldl')
import IntLike.Map (IntLikeMap)
import qualified IntLike.Map as ILM
import IntLike.Set (IntLikeSet)
import qualified IntLike.Set as ILS

type IntLikeMultiMap k v = IntLikeMap k (IntLikeSet v)

empty :: IntLikeMultiMap k v
empty :: forall k v. IntLikeMultiMap k v
empty = IntLikeMap k (IntLikeSet v)
forall x a. IntLikeMap x a
ILM.empty
{-# INLINE empty #-}

size :: IntLikeMultiMap k v -> Int
size :: forall k v. IntLikeMultiMap k v -> Int
size = IntLikeMap k (IntLikeSet v) -> Int
forall x a. IntLikeMap x a -> Int
ILM.size
{-# INLINE size #-}

toList :: Coercible k Int => IntLikeMultiMap k v -> [(k, IntLikeSet v)]
toList :: forall k v.
Coercible k Int =>
IntLikeMultiMap k v -> [(k, IntLikeSet v)]
toList = IntLikeMap k (IntLikeSet v) -> [(k, IntLikeSet v)]
forall x a. Coercible x Int => IntLikeMap x a -> [(x, a)]
ILM.toList
{-# INLINE toList #-}

insert :: (Coercible k Int, Coercible v Int) => k -> v -> IntLikeMultiMap k v -> IntLikeMultiMap k v
insert :: forall k v.
(Coercible k Int, Coercible v Int) =>
k -> v -> IntLikeMultiMap k v -> IntLikeMultiMap k v
insert k
k v
v = (IntLikeSet v -> IntLikeSet v -> IntLikeSet v)
-> k
-> IntLikeSet v
-> IntLikeMap k (IntLikeSet v)
-> IntLikeMap k (IntLikeSet v)
forall x a.
Coercible x Int =>
(a -> a -> a) -> x -> a -> IntLikeMap x a -> IntLikeMap x a
ILM.insertWith IntLikeSet v -> IntLikeSet v -> IntLikeSet v
forall x. IntLikeSet x -> IntLikeSet x -> IntLikeSet x
ILS.union k
k (v -> IntLikeSet v
forall x. Coercible x Int => x -> IntLikeSet x
ILS.singleton v
v)
{-# INLINE insert #-}

member :: (Coercible k Int, Coercible v Int) => k -> v -> IntLikeMultiMap k v -> Bool
member :: forall k v.
(Coercible k Int, Coercible v Int) =>
k -> v -> IntLikeMultiMap k v -> Bool
member k
k v
v = Bool -> (IntLikeSet v -> Bool) -> Maybe (IntLikeSet v) -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False (v -> IntLikeSet v -> Bool
forall x. Coercible x Int => x -> IntLikeSet x -> Bool
ILS.member v
v) (Maybe (IntLikeSet v) -> Bool)
-> (IntLikeMultiMap k v -> Maybe (IntLikeSet v))
-> IntLikeMultiMap k v
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. k -> IntLikeMultiMap k v -> Maybe (IntLikeSet v)
forall x a. Coercible x Int => x -> IntLikeMap x a -> Maybe a
ILM.lookup k
k
{-# INLINE member #-}

invertDisjoint :: (Coercible k Int, Coercible v Int) => IntLikeMultiMap k v -> Either (k, k, v) (IntLikeMap v k)
invertDisjoint :: forall k v.
(Coercible k Int, Coercible v Int) =>
IntLikeMultiMap k v -> Either (k, k, v) (IntLikeMap v k)
invertDisjoint = (IntLikeMap v k
 -> (k, IntLikeSet v) -> Either (k, k, v) (IntLikeMap v k))
-> IntLikeMap v k
-> [(k, IntLikeSet v)]
-> Either (k, k, v) (IntLikeMap v k)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM IntLikeMap v k
-> (k, IntLikeSet v) -> Either (k, k, v) (IntLikeMap v k)
forall {a} {b}.
Coercible a Int =>
IntLikeMap a b
-> (b, IntLikeSet a) -> Either (b, b, a) (IntLikeMap a b)
go1 IntLikeMap v k
forall x a. IntLikeMap x a
ILM.empty ([(k, IntLikeSet v)] -> Either (k, k, v) (IntLikeMap v k))
-> (IntLikeMultiMap k v -> [(k, IntLikeSet v)])
-> IntLikeMultiMap k v
-> Either (k, k, v) (IntLikeMap v k)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntLikeMultiMap k v -> [(k, IntLikeSet v)]
forall x a. Coercible x Int => IntLikeMap x a -> [(x, a)]
ILM.toList
 where
  go1 :: IntLikeMap a b
-> (b, IntLikeSet a) -> Either (b, b, a) (IntLikeMap a b)
go1 IntLikeMap a b
m (b
k, IntLikeSet a
vs) = (IntLikeMap a b -> a -> Either (b, b, a) (IntLikeMap a b))
-> IntLikeMap a b -> [a] -> Either (b, b, a) (IntLikeMap a b)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (b -> IntLikeMap a b -> a -> Either (b, b, a) (IntLikeMap a b)
forall {c} {b}.
Coercible c Int =>
b -> IntLikeMap c b -> c -> Either (b, b, c) (IntLikeMap c b)
go2 b
k) IntLikeMap a b
m (IntLikeSet a -> [a]
forall x. Coercible x Int => IntLikeSet x -> [x]
ILS.toList IntLikeSet a
vs)
  go2 :: b -> IntLikeMap c b -> c -> Either (b, b, c) (IntLikeMap c b)
go2 b
k IntLikeMap c b
m c
v =
    case c -> IntLikeMap c b -> Maybe b
forall x a. Coercible x Int => x -> IntLikeMap x a -> Maybe a
ILM.lookup c
v IntLikeMap c b
m of
      Maybe b
Nothing -> IntLikeMap c b -> Either (b, b, c) (IntLikeMap c b)
forall a b. b -> Either a b
Right (c -> b -> IntLikeMap c b -> IntLikeMap c b
forall x a.
Coercible x Int =>
x -> a -> IntLikeMap x a -> IntLikeMap x a
ILM.insert c
v b
k IntLikeMap c b
m)
      Just b
k' -> (b, b, c) -> Either (b, b, c) (IntLikeMap c b)
forall a b. a -> Either a b
Left (b
k, b
k', c
v)

unsafeInvertDisjoint :: (Coercible k Int, Coercible v Int) => IntLikeMultiMap k v -> IntLikeMap v k
unsafeInvertDisjoint :: forall k v.
(Coercible k Int, Coercible v Int) =>
IntLikeMultiMap k v -> IntLikeMap v k
unsafeInvertDisjoint = (IntLikeMap v k -> (k, IntLikeSet v) -> IntLikeMap v k)
-> IntLikeMap v k -> [(k, IntLikeSet v)] -> IntLikeMap v k
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' IntLikeMap v k -> (k, IntLikeSet v) -> IntLikeMap v k
forall {a} {a}.
Coercible a Int =>
IntLikeMap a a -> (a, IntLikeSet a) -> IntLikeMap a a
go1 IntLikeMap v k
forall x a. IntLikeMap x a
ILM.empty ([(k, IntLikeSet v)] -> IntLikeMap v k)
-> (IntLikeMultiMap k v -> [(k, IntLikeSet v)])
-> IntLikeMultiMap k v
-> IntLikeMap v k
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntLikeMultiMap k v -> [(k, IntLikeSet v)]
forall x a. Coercible x Int => IntLikeMap x a -> [(x, a)]
ILM.toList
 where
  go1 :: IntLikeMap a a -> (a, IntLikeSet a) -> IntLikeMap a a
go1 IntLikeMap a a
m (a
k, IntLikeSet a
vs) = (IntLikeMap a a -> a -> IntLikeMap a a)
-> IntLikeMap a a -> [a] -> IntLikeMap a a
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (a -> IntLikeMap a a -> a -> IntLikeMap a a
forall {x} {a}.
Coercible x Int =>
a -> IntLikeMap x a -> x -> IntLikeMap x a
go2 a
k) IntLikeMap a a
m (IntLikeSet a -> [a]
forall x. Coercible x Int => IntLikeSet x -> [x]
ILS.toList IntLikeSet a
vs)
  go2 :: a -> IntLikeMap x a -> x -> IntLikeMap x a
go2 a
k IntLikeMap x a
m x
v = x -> a -> IntLikeMap x a -> IntLikeMap x a
forall x a.
Coercible x Int =>
x -> a -> IntLikeMap x a -> IntLikeMap x a
ILM.insert x
v a
k IntLikeMap x a
m

fromInvertedMap :: (Coercible k Int, Coercible v Int) => IntLikeMap k v -> IntLikeMultiMap v k
fromInvertedMap :: forall k v.
(Coercible k Int, Coercible v Int) =>
IntLikeMap k v -> IntLikeMultiMap v k
fromInvertedMap = (IntLikeMultiMap v k -> (k, v) -> IntLikeMultiMap v k)
-> IntLikeMultiMap v k -> [(k, v)] -> IntLikeMultiMap v k
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\IntLikeMultiMap v k
m (k
k, v
v) -> v -> k -> IntLikeMultiMap v k -> IntLikeMultiMap v k
forall k v.
(Coercible k Int, Coercible v Int) =>
k -> v -> IntLikeMultiMap k v -> IntLikeMultiMap k v
insert v
v k
k IntLikeMultiMap v k
m) IntLikeMultiMap v k
forall k v. IntLikeMultiMap k v
empty ([(k, v)] -> IntLikeMultiMap v k)
-> (IntLikeMap k v -> [(k, v)])
-> IntLikeMap k v
-> IntLikeMultiMap v k
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntLikeMap k v -> [(k, v)]
forall x a. Coercible x Int => IntLikeMap x a -> [(x, a)]
ILM.toList
{-# INLINE fromInvertedMap #-}