{-|
Module      : KMonad.Util.MultiMap
Description : A `k -> Set v` mapping, with reversing utilities
Copyright   : (c) David Janssen, 2019
License     : MIT
Maintainer  : janssen.dhj@gmail.com
Stability   : experimental
Portability : portable

This datastructure represents a `k -> Set v` mapping: that is to say, each key
can have multiple values (but no duplicates). Additionally, we provide some
operations to reverse this mapping.

In KMonad we use this exclusively to easily define multiple names for the same
'KMonad.Keyboard.Keycode' in a reversible manner.

-}
module KMonad.Util.MultiMap
  ( -- * Types
    -- $typ
    MultiMap
  , mkMultiMap
  , fromSingletons

    -- * Operations on MultiMaps
    -- $ops
  , itemed
  , reverse
  )
where

import KMonad.Prelude hiding (reverse)

import qualified RIO.HashMap as M
import qualified RIO.HashSet as S

--------------------------------------------------------------------------------
-- $typ

-- | All the type constraints required for something to function as a MultiMap
type CanMM k v = (Eq k, Ord v, Hashable k, Hashable v)

-- | The 'MultiMap', which describes a one to many (unique) mapping
newtype MultiMap k v = MultiMap { forall k v. MultiMap k v -> HashMap k (HashSet v)
_unMM :: M.HashMap k (S.HashSet v) }
  deriving Int -> MultiMap k v -> ShowS
[MultiMap k v] -> ShowS
MultiMap k v -> String
(Int -> MultiMap k v -> ShowS)
-> (MultiMap k v -> String)
-> ([MultiMap k v] -> ShowS)
-> Show (MultiMap k v)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall k v. (Show k, Show v) => Int -> MultiMap k v -> ShowS
forall k v. (Show k, Show v) => [MultiMap k v] -> ShowS
forall k v. (Show k, Show v) => MultiMap k v -> String
$cshowsPrec :: forall k v. (Show k, Show v) => Int -> MultiMap k v -> ShowS
showsPrec :: Int -> MultiMap k v -> ShowS
$cshow :: forall k v. (Show k, Show v) => MultiMap k v -> String
show :: MultiMap k v -> String
$cshowList :: forall k v. (Show k, Show v) => [MultiMap k v] -> ShowS
showList :: [MultiMap k v] -> ShowS
Show
makeLenses ''MultiMap

instance (CanMM k v) => Semigroup (MultiMap k v) where
  (MultiMap HashMap k (HashSet v)
a) <> :: MultiMap k v -> MultiMap k v -> MultiMap k v
<> (MultiMap HashMap k (HashSet v)
b) = HashMap k (HashSet v) -> MultiMap k v
forall k v. HashMap k (HashSet v) -> MultiMap k v
MultiMap (HashMap k (HashSet v) -> MultiMap k v)
-> HashMap k (HashSet v) -> MultiMap k v
forall a b. (a -> b) -> a -> b
$ (HashSet v -> HashSet v -> HashSet v)
-> HashMap k (HashSet v)
-> HashMap k (HashSet v)
-> HashMap k (HashSet v)
forall k v.
Eq k =>
(v -> v -> v) -> HashMap k v -> HashMap k v -> HashMap k v
M.unionWith HashSet v -> HashSet v -> HashSet v
forall a. Semigroup a => a -> a -> a
(<>) HashMap k (HashSet v)
a HashMap k (HashSet v)
b
instance (CanMM k v) => Monoid (MultiMap k v) where
  mempty :: MultiMap k v
mempty = HashMap k (HashSet v) -> MultiMap k v
forall k v. HashMap k (HashSet v) -> MultiMap k v
MultiMap HashMap k (HashSet v)
forall k v. HashMap k v
M.empty

type instance Index   (MultiMap k v) = k
type instance IxValue (MultiMap k v) = S.HashSet v

instance CanMM k v => Ixed (MultiMap k v)
instance CanMM k v => At (MultiMap k v) where
  at :: Index (MultiMap k v)
-> Lens' (MultiMap k v) (Maybe (IxValue (MultiMap k v)))
at Index (MultiMap k v)
k = (HashMap k (HashSet v) -> f (HashMap k (HashSet v)))
-> MultiMap k v -> f (MultiMap k v)
forall k v k v (p :: * -> * -> *) (f :: * -> *).
(Profunctor p, Functor f) =>
p (HashMap k (HashSet v)) (f (HashMap k (HashSet v)))
-> p (MultiMap k v) (f (MultiMap k v))
unMM ((HashMap k (HashSet v) -> f (HashMap k (HashSet v)))
 -> MultiMap k v -> f (MultiMap k v))
-> ((Maybe (HashSet v) -> f (Maybe (HashSet v)))
    -> HashMap k (HashSet v) -> f (HashMap k (HashSet v)))
-> (Maybe (HashSet v) -> f (Maybe (HashSet v)))
-> MultiMap k v
-> f (MultiMap k v)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Index (HashMap k (HashSet v))
-> Lens'
     (HashMap k (HashSet v)) (Maybe (IxValue (HashMap k (HashSet v))))
