-- |Module `Containers.PrismMap` include type `PrismMap` and useful functions
--  for working with it.
--
--  `PrismMap` is kind of `Map` with faster access to keys by value. So, it is
--  almost `IsoMap`, but without isomorphism restriction.
module Containers.PrismMap (
    PrismMap,
    size,
    Containers.PrismMap.null,
    empty,
    singleton,
    fromList,
    toList,
    toMap,
    Containers.PrismMap.lookup,
    lookupKeys,
    values,
    keys,
    keysSet,
    value,
    key,
    Containers.PrismMap.map,
    mapKeys,
    Containers.PrismMap.filter,
    filterKeys,
    insert,
    deleteKey,
    deleteValue,
    union,
    intersection,
    difference,
  ) where

import Containers.Set (Set)
import qualified Containers.Set as Set
import Containers.Map (Map)
import qualified Containers.Map as Map
import Lens

import Control.Monad (guard)
import Data.Tuple (swap)
import Data.Maybe (fromMaybe)

newtype PrismMap k a = PrismMap (Map k a, Map a (Set k))
    deriving (PrismMap k a -> PrismMap k a -> Bool
(PrismMap k a -> PrismMap k a -> Bool)
-> (PrismMap k a -> PrismMap k a -> Bool) -> Eq (PrismMap k a)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall k a. (Eq k, Eq a) => PrismMap k a -> PrismMap k a -> Bool
/= :: PrismMap k a -> PrismMap k a -> Bool
$c/= :: forall k a. (Eq k, Eq a) => PrismMap k a -> PrismMap k a -> Bool
== :: PrismMap k a -> PrismMap k a -> Bool
$c== :: forall k a. (Eq k, Eq a) => PrismMap k a -> PrismMap k a -> Bool
Eq)

instance (Show k, Show a) => Show (PrismMap k a) where
    show :: PrismMap k a -> String
show (PrismMap (Map k a
m, Map a (Set k)
_)) = Map k a -> String
forall a. Show a => a -> String
show Map k a
m

(??) :: MonadFail m => Maybe a -> String -> m a
Maybe a
ma ?? :: Maybe a -> String -> m a
?? String
msg = m a -> (a -> m a) -> Maybe a -> m a
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (String -> m a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
msg) a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe a -> m a) -> Maybe a -> m a
forall a b. (a -> b) -> a -> b
$ Maybe a
ma

size :: PrismMap k a -> Int
size :: PrismMap k a -> Int
size (PrismMap (Map k a
m, Map a (Set k)
_)) = Map k a -> Int
forall k a. Map k a -> Int
Map.size Map k a
m

null :: PrismMap k a -> Bool
null :: PrismMap k a -> Bool
null (PrismMap (Map k a
m, Map a (Set k)
_)) = Map k a -> Bool
forall k a. Map k a -> Bool
Map.null Map k a
m

empty :: PrismMap k a
empty :: PrismMap k a
empty = (Map k a, Map a (Set k)) -> PrismMap k a
forall k a. (Map k a, Map a (Set k)) -> PrismMap k a
PrismMap (Map k a
forall k a. Map k a
Map.empty, Map a (Set k)
forall k a. Map k a
Map.empty)

singleton :: k -> a -> PrismMap k a
singleton :: k -> a -> PrismMap k a
singleton k
k a
a = (Map k a, Map a (Set k)) -> PrismMap k a
forall k a. (Map k a, Map a (Set k)) -> PrismMap k a
PrismMap (k -> a -> Map k a
forall k a. k -> a -> Map k a
Map.singleton k
k a
a, a -> Set k -> Map a (Set k)
forall k a. k -> a -> Map k a
Map.singleton a
a (Set k -> Map a (Set k)) -> Set k -> Map a (Set k)
forall a b. (a -> b) -> a -> b
$ k -> Set k
forall a. a -> Set a
Set.singleton k
k)

