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 = forall x a. IntLikeMap x a
ILM.empty
{-# INLINE empty #-}

size :: IntLikeMultiMap k v -> Int
size :: forall k v. IntLikeMultiMap k v -> Int
size = 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 = 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 = forall x a.
Coercible x Int =>
(a -> a -> a) -> x -> a -> IntLikeMap x a -> IntLikeMap x a
ILM.insertWith forall x. IntLikeSet x -> IntLikeSet x -> IntLikeSet x
ILS.union k
k (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 = forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False (forall x. Coercible x Int => x -> IntLikeSet x -> Bool
ILS.member v
v) forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 = forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM forall {a} {b}.
Coercible a Int =>
IntLikeMap a b
-> (b, IntLikeSet a) -> Either (b, b, a) (IntLikeMap a b)
go1 forall x a. IntLikeMap x a
ILM.empty forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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) = forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (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 (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 forall x a. Coercible x Int => x -> IntLikeMap x a -> Maybe a
ILM.lookup c
v IntLikeMap c b
m of
      Maybe b
Nothing -> forall a b. b -> Either a b
Right (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' -> 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 = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' forall {a} {a}.
Coercible a Int =>
IntLikeMap a a -> (a, IntLikeSet a) -> IntLikeMap a a
go1 forall x a. IntLikeMap x a
ILM.empty forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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) = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (forall {x} {a}.
Coercible x Int =>
a -> IntLikeMap x a -> x -> IntLikeMap x a
go2 a
k) IntLikeMap a a
m (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 = 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 = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\IntLikeMultiMap v k
m (k
k, v
v) -> 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) forall k v. IntLikeMultiMap k v
empty forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall x a. Coercible x Int => IntLikeMap x a -> [(x, a)]
ILM.toList
{-# INLINE fromInvertedMap #-}