forall m. At m => Index m -> Lens' m (Maybe (IxValue m))
at Index (HashMap k (HashSet v))
Index (MultiMap k v)
k

-- | Create a new multimap from a foldable of (k, foldable v) pairs.
mkMultiMap :: (Foldable t1, Foldable t2, CanMM k v)
  => t1 (k, t2 v) -> MultiMap k v
mkMultiMap :: forall (t1 :: * -> *) (t2 :: * -> *) k v.
(Foldable t1, Foldable t2, CanMM k v) =>
t1 (k, t2 v) -> MultiMap k v
mkMultiMap = ((k, t2 v) -> MultiMap k v) -> t1 (k, t2 v) -> MultiMap k v
forall m a. Monoid m => (a -> m) -> t1 a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap
  ( HashMap k (HashSet v) -> MultiMap k v
forall k v. HashMap k (HashSet v) -> MultiMap k v
MultiMap
  (HashMap k (HashSet v) -> MultiMap k v)
-> ((k, t2 v) -> HashMap k (HashSet v))
-> (k, t2 v)
-> MultiMap k v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (k -> HashSet v -> HashMap k (HashSet v))
-> (k, HashSet v) -> HashMap k (HashSet v)
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry k -> HashSet v -> HashMap k (HashSet v)
forall k v. Hashable k => k -> v -> HashMap k v
M.singleton
  ((k, HashSet v) -> HashMap k (HashSet v))
-> ((k, t2 v) -> (k, HashSet v))
-> (k, t2 v)
-> HashMap k (HashSet v)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ASetter (k, t2 v) (k, HashSet v) (t2 v) (HashSet v)
-> (t2 v -> HashSet v) -> (k, t2 v) -> (k, HashSet v)
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
over ASetter (k, t2 v) (k, HashSet v) (t2 v) (HashSet v)
forall s t a b. Field2 s t a b => Lens s t a b
Lens (k, t2 v) (k, HashSet v) (t2 v) (HashSet v)
_2 ([v] -> HashSet v
forall a. (Eq a, Hashable a) => [a] -> HashSet a
S.fromList ([v] -> HashSet v) -> (t2 v -> [v]) -> t2 v -> HashSet v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. t2 v -> [v]
forall a. t2 a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList)
  )

-- | Create a new multimap from a foldable of (k, v) pairs
fromSingletons :: (Foldable t, CanMM k v)
  => t (k, v) -> MultiMap k v
fromSingletons :: forall (t :: * -> *) k v.
(Foldable t, CanMM k v) =>
t (k, v) -> MultiMap k v
fromSingletons = [(k, [v])] -> MultiMap k v
forall (t1 :: * -> *) (t2 :: * -> *) k v.
(Foldable t1, Foldable t2, CanMM k v) =>
t1 (k, t2 v) -> MultiMap k v
mkMultiMap ([(k, [v])] -> MultiMap k v)
-> (t (k, v) -> [(k, [v])]) -> t (k, v) -> MultiMap k v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((k, v) -> (k, [v])) -> [(k, v)] -> [(k, [v])]
forall a b. (a -> b) -> [a] -> [b]
map (ASetter (k, v) (k, [v]) v [v] -> (v -> [v]) -> (k, v) -> (k, [v])
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
over ASetter (k, v) (k, [v]) v [v]
forall s t a b. Field2 s t a b => Lens s t a b
Lens (k, v) (k, [v]) v [v]
_2 (v -> [v] -> [v]
forall a. a -> [a] -> [a]
:[])) ([(k, v)] -> [(k, [v])])
-> (t (k, v) -> [(k, v)]) -> t (k, v) -> [(k, [v])]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. t (k, v) -> [(k, v)]
forall a. t a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList



--------------------------------------------------------------------------------
-- $ops

-- | A fold over all the (k, v) pairs in a 'MultiMap'
itemed :: (CanMM k v) => Fold (MultiMap k v) (k, v)
itemed :: forall k v. CanMM k v => Fold (MultiMap k v) (k, v)
itemed = (MultiMap k v -> [(k, v)]) -> Fold (MultiMap k v) (k, v)
forall (f :: * -> *) s a. Foldable f => (s -> f a) -> Fold s a
folding ((MultiMap k v -> [(k, v)]) -> Fold (MultiMap k v) (k, v))
-> (MultiMap k v -> [(k, v)]) -> Fold (MultiMap k v) (k, v)
forall a b. (a -> b) -> a -> b
$ \MultiMap k v
m -> MultiMap k v
m MultiMap k v
-> IndexedGetting k (Endo [(k, v)]) (MultiMap k v) v -> [(k, v)]
forall s i a. s -> IndexedGetting i (Endo [(i, a)]) s a -> [(i, a)]
^@.. (HashMap k (HashSet v)
 -> Const (Endo [(k, v)]) (HashMap k (HashSet v)))
-> MultiMap k v -> Const (Endo [(k, v)]) (MultiMap k v)
forall k v k v (p :: * -> * -> *) (f :: * -> *).
(Profunctor p, Functor f) =>
p (HashMap k (HashSet v)) (f (HashMap k (HashSet v)))
-> p (MultiMap k v) (f (MultiMap k v))
unMM ((HashMap k (HashSet v)
  -> Const (Endo [(k, v)]) (HashMap k (HashSet v)))
 -> MultiMap k v -> Const (Endo [(k, v)]) (MultiMap k v))