sndMapFromList :: (Ord k, Ord a) => [(k, a)] -> Map a (Set k)
sndMapFromList :: [(k, a)] -> Map a (Set k)
sndMapFromList [(k, a)]
xs =
    (Set k -> Set k -> Set k) -> [(a, Set k)] -> Map a (Set k)
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
Map.fromListWith Set k -> Set k -> Set k
forall a. Ord a => Set a -> Set a -> Set a
Set.union ([(a, Set k)] -> Map a (Set k)) -> [(a, Set k)] -> Map a (Set k)
forall a b. (a -> b) -> a -> b
$ ((Set k, a) -> (a, Set k)
forall a b. (a, b) -> (b, a)
swap ((Set k, a) -> (a, Set k))
-> ((k, a) -> (Set k, a)) -> (k, a) -> (a, Set k)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((k -> Identity (Set k)) -> (k, a) -> Identity (Set k, a)
forall s t a b. Field1 s t a b => Lens s t a b
_1 ((k -> Identity (Set k)) -> (k, a) -> Identity (Set k, a))
-> (k -> Set k) -> (k, a) -> (Set k, a)
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
%~ k -> Set k
forall a. a -> Set a
Set.singleton)) ((k, a) -> (a, Set k)) -> [(k, a)] -> [(a, Set k)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(k, a)]
xs

fromList :: (Ord k, Ord a) => [(k, a)] -> PrismMap k a
fromList :: [(k, a)] -> PrismMap k a
fromList [(k, a)]
xs = (Map k a, Map a (Set k)) -> PrismMap k a
forall k a. (Map k a, Map a (Set k)) -> PrismMap k a
PrismMap ([(k, a)] -> Map k a
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [(k, a)]
xs, [(k, a)] -> Map a (Set k)
forall k a. (Ord k, Ord a) => [(k, a)] -> Map a (Set k)
sndMapFromList [(k, a)]
xs)

toList :: PrismMap k a -> [(k, a)]
toList :: PrismMap k a -> [(k, a)]
toList (PrismMap (Map k a
m, Map a (Set k)
_)) = Map k a -> [(k, a)]
forall k a. Map k a -> [(k, a)]
Map.toList Map k a
m

toMap :: PrismMap k a -> Map k a
toMap :: PrismMap k a -> Map k a
toMap (PrismMap (Map k a
m, Map a (Set k)
_)) = Map k a
m

lookup :: (Ord k, MonadFail m) => k -> PrismMap k a -> m a
lookup :: k -> PrismMap k a -> m a
lookup k
k (PrismMap (Map k a
m, Map a (Set k)
_)) = k -> Map k a -> Maybe a
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup k
k Map k a
m Maybe a -> String -> m a
forall (m :: * -> *) a. MonadFail m => Maybe a -> String -> m a
?? String
"Key was not found"

lookupKeys :: Ord a => a -> PrismMap k a -> Set k
lookupKeys :: a -> PrismMap k a -> Set k
lookupKeys a
a (PrismMap (Map k a
_, Map a (Set k)
w)) = Set k -> Maybe (Set k) -> Set k
forall a. a -> Maybe a -> a
fromMaybe Set k
forall a. Set a
Set.empty (Maybe (Set k) -> Set k) -> Maybe (Set k) -> Set k
forall a b. (a -> b) -> a -> b
$ a -> Map a (Set k) -> Maybe (Set k)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup a
a Map a (Set k)
w

values :: PrismMap k a -> [a]
values :: PrismMap k a -> [a]
values (PrismMap (Map k a
m, Map a (Set k)
_)) = Map k a -> [a]
forall k a. Map k a -> [a]
Map.elems Map k a
m

keys :: PrismMap k a -> [k]
keys :: PrismMap k a -> [k]
keys (PrismMap (Map k a
m, Map a (Set k)
_)) = Map k a -> [k]
forall k a. Map k a -> [k]
Map.keys Map k a
m

keysSet :: PrismMap k a -> Set k
keysSet :: PrismMap k a -> Set k
keysSet (PrismMap (Map k a
m, Map a (Set k)
_)) = Map k a -> Set k
forall k a. Map k a -> Set k
Map.keysSet Map k a
m

value :: Ord a => a -> PrismMap k a -> Bool
value :: a -> PrismMap k a -> Bool
value a
a (PrismMap (Map k a
_, Map a (Set k)
w)) = a -> Map a (Set k) -> Bool
forall k a. Ord k => k -> Map k a -> Bool
Map.member a
a Map a (Set k)
w

key :: Ord k => k -> PrismMap k a -> Bool
key :: k -> PrismMap k a -> Bool
key k
k (PrismMap (Map k a
m, Map a (Set k)
_)) = k -> Map k a -> Bool
forall k a. Ord k => k -> Map k a -> Bool
Map.member k
k Map k a
m

