{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE DeriveGeneric #-}
module Data.BimapMany
(
-- * BimapMany type
BimapMany
-- * Construction
, empty
, singleton
, fromMap
, fromSet
-- ** From unordered lists
, fromList
-- * Insertion
, insert
-- * Deletion/Update
, delete
, deleteL
, deleteR
-- * Query
-- ** Lookup
, lookup
, lookupL
, lookupR
, lookupL'
, lookupR'
-- ** Size
, null
, size
, sizeL
, sizeR
-- * Combine
, union
-- * Conversion
-- ** Maps
, toMap
-- ** Lists
, toList
-- * Debugging
, valid
) where

import Prelude hiding (abs, lookup, null)

import Data.Map.Signature (Map)
import qualified Data.Map.Signature as M
import qualified Data.Map.Strict as MS
import Data.Set (Set)
import qualified Data.Set as S

import Control.DeepSeq (NFData)
import GHC.Generics (Generic)

import Data.Function (on)
import Data.List (foldl', groupBy, sort)
import Data.Maybe (fromMaybe)

-- * BimapMany type
-------------------

-- NOTE/TODO: considering different internal representations
--
-- The current implementation uses a Map (a, b) c, but it's far from being the
-- only solution.
--
-- (Map a (Map b c)) (Map b (Set a)
-- It would take less space, having less set-like structures
-- It's asymmetric, which would make a few operations more expensive:
-- * lookupL (but lookupL' could be used instead)
-- Other operations are more expensive due to the lack of Map (a, b) c
-- * toMap
-- * size (but it could be stored separately)
--
-- (Map a (Map b c)) (Map b (Map a c)
-- It would take less space, having less set-like structures
-- Some operations are more expensive due to the lack of Map (a, b) c
-- * toMap
-- * size (but it could be stored separately)
-- It's symetric but it "duplicates" c. On one hand it gets shared anyway, but
-- on the other hand the less used map could accumulate thunks if care is not
-- taken.
--
-- Without analyzing this more in depth, the current option looks safest.

-- NOTE: operations on l and r are always strict
-- When operating on l and r, the Map value is another set structure, not a
-- BimapMany value (c), so operations should be strict.

data BimapMany a b c = BimapMany
  !(MS.Map a (Set b)) -- l
  !(MS.Map b (Set a)) -- r
  !(Map (a, b) c) -- m
  -- Invariants (checked by 'valid'):
  --   (a, bs) ∈ l, b ∈ bs,   => (a, b) ∈ m, (b, as) ∈ r, a ∈ as
  --   (b, as) ∈ r, a ∈ as,   => (a, b) ∈ m, (a, bs) ∈ l, b ∈ bs
  --   ((a, b), _) ∈ m        => (a, bs) ∈ l, (b, as) ∈ r, a ∈ as, b ∈ bs
  -- TODO use proper lazy/strict functions for these instances
  -- MAYBE move the data definition to a separate module and share it, like Map
  -- Or better, provide coercion functions
  deriving (a -> BimapMany a b b -> BimapMany a b a
(a -> b) -> BimapMany a b a -> BimapMany a b b
(forall a b. (a -> b) -> BimapMany a b a -> BimapMany a b b)
-> (forall a b. a -> BimapMany a b b -> BimapMany a b a)
-> Functor (BimapMany a b)
forall a b. a -> BimapMany a b b -> BimapMany a b a
forall a b. (a -> b) -> BimapMany a b a -> BimapMany a b b
forall a b a b. a -> BimapMany a b b -> BimapMany a b a
forall a b a b. (a -> b) -> BimapMany a b a -> BimapMany a b b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: a -> BimapMany a b b -> BimapMany a b a
$c<$ :: forall a b a b. a -> BimapMany a b b -> BimapMany a b a
fmap :: (a -> b) -> BimapMany a b a -> BimapMany a b b
$cfmap :: forall a b a b. (a -> b) -> BimapMany a b a -> BimapMany a b b
Functor, BimapMany a b a -> Bool
(a -> m) -> BimapMany a b a -> m
(a -> b -> b) -> b -> BimapMany a b a -> b
(forall m. Monoid m => BimapMany a b m -> m)
-> (forall m a. Monoid m => (a -> m) -> BimapMany a b a -> m)
-> (forall m a. Monoid m => (a -> m) -> BimapMany a b a -> m)
-> (forall a b. (a -> b -> b) -> b -> BimapMany a b a -> b)
-> (forall a b. (a -> b -> b) -> b -> BimapMany a b a -> b)
-> (forall b a. (b -> a -> b) -> b -> BimapMany a b a -> b)
-> (forall b a. (b -> a -> b) -> b -> BimapMany a b a -> b)
-> (forall a. (a -> a -> a) -> BimapMany a b a -> a)
-> (forall a. (a -> a -> a) -> BimapMany a b a -> a)
-> (forall a. BimapMany a b a -> [a])
-> (forall a. BimapMany a b a -> Bool)
-> (forall a. BimapMany a b a -> Int)
-> (forall a. Eq a => a -> BimapMany a b a -> Bool)
-> (forall a. Ord a => BimapMany a b a -> a)
-> (forall a. Ord a => BimapMany a b a -> a)
-> (forall a. Num a => BimapMany a b a -> a)
-> (forall a. Num a => BimapMany a b a -> a)
-> Foldable (BimapMany a b)
forall a. Eq a => a -> BimapMany a b a -> Bool
forall a. Num a => BimapMany a b a -> a
forall a. Ord a => BimapMany a b a -> a
forall m. Monoid m => BimapMany a b m -> m
forall a. BimapMany a b a -> Bool
forall a. BimapMany a b a -> Int
forall a. BimapMany a b a -> [a]
forall a. (a -> a -> a) -> BimapMany a b a -> a
forall m a. Monoid m => (a -> m) -> BimapMany a b a -> m
forall b a. (b -> a -> b) -> b -> BimapMany a b a -> b
forall a b. (a -> b -> b) -> b -> BimapMany a b a -> b
forall a b a. Eq a => a -> BimapMany a b a -> Bool
forall a b a. Num a => BimapMany a b a -> a
forall a b a. Ord a => BimapMany a b a -> a
forall a b m. Monoid m => BimapMany a b m -> m
forall a b a. BimapMany a b a -> Bool
forall a b a. BimapMany a b a -> Int
forall a b a. BimapMany a b a -> [a]
forall a b a. (a -> a -> a) -> BimapMany a b a -> a
forall a b m a. Monoid m => (a -> m) -> BimapMany a b a -> m
forall a b b a. (b -> a -> b) -> b -> BimapMany a b a -> b
forall a b a b. (a -> b -> b) -> b -> BimapMany a b a -> b
forall (t :: * -> *).
(forall m. Monoid m => t m -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. t a -> [a])
-> (forall a. t a -> Bool)
-> (forall a. t a -> Int)
-> (forall a. Eq a => a -> t a -> Bool)
-> (forall a. Ord a => t a -> a)
-> (forall a. Ord a => t a -> a)
-> (forall a. Num a => t a -> a)
-> (forall a. Num a => t a -> a)
-> Foldable t
product :: BimapMany a b a -> a
$cproduct :: forall a b a. Num a => BimapMany a b a -> a
sum :: BimapMany a b a -> a
$csum :: forall a b a. Num a => BimapMany a b a -> a
minimum :: BimapMany a b a -> a
$cminimum :: forall a b a. Ord a => BimapMany a b a -> a
maximum :: BimapMany a b a -> a
$cmaximum :: forall a b a. Ord a => BimapMany a b a -> a
elem :: a -> BimapMany a b a -> Bool
$celem :: forall a b a. Eq a => a -> BimapMany a b a -> Bool
length :: BimapMany a b a -> Int
$clength :: forall a b a. BimapMany a b a -> Int
null :: BimapMany a b a -> Bool
$cnull :: forall a b a. BimapMany a b a -> Bool
toList :: BimapMany a b a -> [a]
$ctoList :: forall a b a. BimapMany a b a -> [a]
foldl1 :: (a -> a -> a) -> BimapMany a b a -> a
$cfoldl1 :: forall a b a. (a -> a -> a) -> BimapMany a b a -> a
foldr1 :: (a -> a -> a) -> BimapMany a b a -> a
$cfoldr1 :: forall a b a. (a -> a -> a) -> BimapMany a b a -> a
foldl' :: (b -> a -> b) -> b -> BimapMany a b a -> b
$cfoldl' :: forall a b b a. (b -> a -> b) -> b -> BimapMany a b a -> b
foldl :: (b -> a -> b) -> b -> BimapMany a b a -> b
$cfoldl :: forall a b b a. (b -> a -> b) -> b -> BimapMany a b a -> b
foldr' :: (a -> b -> b) -> b -> BimapMany a b a -> b
$cfoldr' :: forall a b a b. (a -> b -> b) -> b -> BimapMany a b a -> b
foldr :: (a -> b -> b) -> b -> BimapMany a b a -> b
$cfoldr :: forall a b a b. (a -> b -> b) -> b -> BimapMany a b a -> b
foldMap' :: (a -> m) -> BimapMany a b a -> m
$cfoldMap' :: forall a b m a. Monoid m => (a -> m) -> BimapMany a b a -> m
foldMap :: (a -> m) -> BimapMany a b a -> m
$cfoldMap :: forall a b m a. Monoid m => (a -> m) -> BimapMany a b a -> m
fold :: BimapMany a b m -> m
$cfold :: forall a b m. Monoid m => BimapMany a b m -> m
Foldable, Functor (BimapMany a b)
Foldable (BimapMany a b)
Functor (BimapMany a b)
-> Foldable (BimapMany a b)
-> (forall (f :: * -> *) a b.
    Applicative f =>
    (a -> f b) -> BimapMany a b a -> f (BimapMany a b b))
-> (forall (f :: * -> *) a.
    Applicative f =>
    BimapMany a b (f a) -> f (BimapMany a b a))
-> (forall (m :: * -> *) a b.
    Monad m =>
    (a -> m b) -> BimapMany a b a -> m (BimapMany a b b))
-> (forall (m :: * -> *) a.
    Monad m =>
    BimapMany a b (m a) -> m (BimapMany a b a))
-> Traversable (BimapMany a b)
(a -> f b) -> BimapMany a b a -> f (BimapMany a b b)
forall a b. Functor (BimapMany a b)
forall a b. Foldable (BimapMany a b)
forall a b (m :: * -> *) a.
Monad m =>
BimapMany a b (m a) -> m (BimapMany a b a)
forall a b (f :: * -> *) a.
Applicative f =>
BimapMany a b (f a) -> f (BimapMany a b a)
forall a b (m :: * -> *) a b.
Monad m =>
(a -> m b) -> BimapMany a b a -> m (BimapMany a b b)
forall a b (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> BimapMany a b a -> f (BimapMany a b b)
forall (t :: * -> *).
Functor t
-> Foldable t
-> (forall (f :: * -> *) a b.
    Applicative f =>
    (a -> f b) -> t a -> f (t b))
-> (forall (f :: * -> *) a. Applicative f => t (f a) -> f (t a))
-> (forall (m :: * -> *) a b.
    Monad m =>
    (a -> m b) -> t a -> m (t b))
-> (forall (m :: * -> *) a. Monad m => t (m a) -> m (t a))
-> Traversable t
forall (m :: * -> *) a.
Monad m =>
BimapMany a b (m a) -> m (BimapMany a b a)
forall (f :: * -> *) a.
Applicative f =>
BimapMany a b (f a) -> f (BimapMany a b a)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> BimapMany a b a -> m (BimapMany a b b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> BimapMany a b a -> f (BimapMany a b b)
sequence :: BimapMany a b (m a) -> m (BimapMany a b a)
$csequence :: forall a b (m :: * -> *) a.
Monad m =>
BimapMany a b (m a) -> m (BimapMany a b a)
mapM :: (a -> m b) -> BimapMany a b a -> m (BimapMany a b b)
$cmapM :: forall a b (m :: * -> *) a b.
Monad m =>
(a -> m b) -> BimapMany a b a -> m (BimapMany a b b)
sequenceA :: BimapMany a b (f a) -> f (BimapMany a b a)
$csequenceA :: forall a b (f :: * -> *) a.
Applicative f =>
BimapMany a b (f a) -> f (BimapMany a b a)
traverse :: (a -> f b) -> BimapMany a b a -> f (BimapMany a b b)
$ctraverse :: forall a b (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> BimapMany a b a -> f (BimapMany a b b)
$cp2Traversable :: forall a b. Foldable (BimapMany a b)
$cp1Traversable :: forall a b. Functor (BimapMany a b)
Traversable, (forall x. BimapMany a b c -> Rep (BimapMany a b c) x)
-> (forall x. Rep (BimapMany a b c) x -> BimapMany a b c)
-> Generic (BimapMany a b c)
forall x. Rep (BimapMany a b c) x -> BimapMany a b c
forall x. BimapMany a b c -> Rep (BimapMany a b c) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall a b c x. Rep (BimapMany a b c) x -> BimapMany a b c
forall a b c x. BimapMany a b c -> Rep (BimapMany a b c) x
$cto :: forall a b c x. Rep (BimapMany a b c) x -> BimapMany a b c
$cfrom :: forall a b c x. BimapMany a b c -> Rep (BimapMany a b c) x
Generic)

instance (Show a, Show b, Show c) => Show (BimapMany a b c) where
  show :: BimapMany a b c -> String
show BimapMany a b c
x = String
"fromList " String -> ShowS
forall a. [a] -> [a] -> [a]
++ [(a, b, c)] -> String
forall a. Show a => a -> String
show (BimapMany a b c -> [(a, b, c)]
forall a b c. BimapMany a b c -> [(a, b, c)]
toList BimapMany a b c
x)

instance (Eq a, Eq b, Eq c) => Eq (BimapMany a b c) where
  {-# INLINABLE (==) #-}
  == :: BimapMany a b c -> BimapMany a b c -> Bool
(==) = Map (a, b) c -> Map (a, b) c -> Bool
forall a. Eq a => a -> a -> Bool
(==) (Map (a, b) c -> Map (a, b) c -> Bool)
-> (BimapMany a b c -> Map (a, b) c)
-> BimapMany a b c
-> BimapMany a b c
-> Bool
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` BimapMany a b c -> Map (a, b) c
forall a b c. BimapMany a b c -> Map (a, b) c
toMap

instance (Ord a, Ord b, Ord c) => Ord (BimapMany a b c) where
  {-# INLINABLE compare #-}
  compare :: BimapMany a b c -> BimapMany a b c -> Ordering
compare = Map (a, b) c -> Map (a, b) c -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (Map (a, b) c -> Map (a, b) c -> Ordering)
-> (BimapMany a b c -> Map (a, b) c)
-> BimapMany a b c
-> BimapMany a b c
-> Ordering
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` BimapMany a b c -> Map (a, b) c
forall a b c. BimapMany a b c -> Map (a, b) c
toMap

instance (Ord a, Ord b) => Semigroup (BimapMany a b c) where
  {-# INLINABLE (<>) #-}
  <> :: BimapMany a b c -> BimapMany a b c -> BimapMany a b c
(<>) = BimapMany a b c -> BimapMany a b c -> BimapMany a b c
forall a b c.
(Ord a, Ord b) =>
BimapMany a b c -> BimapMany a b c -> BimapMany a b c
union

instance (Ord a, Ord b) => Monoid (BimapMany a b c) where
  {-# INLINABLE mempty #-}
  mempty :: BimapMany a b c
mempty = BimapMany a b c
forall a b c. BimapMany a b c
empty

instance (NFData a, NFData b, NFData c) => NFData (BimapMany a b c)

-- * Construction
-----------------

{-# INLINABLE empty #-}
empty :: BimapMany a b c
empty :: BimapMany a b c
empty = Map a (Set b) -> Map b (Set a) -> Map (a, b) c -> BimapMany a b c
forall a b c.
Map a (Set b) -> Map b (Set a) -> Map (a, b) c -> BimapMany a b c
BimapMany Map a (Set b)
forall k a. Map k a
MS.empty Map b (Set a)
forall k a. Map k a
MS.empty Map (a, b) c
forall k a. Map k a
M.empty

{-# INLINABLE singleton #-}
singleton :: a -> b -> c -> BimapMany a b c
singleton :: a -> b -> c -> BimapMany a b c
singleton a
a b
b c
c = Map a (Set b) -> Map b (Set a) -> Map (a, b) c -> BimapMany a b c
forall a b c.
Map a (Set b) -> Map b (Set a) -> Map (a, b) c -> BimapMany a b c
BimapMany Map a (Set b)
l Map b (Set a)
r Map (a, b) c
m
  where
    l :: Map a (Set b)
l = a -> Set b -> Map a (Set b)
forall k a. k -> a -> Map k a
MS.singleton a
a (Set b -> Map a (Set b)) -> Set b -> Map a (Set b)
forall a b. (a -> b) -> a -> b
$ b -> Set b
forall a. a -> Set a
S.singleton b
b
    r :: Map b (Set a)
r = b -> Set a -> Map b (Set a)
forall k a. k -> a -> Map k a
MS.singleton b
b (Set a -> Map b (Set a)) -> Set a -> Map b (Set a)
forall a b. (a -> b) -> a -> b
$ a -> Set a
forall a. a -> Set a
S.singleton a
a
    m :: Map (a, b) c
m = (a, b) -> c -> Map (a, b) c
forall k a. k -> a -> Map k a
M.singleton (a
a, b
b) c
c

-- TODO write a warning about strict/lazy map that must correspond to BimapMany
{-# INLINABLE fromMap #-}
fromMap :: (Ord a, Ord b) => Map (a, b) c -> BimapMany a b c
fromMap :: Map (a, b) c -> BimapMany a b c
fromMap Map (a, b) c
m = Map a (Set b) -> Map b (Set a) -> Map (a, b) c -> BimapMany a b c
forall a b c.
Map a (Set b) -> Map b (Set a) -> Map (a, b) c -> BimapMany a b c
BimapMany Map a (Set b)
l Map b (Set a)
r Map (a, b) c
m
  where
    abs :: [(a, b)]
abs = Map (a, b) c -> [(a, b)]
forall k a. Map k a -> [k]
M.keys Map (a, b) c
m
    l :: Map a (Set b)
l = [(a, b)] -> Map a (Set b)
forall a b. (Ord a, Ord b) => [(a, b)] -> Map a (Set b)
ascListToMapSet [(a, b)]
abs
    bas :: [(b, a)]
bas = [(b, a)] -> [(b, a)]
forall a. Ord a => [a] -> [a]
sort ([(b, a)] -> [(b, a)]) -> [(b, a)] -> [(b, a)]
forall a b. (a -> b) -> a -> b
$ (\(a
a, b
b) -> (b
b, a
a)) ((a, b) -> (b, a)) -> [(a, b)] -> [(b, a)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(a, b)]
abs
    r :: Map b (Set a)
r = [(b, a)] -> Map b (Set a)
forall a b. (Ord a, Ord b) => [(a, b)] -> Map a (Set b)
ascListToMapSet [(b, a)]
bas

{-# INLINABLE fromSet #-}
fromSet :: (Ord a, Ord b) => (a -> b -> c) -> Set (a, b) -> BimapMany a b c
fromSet :: (a -> b -> c) -> Set (a, b) -> BimapMany a b c
fromSet a -> b -> c
f Set (a, b)
s = Map a (Set b) -> Map b (Set a) -> Map (a, b) c -> BimapMany a b c
forall a b c.
Map a (Set b) -> Map b (Set a) -> Map (a, b) c -> BimapMany a b c
BimapMany Map a (Set b)
l Map b (Set a)
r Map (a, b) c
m
  where
    abs :: [(a, b)]
abs = Set (a, b) -> [(a, b)]
forall a. Set a -> [a]
S.toAscList Set (a, b)
s
    l :: Map a (Set b)
l = [(a, b)] -> Map a (Set b)
forall a b. (Ord a, Ord b) => [(a, b)] -> Map a (Set b)
ascListToMapSet [(a, b)]
abs
    bas :: [(b, a)]
bas = [(b, a)] -> [(b, a)]
forall a. Ord a => [a] -> [a]
sort ([(b, a)] -> [(b, a)]) -> [(b, a)] -> [(b, a)]
forall a b. (a -> b) -> a -> b
$ (\(a
a, b
b) -> (b
b, a
a)) ((a, b) -> (b, a)) -> [(a, b)] -> [(b, a)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(a, b)]
abs
    r :: Map b (Set a)
r = [(b, a)] -> Map b (Set a)
forall a b. (Ord a, Ord b) => [(a, b)] -> Map a (Set b)
ascListToMapSet [(b, a)]
bas
    m :: Map (a, b) c
m = ((a, b) -> c) -> Set (a, b) -> Map (a, b) c
forall k a. (k -> a) -> Set k -> Map k a
M.fromSet ((a -> b -> c) -> (a, b) -> c
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> b -> c
f) Set (a, b)
s

-- ** From unordered lists
--------------------------

{-# INLINABLE fromList #-}
fromList :: (Ord a, Ord b) => [(a, b, c)] -> BimapMany a b c
fromList :: [(a, b, c)] -> BimapMany a b c
fromList = (BimapMany a b c -> (a, b, c) -> BimapMany a b c)
-> BimapMany a b c -> [(a, b, c)] -> BimapMany a b c
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\BimapMany a b c
m (a
a, b
b, c
c) -> a -> b -> c -> BimapMany a b c -> BimapMany a b c
forall a b c.
(Ord a, Ord b) =>
a -> b -> c -> BimapMany a b c -> BimapMany a b c
insert a
a b
b c
c BimapMany a b c
m) BimapMany a b c
forall a b c. BimapMany a b c
empty

-- TODO fromListWith(Key)
-- TODO asc and desc unsafe construction

-- * Insertion
--------------

{-# INLINABLE insert #-}
insert :: (Ord a, Ord b) => a -> b -> c -> BimapMany a b c -> BimapMany a b c
insert :: a -> b -> c -> BimapMany a b c -> BimapMany a b c
insert a
a b
b c
c (BimapMany Map a (Set b)
l Map b (Set a)
r Map (a, b) c
m) = Map a (Set b) -> Map b (Set a) -> Map (a, b) c -> BimapMany a b c
forall a b c.
Map a (Set b) -> Map b (Set a) -> Map (a, b) c -> BimapMany a b c
BimapMany Map a (Set b)
l' Map b (Set a)
r' Map (a, b) c
m'
  where
    l' :: Map a (Set b)
l' = (Maybe (Set b) -> Maybe (Set b))
-> a -> Map a (Set b) -> Map a (Set b)
forall k a.
Ord k =>
(Maybe a -> Maybe a) -> k -> Map k a -> Map k a
MS.alter (Set b -> Maybe (Set b)
forall a. a -> Maybe a
Just (Set b -> Maybe (Set b))
-> (Maybe (Set b) -> Set b) -> Maybe (Set b) -> Maybe (Set b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Set b -> (Set b -> Set b) -> Maybe (Set b) -> Set b
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (b -> Set b
forall a. a -> Set a
S.singleton b
b) (b -> Set b -> Set b
forall a. Ord a => a -> Set a -> Set a
S.insert b
b)) a
a Map a (Set b)
l
    r' :: Map b (Set a)
r' = (Maybe (Set a) -> Maybe (Set a))
-> b -> Map b (Set a) -> Map b (Set a)
forall k a.
Ord k =>
(Maybe a -> Maybe a) -> k -> Map k a -> Map k a
MS.alter (Set a -> Maybe (Set a)
forall a. a -> Maybe a
Just (Set a -> Maybe (Set a))
-> (Maybe (Set a) -> Set a) -> Maybe (Set a) -> Maybe (Set a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Set a -> (Set a -> Set a) -> Maybe (Set a) -> Set a
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (a -> Set a
forall a. a -> Set a
S.singleton a
a) (a -> Set a -> Set a
forall a. Ord a => a -> Set a -> Set a
S.insert a
a)) b
b Map b (Set a)
r
    m' :: Map (a, b) c
m' = (a, b) -> c -> Map (a, b) c -> Map (a, b) c
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert (a
a, b
b) c
c Map (a, b) c
m

-- TODO insertWith etc

-- * Deletion/Update
--------------------

{-# INLINABLE delete #-}
delete :: (Ord a, Ord b) => a -> b -> BimapMany a b c -> BimapMany a b c
delete :: a -> b -> BimapMany a b c -> BimapMany a b c
delete a
a b
b (BimapMany Map a (Set b)
l Map b (Set a)
r Map (a, b) c
m) = Map a (Set b) -> Map b (Set a) -> Map (a, b) c -> BimapMany a b c
forall a b c.
Map a (Set b) -> Map b (Set a) -> Map (a, b) c -> BimapMany a b c
BimapMany Map a (Set b)
l' Map b (Set a)
r' Map (a, b) c
m'
  where
    l' :: Map a (Set b)
l' = (Set b -> Maybe (Set b)) -> a -> Map a (Set b) -> Map a (Set b)
forall k a. Ord k => (a -> Maybe a) -> k -> Map k a -> Map k a
MS.update (b -> Set b -> Maybe (Set b)
forall a. Ord a => a -> Set a -> Maybe (Set a)
setDelete' b
b) a
a Map a (Set b)
l
    r' :: Map b (Set a)
r' = (Set a -> Maybe (Set a)) -> b -> Map b (Set a) -> Map b (Set a)
forall k a. Ord k => (a -> Maybe a) -> k -> Map k a -> Map k a
MS.update (a -> Set a -> Maybe (Set a)
forall a. Ord a => a -> Set a -> Maybe (Set a)
setDelete' a
a) b
b Map b (Set a)
r
    m' :: Map (a, b) c
m' = (a, b) -> Map (a, b) c -> Map (a, b) c
forall k a. Ord k => k -> Map k a -> Map k a
M.delete (a
a, b
b) Map (a, b) c
m

{-# INLINABLE deleteL #-}
deleteL :: (Ord a, Ord b) => a -> BimapMany a b c -> BimapMany a b c
deleteL :: a -> BimapMany a b c -> BimapMany a b c
deleteL a
a (BimapMany Map a (Set b)
l Map b (Set a)
r Map (a, b) c
m) = Map a (Set b) -> Map b (Set a) -> Map (a, b) c -> BimapMany a b c
forall a b c.
Map a (Set b) -> Map b (Set a) -> Map (a, b) c -> BimapMany a b c
BimapMany Map a (Set b)
l' Map b (Set a)
r' Map (a, b) c
m'
  where
    bs :: Set b
bs = Set b -> Maybe (Set b) -> Set b
forall a. a -> Maybe a -> a
fromMaybe Set b
forall a. Set a
S.empty (Maybe (Set b) -> Set b) -> Maybe (Set b) -> Set b
forall a b. (a -> b) -> a -> b
$ a -> Map a (Set b) -> Maybe (Set b)
forall k a. Ord k => k -> Map k a -> Maybe a
MS.lookup a
a Map a (Set b)
l
    l' :: Map a (Set b)
l' = a -> Map a (Set b) -> Map a (Set b)
forall k a. Ord k => k -> Map k a -> Map k a
MS.delete a
a Map a (Set b)
l
    r' :: Map b (Set a)
r' = (b -> Map b (Set a) -> Map b (Set a))
-> Map b (Set a) -> Set b -> Map b (Set a)
forall a b. (a -> b -> b) -> b -> Set a -> b
S.foldr' ((Set a -> Maybe (Set a)) -> b -> Map b (Set a) -> Map b (Set a)
forall k a. Ord k => (a -> Maybe a) -> k -> Map k a -> Map k a
MS.update ((Set a -> Maybe (Set a)) -> b -> Map b (Set a) -> Map b (Set a))
-> (Set a -> Maybe (Set a)) -> b -> Map b (Set a) -> Map b (Set a)
forall a b. (a -> b) -> a -> b
$ a -> Set a -> Maybe (Set a)
forall a. Ord a => a -> Set a -> Maybe (Set a)
setDelete' a
a) Map b (Set a)
r Set b
bs
    m' :: Map (a, b) c
m' = (b -> Map (a, b) c -> Map (a, b) c)
-> Map (a, b) c -> Set b -> Map (a, b) c
forall a b. (a -> b -> b) -> b -> Set a -> b
S.foldr' (\b
b -> (a, b) -> Map (a, b) c -> Map (a, b) c
forall k a. Ord k => k -> Map k a -> Map k a
M.delete (a
a, b
b)) Map (a, b) c
m Set b
bs

{-# INLINABLE deleteR #-}
deleteR :: (Ord a, Ord b) => b -> BimapMany a b c -> BimapMany a b c
deleteR :: b -> BimapMany a b c -> BimapMany a b c
deleteR b
b (BimapMany Map a (Set b)
l Map b (Set a)
r Map (a, b) c
m) = Map a (Set b) -> Map b (Set a) -> Map (a, b) c -> BimapMany a b c
forall a b c.
Map a (Set b) -> Map b (Set a) -> Map (a, b) c -> BimapMany a b c
BimapMany Map a (Set b)
l' Map b (Set a)
r' Map (a, b) c
m'
  where
    as :: Set a
as = Set a -> Maybe (Set a) -> Set a
forall a. a -> Maybe a -> a
fromMaybe Set a
forall a. Set a
S.empty (Maybe (Set a) -> Set a) -> Maybe (Set a) -> Set a
forall a b. (a -> b) -> a -> b
$ b -> Map b (Set a) -> Maybe (Set a)
forall k a. Ord k => k -> Map k a -> Maybe a
MS.lookup b
b Map b (Set a)
r
    r' :: Map b (Set a)
r' = b -> Map b (Set a) -> Map b (Set a)
forall k a. Ord k => k -> Map k a -> Map k a
MS.delete b
b Map b (Set a)
r
    l' :: Map a (Set b)
l' = (a -> Map a (Set b) -> Map a (Set b))
-> Map a (Set b) -> Set a -> Map a (Set b)
forall a b. (a -> b -> b) -> b -> Set a -> b
S.foldr' ((Set b -> Maybe (Set b)) -> a -> Map a (Set b) -> Map a (Set b)
forall k a. Ord k => (a -> Maybe a) -> k -> Map k a -> Map k a
MS.update ((Set b -> Maybe (Set b)) -> a -> Map a (Set b) -> Map a (Set b))
-> (Set b -> Maybe (Set b)) -> a -> Map a (Set b) -> Map a (Set b)
forall a b. (a -> b) -> a -> b
$ b -> Set b -> Maybe (Set b)
forall a. Ord a => a -> Set a -> Maybe (Set a)
setDelete' b
b) Map a (Set b)
l Set a
as
    m' :: Map (a, b) c
m' = (a -> Map (a, b) c -> Map (a, b) c)
-> Map (a, b) c -> Set a -> Map (a, b) c
forall a b. (a -> b -> b) -> b -> Set a -> b
S.foldr' (\a
a -> (a, b) -> Map (a, b) c -> Map (a, b) c
forall k a. Ord k => k -> Map k a -> Map k a
M.delete (a
a, b
b)) Map (a, b) c
m Set a
as

-- TODO updating functions

-- * Query
----------

-- ** Lookup
------------

{-# INLINABLE lookup #-}
lookup :: (Ord a, Ord b) => a -> b -> BimapMany a b c -> Maybe c
lookup :: a -> b -> BimapMany a b c -> Maybe c
lookup a
a b
b (BimapMany Map a (Set b)
_ Map b (Set a)
_ Map (a, b) c
m) = (a, b) -> Map (a, b) c -> Maybe c
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (a
a, b
b) Map (a, b) c
m

{-# INLINABLE lookupL #-}
lookupL :: Ord a => a -> BimapMany a b c -> Set b
lookupL :: a -> BimapMany a b c -> Set b
lookupL a
a (BimapMany Map a (Set b)
l Map b (Set a)
_ Map (a, b) c
_) = Maybe (Set b) -> Set b
forall a. Maybe (Set a) -> Set a
mSetToSet (Maybe (Set b) -> Set b) -> Maybe (Set b) -> Set b
forall a b. (a -> b) -> a -> b
$ a -> Map a (Set b) -> Maybe (Set b)
forall k a. Ord k => k -> Map k a -> Maybe a
MS.lookup a
a Map a (Set b)
l

{-# INLINABLE lookupR #-}
lookupR :: Ord b => b -> BimapMany a b c -> Set a
lookupR :: b -> BimapMany a b c -> Set a
lookupR b
b (BimapMany Map a (Set b)
_ Map b (Set a)
r Map (a, b) c
_) = Maybe (Set a) -> Set a
forall a. Maybe (Set a) -> Set a
mSetToSet (Maybe (Set a) -> Set a) -> Maybe (Set a) -> Set a
forall a b. (a -> b) -> a -> b
$ b -> Map b (Set a) -> Maybe (Set a)
forall k a. Ord k => k -> Map k a -> Maybe a
MS.lookup b
b Map b (Set a)
r

{-# INLINABLE lookupL' #-}
lookupL' :: (Ord a, Ord b) => a -> BimapMany a b c -> Map b c
lookupL' :: a -> BimapMany a b c -> Map b c
lookupL' a
a (BimapMany Map a (Set b)
l Map b (Set a)
_ Map (a, b) c
m) = (b -> c) -> Set b -> Map b c
forall k a. (k -> a) -> Set k -> Map k a
M.fromSet (\b
b -> Map (a, b) c
m Map (a, b) c -> (a, b) -> c
forall k a. Ord k => Map k a -> k -> a
M.! (a
a, b
b)) Set b
bs
  where bs :: Set b
bs = Maybe (Set b) -> Set b
forall a. Maybe (Set a) -> Set a
mSetToSet (Maybe (Set b) -> Set b) -> Maybe (Set b) -> Set b
forall a b. (a -> b) -> a -> b
$ a -> Map a (Set b) -> Maybe (Set b)
forall k a. Ord k => k -> Map k a -> Maybe a
MS.lookup a
a Map a (Set b)
l

{-# INLINABLE lookupR' #-}
lookupR' :: (Ord a, Ord b) => b -> BimapMany a b c -> Map a c
lookupR' :: b -> BimapMany a b c -> Map a c
lookupR' b
b (BimapMany Map a (Set b)
_ Map b (Set a)
r Map (a, b) c
m) = (a -> c) -> Set a -> Map a c
forall k a. (k -> a) -> Set k -> Map k a
M.fromSet (\a
a -> Map (a, b) c
m Map (a, b) c -> (a, b) -> c
forall k a. Ord k => Map k a -> k -> a
M.! (a
a, b
b)) Set a
as
  where as :: Set a
as = Maybe (Set a) -> Set a
forall a. Maybe (Set a) -> Set a
mSetToSet (Maybe (Set a) -> Set a) -> Maybe (Set a) -> Set a
forall a b. (a -> b) -> a -> b
$ b -> Map b (Set a) -> Maybe (Set a)
forall k a. Ord k => k -> Map k a -> Maybe a
MS.lookup b
b Map b (Set a)
r

-- ** Size
----------

{-# INLINABLE null #-}
null :: BimapMany a b c -> Bool
null :: BimapMany a b c -> Bool
null (BimapMany Map a (Set b)
_ Map b (Set a)
_ Map (a, b) c
m) = Map (a, b) c -> Bool
forall k a. Map k a -> Bool
M.null Map (a, b) c
m

{-# INLINABLE size #-}
size :: BimapMany a b c -> Int
size :: BimapMany a b c -> Int
size (BimapMany Map a (Set b)
_ Map b (Set a)
_ Map (a, b) c
m) = Map (a, b) c -> Int
forall k a. Map k a -> Int
M.size Map (a, b) c
m

{-# INLINABLE sizeL #-}
sizeL :: BimapMany a b c -> Int
sizeL :: BimapMany a b c -> Int
sizeL (BimapMany Map a (Set b)
l Map b (Set a)
_ Map (a, b) c
_) = Map a (Set b) -> Int
forall k a. Map k a -> Int
MS.size Map a (Set b)
l

{-# INLINABLE sizeR #-}
sizeR :: BimapMany a b c -> Int
sizeR :: BimapMany a b c -> Int
sizeR (BimapMany Map a (Set b)
_ Map b (Set a)
r Map (a, b) c
_) = Map b (Set a) -> Int
forall k a. Map k a -> Int
MS.size Map b (Set a)
r

-- * Combine
------------

{-# INLINABLE union #-}
union :: (Ord a, Ord b) => BimapMany a b c -> BimapMany a b c -> BimapMany a b c
union :: BimapMany a b c -> BimapMany a b c -> BimapMany a b c
union (BimapMany Map a (Set b)
l1 Map b (Set a)
r1 Map (a, b) c
m1) (BimapMany Map a (Set b)
l2 Map b (Set a)
r2 Map (a, b) c
m2) = Map a (Set b) -> Map b (Set a) -> Map (a, b) c -> BimapMany a b c
forall a b c.
Map a (Set b) -> Map b (Set a) -> Map (a, b) c -> BimapMany a b c
BimapMany Map a (Set b)
l Map b (Set a)
r Map (a, b) c
m
  where
    l :: Map a (Set b)
l = (Set b -> Set b -> Set b)
-> Map a (Set b) -> Map a (Set b) -> Map a (Set b)
forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
MS.unionWith Set b -> Set b -> Set b
forall a. Ord a => Set a -> Set a -> Set a
S.union Map a (Set b)
l1 Map a (Set b)
l2
    r :: Map b (Set a)
r = (Set a -> Set a -> Set a)
-> Map b (Set a) -> Map b (Set a) -> Map b (Set a)
forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
MS.unionWith Set a -> Set a -> Set a
forall a. Ord a => Set a -> Set a -> Set a
S.union Map b (Set a)
r1 Map b (Set a)
r2
    m :: Map (a, b) c
m = Map (a, b) c -> Map (a, b) c -> Map (a, b) c
forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union Map (a, b) c
m1 Map (a, b) c
m2

-- * Conversion
---------------

-- ** Maps
----------

{-# INLINABLE toMap #-}
toMap :: BimapMany a b c -> Map (a, b) c
toMap :: BimapMany a b c -> Map (a, b) c
toMap (BimapMany Map a (Set b)
_ Map b (Set a)
_ Map (a, b) c
m) = Map (a, b) c
m

-- MAYBE BimapMany a b c -> Map a (Set b) (and mirrored)
-- MAYBE toNestedMapL :: BimapMany a b c -> Map a (Map b c) (and mirrored)

-- ** Lists
-----------

{-# INLINABLE toList #-}
toList :: BimapMany a b c -> [(a, b, c)]
toList :: BimapMany a b c -> [(a, b, c)]
toList (BimapMany Map a (Set b)
_ Map b (Set a)
_ Map (a, b) c
m) = (\((a
a, b
b), c
c) -> (a
a, b
b, c
c)) (((a, b), c) -> (a, b, c)) -> [((a, b), c)] -> [(a, b, c)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map (a, b) c -> [((a, b), c)]
forall k a. Map k a -> [(k, a)]
M.toList Map (a, b) c
m

-- * Debugging
--------------

valid :: (Ord a, Ord b) => BimapMany a b c -> Bool
valid :: BimapMany a b c -> Bool
valid (BimapMany Map a (Set b)
l Map b (Set a)
r Map (a, b) c
m) = Bool
prop1 Bool -> Bool -> Bool
&& Bool
prop2 Bool -> Bool -> Bool
&& Bool
prop3
  where
    prop1 :: Bool
prop1 = ((a, b) -> Bool) -> [(a, b)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\(a
a, b
b) -> (a, b) -> Map (a, b) c -> Bool
forall k a. Ord k => k -> Map k a -> Bool
M.member (a
a, b
b) Map (a, b) c
m Bool -> Bool -> Bool
&&
                            Bool -> (Set a -> Bool) -> Maybe (Set a) -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False (a -> Set a -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member a
a) (Map b (Set a)
r Map b (Set a) -> b -> Maybe (Set a)
forall k a. Ord k => Map k a -> k -> Maybe a
MS.!? b
b)) ([(a, b)] -> Bool) -> [(a, b)] -> Bool
forall a b. (a -> b) -> a -> b
$
                ((a, Set b) -> [(a, b)]) -> [(a, Set b)] -> [(a, b)]
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (\(a
a, Set b
bs) -> (,) a
a (b -> (a, b)) -> [b] -> [(a, b)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Set b -> [b]
forall a. Set a -> [a]
S.toList Set b
bs) ([(a, Set b)] -> [(a, b)]) -> [(a, Set b)] -> [(a, b)]
forall a b. (a -> b) -> a -> b
$ Map a (Set b) -> [(a, Set b)]
forall k a. Map k a -> [(k, a)]
MS.toList Map a (Set b)
l
    prop2 :: Bool
prop2 = ((b, a) -> Bool) -> [(b, a)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\(b
b, a
a) -> (a, b) -> Map (a, b) c -> Bool
forall k a. Ord k => k -> Map k a -> Bool
M.member (a
a, b
b) Map (a, b) c
m Bool -> Bool -> Bool
&&
                            Bool -> (Set b -> Bool) -> Maybe (Set b) -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False (b -> Set b -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member b
b) (Map a (Set b)
l Map a (Set b) -> a -> Maybe (Set b)
forall k a. Ord k => Map k a -> k -> Maybe a
MS.!? a
a)) ([(b, a)] -> Bool) -> [(b, a)] -> Bool
forall a b. (a -> b) -> a -> b
$
                ((b, Set a) -> [(b, a)]) -> [(b, Set a)] -> [(b, a)]
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (\(b
b, Set a
as) -> (,) b
b (a -> (b, a)) -> [a] -> [(b, a)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Set a -> [a]
forall a. Set a -> [a]
S.toList Set a
as) ([(b, Set a)] -> [(b, a)]) -> [(b, Set a)] -> [(b, a)]
forall a b. (a -> b) -> a -> b
$ Map b (Set a) -> [(b, Set a)]
forall k a. Map k a -> [(k, a)]
MS.toList Map b (Set a)
r
    prop3 :: Bool
prop3 = ((a, b) -> Bool) -> [(a, b)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\(a
a, b
b) -> Bool -> (Set b -> Bool) -> Maybe (Set b) -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False (b -> Set b -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member b
b) (Map a (Set b)
l Map a (Set b) -> a -> Maybe (Set b)
forall k a. Ord k => Map k a -> k -> Maybe a
MS.!? a
a) Bool -> Bool -> Bool
&&
                            Bool -> (Set a -> Bool) -> Maybe (Set a) -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False (a -> Set a -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member a
a) (Map b (Set a)
r Map b (Set a) -> b -> Maybe (Set a)
forall k a. Ord k => Map k a -> k -> Maybe a
MS.!? b
b)) ([(a, b)] -> Bool) -> [(a, b)] -> Bool
forall a b. (a -> b) -> a -> b
$
                Map (a, b) c -> [(a, b)]
forall k a. Map k a -> [k]
M.keys Map (a, b) c
m

-- Internal utilities
---------------------

{-# INLINE mSetToSet #-}
mSetToSet :: Maybe (Set a) -> Set a
mSetToSet :: Maybe (Set a) -> Set a
mSetToSet Maybe (Set a)
Nothing = Set a
forall a. Set a
S.empty
mSetToSet (Just Set a
set) = Set a
set

{-# INLINE ascListToMapSet #-}
ascListToMapSet :: (Ord a, Ord b) => [(a, b)] -> MS.Map a (Set b)
ascListToMapSet :: [(a, b)] -> Map a (Set b)
ascListToMapSet [(a, b)]
abs = [(a, Set b)] -> Map a (Set b)
forall k a. [(k, a)] -> Map k a
MS.fromDistinctAscList [(a, Set b)]
sets
  where
    grouped :: [[(a, b)]]
grouped = ((a, b) -> (a, b) -> Bool) -> [(a, b)] -> [[(a, b)]]
forall a. (a -> a -> Bool) -> [a] -> [[a]]
groupBy (a -> a -> Bool
forall a. Eq a => a -> a -> Bool
(==) (a -> a -> Bool) -> ((a, b) -> a) -> (a, b) -> (a, b) -> Bool
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` (a, b) -> a
forall a b. (a, b) -> a
fst) [(a, b)]
abs
    sets :: [(a, Set b)]
sets = (\[(a, b)]
xs -> ((a, b) -> a
forall a b. (a, b) -> a
fst ((a, b) -> a) -> (a, b) -> a
forall a b. (a -> b) -> a -> b
$ [(a, b)] -> (a, b)
forall a. [a] -> a
head [(a, b)]
xs, [b] -> Set b
forall a. Eq a => [a] -> Set a
S.fromAscList ((a, b) -> b
forall a b. (a, b) -> b
snd ((a, b) -> b) -> [(a, b)] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(a, b)]
xs))) ([(a, b)] -> (a, Set b)) -> [[(a, b)]] -> [(a, Set b)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [[(a, b)]]
grouped

{-# INLINE setDelete' #-}
setDelete' :: Ord a => a -> Set a -> Maybe (Set a)
setDelete' :: a -> Set a -> Maybe (Set a)
setDelete' a
x Set a
s = if Set a -> Bool
forall a. Set a -> Bool
S.null Set a
s' then Maybe (Set a)
forall a. Maybe a
Nothing else Set a -> Maybe (Set a)
forall a. a -> Maybe a
Just Set a
s'
  where s' :: Set a
s' = a -> Set a -> Set a
forall a. Ord a => a -> Set a -> Set a
S.delete a
x Set a
s