-> (Indexed k v (Const (Endo [(k, v)]) v)
    -> HashMap k (HashSet v)
    -> Const (Endo [(k, v)]) (HashMap k (HashSet v)))
-> IndexedGetting k (Endo [(k, v)]) (MultiMap k v) v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Indexed k (HashSet v) (Const (Endo [(k, v)]) (HashSet v))
-> HashMap k (HashSet v)
-> Const (Endo [(k, v)]) (HashMap k (HashSet v))
forall i (f :: * -> *) a.
FoldableWithIndex i f =>
IndexedFold i (f a) a
IndexedFold k (HashMap k (HashSet v)) (HashSet v)
ifolded (Indexed k (HashSet v) (Const (Endo [(k, v)]) (HashSet v))
 -> HashMap k (HashSet v)
 -> Const (Endo [(k, v)]) (HashMap k (HashSet v)))
-> ((v -> Const (Endo [(k, v)]) v)
    -> HashSet v -> Const (Endo [(k, v)]) (HashSet v))
-> Indexed k v (Const (Endo [(k, v)]) v)
-> HashMap k (HashSet v)
-> Const (Endo [(k, v)]) (HashMap k (HashSet v))
forall i (p :: * -> * -> *) s t r a b.
Indexable i p =>
(Indexed i s t -> r) -> ((a -> b) -> s -> t) -> p a b -> r
<. (v -> Const (Endo [(k, v)]) v)
-> HashSet v -> Const (Endo [(k, v)]) (HashSet v)
forall (f :: * -> *) a. Foldable f => IndexedFold Int (f a) a
IndexedFold Int (HashSet v) v
folded

-- | Reverse a MultiMap. Note: this is not necessarily a lossless conversion.
reverse :: (CanMM k v, CanMM v k) => MultiMap k v -> MultiMap v k
reverse :: forall k v. (CanMM k v, CanMM v k) => MultiMap k v -> MultiMap v k
reverse MultiMap k v
m = [(v, [k])] -> MultiMap v k
forall (t1 :: * -> *) (t2 :: * -> *) k v.
(Foldable t1, Foldable t2, CanMM k v) =>
t1 (k, t2 v) -> MultiMap k v
mkMultiMap ([(v, [k])] -> MultiMap v k) -> [(v, [k])] -> MultiMap v k
forall a b. (a -> b) -> a -> b
$ MultiMap k v
m MultiMap k v
-> Getting (Endo [(v, [k])]) (MultiMap k v) (v, [k]) -> [(v, [k])]
forall s a. s -> Getting (Endo [a]) s a -> [a]
^.. ((k, v) -> Const (Endo [(v, [k])]) (k, v))
-> MultiMap k v -> Const (Endo [(v, [k])]) (MultiMap k v)
forall k v. CanMM k v => Fold (MultiMap k v) (k, v)
Fold (MultiMap k v) (k, v)
itemed (((k, v) -> Const (Endo [(v, [k])]) (k, v))
 -> MultiMap k v -> Const (Endo [(v, [k])]) (MultiMap k v))
-> (((v, [k]) -> Const (Endo [(v, [k])]) (v, [k]))
    -> (k, v) -> Const (Endo [(v, [k])]) (k, v))
-> Getting (Endo [(v, [k])]) (MultiMap k v) (v, [k])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((v, k) -> Const (Endo [(v, [k])]) (v, k))
-> (k, v) -> Const (Endo [(v, [k])]) (k, v)
forall (p :: * -> * -> *) a b c d.
Swap p =>
Iso (p a b) (p c d) (p b a) (p d c)
Iso (k, v) (k, v) (v, k) (v, k)
swapped (((v, k) -> Const (Endo [(v, [k])]) (v, k))
 -> (k, v) -> Const (Endo [(v, [k])]) (k, v))
-> (((v, [k]) -> Const (Endo [(v, [k])]) (v, [k]))
    -> (v, k) -> Const (Endo [(v, [k])]) (v, k))
-> ((v, [k]) -> Const (Endo [(v, [k])]) (v, [k]))
-> (k, v)
-> Const (Endo [(v, [k])]) (k, v)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((v, k) -> (v, [k]))
-> ((v, [k]) -> Const (Endo [(v, [k])]) (v, [k]))
-> (v, k)
-> Const (Endo [(v, [k])]) (v, k)
forall (p :: * -> * -> *) (f :: * -> *) s a.
(Profunctor p, Contravariant f) =>
(s -> a) -> Optic' p f s a
to (ASetter (v, k) (v, [k]) k [k] -> (k -> [k]) -> (v, k) -> (v, [k])
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
over ASetter (v, k) (v, [k]) k [k]
forall s t a b. Field2 s t a b => Lens s t a b
Lens (v, k) (v, [k]) k [k]
_2 (k -> [k] -> [k]
forall a. a -> [a] -> [a]
:[]))