map :: Ord b => (a -> b) -> PrismMap k a -> PrismMap k b
map :: (a -> b) -> PrismMap k a -> PrismMap k b
map a -> b
f (PrismMap (Map k a
m, Map a (Set k)
w)) = (Map k b, Map b (Set k)) -> PrismMap k b
forall k a. (Map k a, Map a (Set k)) -> PrismMap k a
PrismMap ((a -> b) -> Map k a -> Map k b
forall a b k. (a -> b) -> Map k a -> Map k b
Map.map a -> b
f Map k a
m, (a -> b) -> Map a (Set k) -> Map b (Set k)
forall k2 k1 a. Ord k2 => (k1 -> k2) -> Map k1 a -> Map k2 a
Map.mapKeys a -> b
f Map a (Set k)
w)

mapKeys :: Ord l => (k -> l) -> PrismMap k a -> PrismMap l a
mapKeys :: (k -> l) -> PrismMap k a -> PrismMap l a
mapKeys k -> l
f (PrismMap (Map k a
m, Map a (Set k)
w)) = (Map l a, Map a (Set l)) -> PrismMap l a
forall k a. (Map k a, Map a (Set k)) -> PrismMap k a
PrismMap ((k -> l) -> Map k a -> Map l a
forall k2 k1 a. Ord k2 => (k1 -> k2) -> Map k1 a -> Map k2 a
Map.mapKeys k -> l
f Map k a
m, ((Set k -> Set l) -> Map a (Set k) -> Map a (Set l)
forall a b k. (a -> b) -> Map k a -> Map k b
Map.map((Set k -> Set l) -> Map a (Set k) -> Map a (Set l))
-> ((k -> l) -> Set k -> Set l)
-> (k -> l)
-> Map a (Set k)
-> Map a (Set l)
forall b c a. (b -> c) -> (a -> b) -> a -> c
.(k -> l) -> Set k -> Set l
forall b a. Ord b => (a -> b) -> Set a -> Set b
Set.map) k -> l
f Map a (Set k)
w)

filter :: (a -> Bool) -> PrismMap k a -> PrismMap k a
filter :: (a -> Bool) -> PrismMap k a -> PrismMap k a
filter a -> Bool
p (PrismMap (Map k a
m, Map a (Set k)
w)) =
    (Map k a, Map a (Set k)) -> PrismMap k a
forall k a. (Map k a, Map a (Set k)) -> PrismMap k a
PrismMap ((a -> Bool) -> Map k a -> Map k a
forall a k. (a -> Bool) -> Map k a -> Map k a
Map.filter a -> Bool
p Map k a
m, (a -> Set k -> Bool) -> Map a (Set k) -> Map a (Set k)
forall k a. (k -> a -> Bool) -> Map k a -> Map k a
Map.filterWithKey (Bool -> Set k -> Bool
forall a b. a -> b -> a
const(Bool -> Set k -> Bool) -> (a -> Bool) -> a -> Set k -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
.a -> Bool
p) Map a (Set k)
w)

filterKeys :: (k -> Bool) -> PrismMap k a -> PrismMap k a
filterKeys :: (k -> Bool) -> PrismMap k a -> PrismMap k a
filterKeys k -> Bool
p (PrismMap (Map k a
m, Map a (Set k)
w)) =
    let w' :: Map a (Set k)
w' = (Set k -> Set k) -> Map a (Set k) -> Map a (Set k)
forall a b k. (a -> b) -> Map k a -> Map k b
Map.map ((k -> Bool) -> Set k -> Set k
forall a. (a -> Bool) -> Set a -> Set a
Set.filter k -> Bool
p) Map a (Set k)
w
    in  (Map k a, Map a (Set k)) -> PrismMap k a
