{-# LANGUAGE ExistentialQuantification, NamedFieldPuns, ScopedTypeVariables, RecordWildCards, ApplicativeDo #-}
{-|
Module : Data.MultiKeyedMap
Description : A map with possibly multiple keys per value
Maintainer : Profpatsch
Stability : experimental

Still very much experimental and missing lots of functions and testing.

Internally, a 'MKMap' is two maps, a @keyMap@ referencing an intermediate key
(whose type can be chosen freely and which is incremented sequentially), and
a @valueMap@ going from intermediate key to final value.

A correct implementation guarantees that

(1) the internal structure can’t be corrupted by operations declared safe
(2) adding and removing keys does not make values inaccessible
    (thus leaking memory) and doesn’t insert unnecessary values
-}
module Data.MultiKeyedMap
( MKMap
, at, (!)
, mkMKMap, fromList, toList
, insert
, flattenKeys, keys, values
) where

import qualified Data.Map.Strict as M
import Data.Monoid (All(..))
import Data.Semigroup ((<>))
import Data.Foldable (foldl')
import qualified Data.List.NonEmpty as NE
import Data.Proxy (Proxy(..))
import qualified Data.Tuple as Tuple
import GHC.Stack (HasCallStack)
import qualified Text.Show as Show

-- TODO: add time behaviour of functions to docstrings

-- | A `Map`-like structure where multiple keys can point
--   to the same value, with corresponding abstracted interface.
--
-- Internally, we use two maps connected by an intermediate key.
-- The intermediate key (@ik@) can be anything implementing
-- 'Ord' (for 'Map'), 'Bounded' (to get the first value)
-- and 'Enum' (for 'succ').
data MKMap k v = forall ik. (Ord ik, Enum ik)
              => MKMap
                 { ()
keyMap :: M.Map k ik
                 , ()
highestIk :: ik
                 , ()
valMap :: M.Map ik v }

-- TODO: is it possible without (Ord k)?
instance (Eq k, Ord k, Eq v) => Eq (MKMap k v) where
  -- TODO: not sure if that’s correct, add tests
  == :: MKMap k v -> MKMap k v -> Bool
(==) m1 :: MKMap k v
m1@(MKMap { keyMap :: ()
keyMap = Map k ik
km1
                 , valMap :: ()
valMap = Map ik v
vm1 })
       m2 :: MKMap k v
m2@(MKMap { keyMap :: ()
keyMap = Map k ik
km2
                 , valMap :: ()
valMap = Map ik v
vm2 })
    = All -> Bool
getAll (All -> Bool) -> All -> Bool
forall a b. (a -> b) -> a -> b
$ (Bool -> All) -> [Bool] -> All
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Bool -> All
All
        ([Bool] -> All) -> [Bool] -> All
forall a b. (a -> b) -> a -> b
$ let ks1 :: [k]
ks1 = Map k ik -> [k]
forall k a. Map k a -> [k]
M.keys Map k ik
km1 in
        -- shortcut if the length of the value map is not equal
        [ Map ik v -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Map ik v
vm1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Map ik v -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Map ik v
vm2
        -- the keys have to be equal (the lists are ascending)
        ,        [k]
ks1 [k] -> [k] -> Bool
forall a. Eq a => a -> a -> Bool
== Map k ik -> [k]
forall k a. Map k a -> [k]
M.keys Map k ik
km2 ]
        -- now test whether every key leads to the same value
        -- I wonder if there is a more efficient way?
        [Bool] -> [Bool] -> [Bool]
forall a. [a] -> [a] -> [a]
++ (k -> Bool) -> [k] -> [Bool]
forall a b. (a -> b) -> [a] -> [b]
map (\k
k -> MKMap k v
m1 MKMap k v -> k -> v
forall k v. (HasCallStack, Ord k) => MKMap k v -> k -> v
! k
k v -> v -> Bool
forall a. Eq a => a -> a -> Bool
== MKMap k v
m2 MKMap k v -> k -> v
forall k v. (HasCallStack, Ord k) => MKMap k v -> k -> v
! k
k) [k]
ks1
        -- we could test whether the values are equal,
        -- but if the implementation is correct they should
        -- all be reachable from the keys (TODO: add invariants)
   -- TODO: can (/=) be implemented more efficient than not.(==)?


instance Functor (MKMap k) where
  fmap :: (a -> b) -> MKMap k a -> MKMap k b
fmap a -> b
f (MKMap{ik
Map k ik
Map ik a
valMap :: Map ik a
highestIk :: ik
keyMap :: Map k ik
valMap :: ()
highestIk :: ()
keyMap :: ()
..}) = MKMap :: forall k v ik.
(Ord ik, Enum ik) =>
Map k ik -> ik -> Map ik v -> MKMap k v
MKMap { valMap :: Map ik b
valMap = (a -> b) -> Map ik a -> Map ik b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f Map ik a
valMap, ik
Map k ik
highestIk :: ik
keyMap :: Map k ik
highestIk :: ik
keyMap :: Map k ik
.. }
  {-# INLINE fmap #-}

-- TODO implement all functions Data.Map also implements for Foldable
instance Foldable (MKMap k) where
  foldMap :: (a -> m) -> MKMap k a -> m
foldMap a -> m
f (MKMap{ik
Map k ik
Map ik a
valMap :: Map ik a
highestIk :: ik
keyMap :: Map k ik
valMap :: ()
highestIk :: ()
keyMap :: ()
..}) = (a -> m) -> Map ik a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap a -> m
f Map ik a
valMap
  {-# INLINE foldMap #-}

instance Traversable (MKMap k) where
  traverse :: (a -> f b) -> MKMap k a -> f (MKMap k b)
traverse a -> f b
f (MKMap{ik
Map k ik
Map ik a
valMap :: Map ik a
highestIk :: ik
keyMap :: Map k ik
valMap :: ()
highestIk :: ()
keyMap :: ()
..}) = do
    Map ik b
val <- (a -> f b) -> Map ik a -> f (Map ik b)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse a -> f b
f Map ik a
valMap
    pure $ MKMap :: forall k v ik.
(Ord ik, Enum ik) =>
Map k ik -> ik -> Map ik v -> MKMap k v
MKMap{ valMap :: Map ik b
valMap=Map ik b
val, ik
Map k ik
highestIk :: ik
keyMap :: Map k ik
highestIk :: ik
keyMap :: Map k ik
.. }
  {-# INLINE traverse #-}

-- | Find value at key. Partial. See 'M.!'.
at :: (HasCallStack, Ord k) => MKMap k v -> k -> v
at :: MKMap k v -> k -> v
at MKMap{Map k ik
keyMap :: Map k ik
keyMap :: ()
keyMap, Map ik v
valMap :: Map ik v
valMap :: ()
valMap} k
k =  Map ik v
valMap Map ik v -> ik -> v
forall k a. Ord k => Map k a -> k -> a
M.! (Map k ik
keyMap Map k ik -> k -> ik
forall k a. Ord k => Map k a -> k -> a
M.! k
k)
-- | Operator alias of 'at'.
(!) :: (HasCallStack, Ord k) => MKMap k v -> k -> v
(!) = MKMap k v -> k -> v
forall k v. (HasCallStack, Ord k) => MKMap k v -> k -> v
at
{-# INLINABLE (!) #-}
{-# INLINABLE at #-}
infixl 9 !

-- | Create a 'MKMap' given a type for the internally used intermediate key.
mkMKMap :: forall k ik v. (Ord k, Ord ik, Enum ik, Bounded ik)
        => (Proxy ik) -- ^ type of intermediate key
        -> MKMap k v -- ^ new map
mkMKMap :: Proxy ik -> MKMap k v
mkMKMap Proxy ik
_ = Map k ik -> ik -> Map ik v -> MKMap k v
forall k v ik.
(Ord ik, Enum ik) =>
Map k ik -> ik -> Map ik v -> MKMap k v
MKMap Map k ik
forall a. Monoid a => a
mempty (ik
forall a. Bounded a => a
minBound :: ik) Map ik v
forall a. Monoid a => a
mempty
{-# INLINE mkMKMap #-}

instance (Show k, Show v) => Show (MKMap k v) where
  showsPrec :: Int -> MKMap k v -> ShowS
showsPrec Int
d MKMap k v
m = String -> ShowS
Show.showString String
"fromList " ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> [(NonEmpty k, v)] -> ShowS
forall a. Show a => Int -> a -> ShowS
showsPrec Int
d ([(NonEmpty k, v)] -> ShowS) -> [(NonEmpty k, v)] -> ShowS
forall a b. (a -> b) -> a -> b
$ MKMap k v -> [(NonEmpty k, v)]
forall k v. MKMap k v -> [(NonEmpty k, v)]
toList MKMap k v
m)

-- | Build a map from a list of key\/value pairs.
fromList :: forall ik k v. (Ord k, Ord ik, Enum ik, Bounded ik)
         => (Proxy ik) -- ^ type of intermediate key
         -> [(NE.NonEmpty k, v)] -- ^ list of @(key, value)@
         -> MKMap k v  -- ^ new map

-- TODO: it’s probably better to implement with M.fromList
fromList :: Proxy ik -> [(NonEmpty k, v)] -> MKMap k v
fromList Proxy ik
p = (MKMap k v -> (NonEmpty k, v) -> MKMap k v)
-> MKMap k v -> [(NonEmpty k, v)] -> MKMap k v
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\MKMap k v
m (NonEmpty k
ks, v
v) -> NonEmpty k -> v -> MKMap k v -> MKMap k v
forall k v. Ord k => NonEmpty k -> v -> MKMap k v -> MKMap k v
newVal NonEmpty k
ks v
v MKMap k v
m) (Proxy ik -> MKMap k v
forall k ik v.
(Ord k, Ord ik, Enum ik, Bounded ik) =>
Proxy ik -> MKMap k v
mkMKMap Proxy ik
p)

-- | Convert the map to a list of key\/value pairs.
toList :: MKMap k v -> [(NE.NonEmpty k, v)]
toList :: MKMap k v -> [(NonEmpty k, v)]
toList MKMap{Map k ik
keyMap :: Map k ik
keyMap :: ()
keyMap, Map ik v
valMap :: Map ik v
valMap :: ()
valMap} =
  ((ik, NonEmpty k) -> (NonEmpty k, v))
-> [(ik, NonEmpty k)] -> [(NonEmpty k, v)]
forall a b. (a -> b) -> [a] -> [b]
map ((ik -> v) -> (NonEmpty k, ik) -> (NonEmpty k, v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Map ik v
valMap Map ik v -> ik -> v
forall k a. Ord k => Map k a -> k -> a
M.!) ((NonEmpty k, ik) -> (NonEmpty k, v))
-> ((ik, NonEmpty k) -> (NonEmpty k, ik))
-> (ik, NonEmpty k)
-> (NonEmpty k, v)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ik, NonEmpty k) -> (NonEmpty k, ik)
forall a b. (a, b) -> (b, a)
Tuple.swap) ([(ik, NonEmpty k)] -> [(NonEmpty k, v)])
-> (Map k ik -> [(ik, NonEmpty k)])
-> Map k ik
-> [(NonEmpty k, v)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map ik (NonEmpty k) -> [(ik, NonEmpty k)]
forall k a. Map k a -> [(k, a)]
M.assocs (Map ik (NonEmpty k) -> [(ik, NonEmpty k)])
-> (Map k ik -> Map ik (NonEmpty k))
-> Map k ik
-> [(ik, NonEmpty k)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map k ik -> Map ik (NonEmpty k)
forall k ik. (Ord ik, Enum ik) => Map k ik -> Map ik (NonEmpty k)
aggregateIk (Map k ik -> [(NonEmpty k, v)]) -> Map k ik -> [(NonEmpty k, v)]
forall a b. (a -> b) -> a -> b
$ Map k ik
keyMap
    where
      aggregateIk :: forall k ik. (Ord ik, Enum ik)
                  => M.Map k ik
                  -> M.Map ik (NE.NonEmpty k)
      aggregateIk :: Map k ik -> Map ik (NonEmpty k)
aggregateIk = (Map ik (NonEmpty k) -> k -> ik -> Map ik (NonEmpty k))
-> Map ik (NonEmpty k) -> Map k ik -> Map ik (NonEmpty k)
forall a k b. (a -> k -> b -> a) -> a -> Map k b -> a
M.foldlWithKey
            (\Map ik (NonEmpty k)
m k
k ik
ik -> (NonEmpty k -> NonEmpty k -> NonEmpty k)
-> ik -> NonEmpty k -> Map ik (NonEmpty k) -> Map ik (NonEmpty k)
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWith NonEmpty k -> NonEmpty k -> NonEmpty k
forall a. Semigroup a => a -> a -> a
(<>) ik
ik (k -> NonEmpty k
forall (f :: * -> *) a. Applicative f => a -> f a
pure k
k) Map ik (NonEmpty k)
m) Map ik (NonEmpty k)
forall a. Monoid a => a
mempty

-- | “Unlink” keys that are pointing to the same value.
--
-- Returns a normal map.
flattenKeys :: (Ord k) => MKMap k v -> M.Map k v
flattenKeys :: MKMap k v -> Map k v
flattenKeys MKMap{Map k ik
keyMap :: Map k ik
keyMap :: ()
keyMap, Map ik v
valMap :: Map ik v
valMap :: ()
valMap} =
  (Map k v -> k -> ik -> Map k v) -> Map k v -> Map k ik -> Map k v
forall a k b. (a -> k -> b -> a) -> a -> Map k b -> a
M.foldlWithKey' (\Map k v
m k
k ik
ik -> k -> v -> Map k v -> Map k v
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
k (Map ik v
valMap Map ik v -> ik -> v
forall k a. Ord k => Map k a -> k -> a
M.! ik
ik) Map k v
m) Map k v
forall a. Monoid a => a
mempty Map k ik
keyMap

-- | Return a list of all keys.
keys :: (Ord k) => MKMap k v -> [k]
keys :: MKMap k v -> [k]
keys = Map k v -> [k]
forall k a. Map k a -> [k]
M.keys (Map k v -> [k]) -> (MKMap k v -> Map k v) -> MKMap k v -> [k]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MKMap k v -> Map k v
forall k v. Ord k => MKMap k v -> Map k v
flattenKeys

-- | Return a list of all values.
values :: MKMap k v -> [v]
values :: MKMap k v -> [v]
values (MKMap Map k ik
_ ik
_ Map ik v
valMap) = Map ik v -> [v]
forall k a. Map k a -> [a]
M.elems Map ik v
valMap

-- TODO: this is like normal insert, it doesn’t search if the value
-- already exists (where it might want to add the key instead).
-- Of course that would be O(n) in the naive implementation.
-- In that case the keyMap should probably be changed to a bimap.
-- also, naming
-- | Equivalent to 'M.insert', if the key doesn’t exist a new
-- singleton key is added.
insert :: (Ord k) => k -> v -> MKMap k v -> MKMap k v
insert :: k -> v -> MKMap k v -> MKMap k v
insert k
k v
v m :: MKMap k v
m@MKMap{Map k ik
keyMap :: Map k ik
keyMap :: ()
keyMap, ik
highestIk :: ik
highestIk :: ()
highestIk, Map ik v
valMap :: Map ik v
valMap :: ()
valMap} =
  MKMap k v -> (ik -> MKMap k v) -> Maybe ik -> MKMap k v
forall b a. b -> (a -> b) -> Maybe a -> b
maybe MKMap k v
ins ik -> MKMap k v
upd (Maybe ik -> MKMap k v) -> Maybe ik -> MKMap k v
forall a b. (a -> b) -> a -> b
$ k -> Map k ik -> Maybe ik
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup k
k Map k ik
keyMap
  where
    ins :: MKMap k v
ins = NonEmpty k -> v -> MKMap k v -> MKMap k v
forall k v. Ord k => NonEmpty k -> v -> MKMap k v -> MKMap k v
newVal (k -> NonEmpty k
forall (f :: * -> *) a. Applicative f => a -> f a
pure k
k) v
v MKMap k v
m
    upd :: ik -> MKMap k v
upd ik
ik = MKMap :: forall k v ik.
(Ord ik, Enum ik) =>
Map k ik -> ik -> Map ik v -> MKMap k v
MKMap { Map k ik
keyMap :: Map k ik
keyMap :: Map k ik
keyMap, ik
highestIk :: ik
highestIk :: ik
highestIk, valMap :: Map ik v
valMap = ik -> v -> Map ik v -> Map ik v
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert ik
ik v
v Map ik v
valMap }


-- | Helper, assumes there is no such value already.
-- Will leak space otherwise!
--
-- Insert every key into the keyMap, increase the intermediate counter,
-- insert the value at new intermediate counter.
-- Overwrites all already existing keys!
newVal :: (Ord k) => NE.NonEmpty k -> v -> MKMap k v -> MKMap k v
newVal :: NonEmpty k -> v -> MKMap k v -> MKMap k v
newVal NonEmpty k
ks v
v MKMap{Map k ik
keyMap :: Map k ik
keyMap :: ()
keyMap, ik
highestIk :: ik
highestIk :: ()
highestIk, Map ik v
valMap :: Map ik v
valMap :: ()
valMap} =
  MKMap :: forall k v ik.
(Ord ik, Enum ik) =>
Map k ik -> ik -> Map ik v -> MKMap k v
MKMap { keyMap :: Map k ik
keyMap = (Map k ik -> k -> Map k ik) -> Map k ik -> NonEmpty k -> Map k ik
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Map k ik
m k
k -> k -> ik -> Map k ik -> Map k ik
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
k ik
next Map k ik
m) Map k ik
keyMap NonEmpty k
ks
        , highestIk :: ik
highestIk = ik
next
        , valMap :: Map ik v
valMap = ik -> v -> Map ik v -> Map ik v
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert ik
next v
v Map ik v
valMap }
  where next :: ik
next = ik -> ik
forall a. Enum a => a -> a
succ ik
highestIk