forall k a. (Map k a, Map a (Set k)) -> PrismMap k a
PrismMap ((k -> a -> Bool) -> Map k a -> Map k a
forall k a. (k -> a -> Bool) -> Map k a -> Map k a
Map.filterWithKey (Bool -> a -> Bool
forall a b. a -> b -> a
const(Bool -> a -> Bool) -> (k -> Bool) -> k -> a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
.k -> Bool
p) Map k a
m, (Set k -> Bool) -> Map a (Set k) -> Map a (Set k)
forall a k. (a -> Bool) -> Map k a -> Map k a
Map.filter (Bool -> Bool
not(Bool -> Bool) -> (Set k -> Bool) -> Set k -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
.Set k -> Bool
forall a. Set a -> Bool
Set.null) Map a (Set k)
w')

insert :: (Ord k, Ord a) => k -> a -> PrismMap k a -> PrismMap k a
insert :: k -> a -> PrismMap k a -> PrismMap k a
insert k
k a
a (PrismMap (Map k a
m, Map a (Set k)
w)) =
    (Map k a, Map a (Set k)) -> PrismMap k a
forall k a. (Map k a, Map a (Set k)) -> PrismMap k a
PrismMap (k -> a -> Map k a -> Map k a
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert k
k a
a Map k a
m, (Set k -> Maybe (Set k)) -> a -> Map a (Set k) -> Map a (Set k)
forall k a. Ord k => (a -> Maybe a) -> k -> Map k a -> Map k a
Map.update (Set k -> Maybe (Set k)
forall a. a -> Maybe a
Just (Set k -> Maybe (Set k))
-> (Set k -> Set k) -> Set k -> Maybe (Set k)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. k -> Set k -> Set k
forall a. Ord a => a -> Set a -> Set a
Set.insert k
k) a
a Map a (Set k)
w)

deleteKey :: (Ord k, Ord a, MonadFail m) =>
    k -> PrismMap k a -> m (PrismMap k a)
deleteKey :: k -> PrismMap k a -> m (PrismMap k a)
deleteKey k
k (PrismMap (Map k a
m, Map a (Set k)
w)) = (Maybe (PrismMap k a) -> String -> m (PrismMap k a)
forall (m :: * -> *) a. MonadFail m => Maybe a -> String -> m a
?? String
"Key isn't in PrismMap") (Maybe (PrismMap k a) -> m (PrismMap k a))
-> Maybe (PrismMap k a) -> m (PrismMap k a)
forall a b. (a -> b) -> a -> b
$ do
    a
a <- Map k a
m Map k a -> k -> Maybe a
forall k a. Ord k => Map k a -> k -> Maybe a
Map.!? k
k
    let delKey :: Set k -> m (Set k)
delKey Set k
s = do
            let s' :: Set k
s' = k -> Set k -> Set k
forall a. Ord a => a -> Set a -> Set a
Set.delete k
k Set k
s
            Bool -> m ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> m ()) -> Bool -> m ()
forall a b. (a -> b) -> a -> b
$ Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Set k -> Bool
forall a. Set a -> Bool
Set.null Set k
s'
            Set k -> m (Set k)
forall (m :: * -> *) a. Monad m => a -> m a
return (Set k -> m (Set k)) -> Set k -> m (Set k)
forall a b. (a -> b) -> a -> b
$ Set k
s'
    PrismMap k a -> Maybe (PrismMap k a)
forall (m :: * -> *) a. Monad m => a -> m a
return (PrismMap k a -> Maybe (PrismMap k a))
-> PrismMap k a -> Maybe (PrismMap k a)
forall a b. (a -> b) -> a -> b
$
        (Map k a, Map a (Set k)) -> PrismMap k a
forall k a. (Map k a, Map a (Set k)) -> PrismMap k a
PrismMap (k -> Map k a -> Map k a
forall k a. Ord k => k -> Map k a -> Map k a
Map.delete k
k Map k a
m, (Set k -> Maybe (Set k)) -> a -> Map a (Set k) -> Map a (Set k)
forall k a. Ord k => (a -> Maybe a) -> k -> Map k a -> Map k a
Map.update Set k -> Maybe (Set k)
forall (m :: * -> *).
(Monad m, Alternative m) =>
Set k -> m (Set k)
delKey a
a Map a (Set k)
w)

deleteValue :: (Ord k, Ord a, MonadFail m) =>
    a -> PrismMap k a -> m (PrismMap k a)
deleteValue :: a -> PrismMap k a -> m (PrismMap k a)
deleteValue a
a (PrismMap (Map k a
m, Map a (Set k)
w)) = (Maybe (PrismMap k a) -> String -> m (PrismMap k a)
forall (m :: * -> *) a. MonadFail m => Maybe a -> String -> m a
?? String
"Value isn't in PrismMap") (Maybe (PrismMap k a) -> m (PrismMap k a))
-> Maybe (PrismMap k a) -> m (PrismMap k a)
forall a b. (a -> b) -> a -> b
$ do
    Set k
ks <- Map a (Set k)
w Map a (Set k) -> a -> Maybe (Set k)
forall k a. Ord k => Map k a -> k -> Maybe a
Map.!? a
a
    PrismMap k a -> Maybe (PrismMap k a)
forall (m :: * -> *) a. Monad m => a -> m a
return (PrismMap k a -> Maybe (PrismMap k a))
-> PrismMap k a -> Maybe (PrismMap k a)
forall a b. (a -> b) -> a -> b
$
        (Map k a, Map a (Set k)) -> PrismMap k a
forall k a. (Map k a, Map a (Set k)) -> PrismMap k a
PrismMap (Map k a -> Set k -> Map k a
forall k a. Ord k => Map k a -> Set k -> Map k a
Map.withoutKeys Map k a
m Set k
ks, a -> Map a (Set k) -> Map a (Set k)
forall k a. Ord k => k -> Map k a -> Map k a
Map.delete a
a Map a (Set k)
w)

union :: (Ord k, Ord a) => PrismMap k a -> PrismMap k a -> PrismMap k a
union :: PrismMap k a -> PrismMap k a -> PrismMap k a
union (PrismMap (Map k a
m1, Map a (Set k)
_)) (PrismMap (Map k a
m2, Map a (Set k)
_)) =
    let m :: Map k a
m = Map k a -> Map k a -> Map k a
forall k a. Ord k => Map k a -> Map k a -> Map k a
Map.union Map k a
m1 Map k a
m2
    in  (Map k a, Map a (Set k)) -> PrismMap k a
forall k a. (Map k a, Map a (Set k)) -> PrismMap k a
PrismMap (Map k a
m, [(k, a)] -> Map a (Set k)
forall k a. (Ord k, Ord a) => [(k, a)] -> Map a (Set k)
sndMapFromList ([(k, a)] -> Map a (Set k)) -> [(k, a)] -> Map a (Set k)
forall a b. (a -> b) -> a -> b
$ Map k a -> [(k, a)]
forall k a. Map k a -> [(k, a)]
Map.toList Map k a
m)

intersection :: (Ord k, Ord a) => PrismMap k a -> PrismMap k a -> PrismMap k a
intersection :: PrismMap k a -> PrismMap k a -> PrismMap k a
intersection (PrismMap (Map k a
m1, Map a (Set k)
_)) (PrismMap (Map k a
m2, Map a (Set k)
_)) =
    let m :: Map k a
m = Map k a -> Map k a -> Map k a
forall k a b. Ord k => Map k a -> Map k b -> Map k a
Map.intersection Map k a
m1 Map k a
m2
    in  (Map k a, Map a (Set k)) -> PrismMap k a
forall k a. (Map k a, Map a (Set k)) -> PrismMap k a
PrismMap (Map k a
m, [(k, a)] -> Map a (Set k)
forall k a. (Ord k, Ord a) => [(k, a)] -> Map a (Set k)
sndMapFromList ([(k, a)] -> Map a (Set k)) -> [(k, a)] -> Map a (Set k)
forall a b. (a -> b) -> a -> b
$ Map k a -> [(k, a)]
forall k a. Map k a -> [(k, a)]
Map.toList Map k a
m)

difference :: (Ord k, Ord a) => PrismMap k a -> PrismMap k a -> PrismMap k a
difference :: PrismMap k a -> PrismMap k a -> PrismMap k a
difference (PrismMap (Map k a
m1, Map a (Set k)
_)) (PrismMap (Map k a
m2, Map a (Set k)
_)) =
    let m :: Map k a
m = Map k a -> Map k a -> Map k a
forall k a b. Ord k => Map k a -> Map k b -> Map k a
Map.difference Map k a
m1 Map k a
m2
    in  (Map k a, Map a (Set k)) -> PrismMap k a
forall k a. (Map k a, Map a (Set k)) -> PrismMap k a
PrismMap (Map k a
m, [(k, a)] -> Map a (Set k)
forall k a. (Ord k, Ord a) => [(k, a)] -> Map a (Set k)
sndMapFromList ([(k, a)] -> Map a (Set k)) -> [(k, a)] -> Map a (Set k)
forall a b. (a -> b) -> a -> b
$ Map k a -> [(k, a)]
forall k a. Map k a -> [(k, a)]
Map.toList Map k a
m)