content: {-# 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.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
    (==) m1@(MKMap { keyMap = km1
                   , valMap = vm1 })
         m2@(MKMap { keyMap = km2
                   , valMap = vm2 })
      = getAll $ foldMap All
          $ let ks1 = M.keys km1 in
          -- shortcut if the length of the value map is not equal
          [ length vm1 == length vm2
          -- the keys have to be equal (the lists are ascending)
          ,        ks1 == M.keys km2 ]
          -- now test whether every key leads to the same value
          -- I wonder if there is a more efficient way?
          ++ map (\k -> m1 ! k == m2 ! 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 f (MKMap{..}) = MKMap { valMap = fmap f valMap, .. }
    {-# INLINE fmap #-}
  
  -- TODO implement all functions Data.Map also implements for Foldable
  instance Foldable (MKMap k) where
    foldMap f (MKMap{..}) = foldMap f valMap
    {-# INLINE foldMap #-}
  
  instance Traversable (MKMap k) where
    traverse f (MKMap{..}) = do
      val <- traverse f valMap
      pure $ MKMap{ valMap=val, .. }
    {-# INLINE traverse #-}
  
  -- | Find value at key. Partial. See 'M.!'.
  at :: (HasCallStack, Ord k) => MKMap k v -> k -> v
  at MKMap{keyMap, valMap} k =  valMap M.! (keyMap M.! k)
  -- | Operator alias of 'at'.
  (!) :: (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 _ = MKMap mempty (minBound :: ik) mempty
  {-# INLINE mkMKMap #-}
  
  instance (Show k, Show v) => Show (MKMap k v) where
    showsPrec d m = Show.showString "fromList " . (showsPrec d $ toList 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 p = foldl' (\m (ks, v) -> newVal ks v m) (mkMKMap p)
  
  -- | Convert the map to a list of key\/value pairs.
  toList :: MKMap k v -> [(NE.NonEmpty k, v)]
  toList MKMap{keyMap, valMap} =
    map (fmap (valMap M.!) . Tuple.swap) . M.assocs . aggregateIk $ keyMap
      where
        aggregateIk :: forall k ik. (Ord ik, Enum ik)
                    => M.Map k ik
                    -> M.Map ik (NE.NonEmpty k)
        aggregateIk = M.foldlWithKey
              (\m k ik -> M.insertWith (<>) ik (pure k) m) 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{keyMap, valMap} =
    M.foldlWithKey' (\m k ik -> M.insert k (valMap M.! ik) m) mempty keyMap
  
  -- | Return a list of all keys.
  keys :: (Ord k) => MKMap k v -> [k]
  keys = M.keys . flattenKeys
  
  -- | Return a list of all values.
  values :: MKMap k v -> [v]
  values (MKMap _ _ valMap) = M.elems 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 m@MKMap{keyMap, highestIk, valMap} =
    case M.lookup k keyMap of
      Nothing -> ins
      Just ik -> upd ik
    where
      ins = newVal (pure k) v m
      upd ik = MKMap { keyMap, highestIk, valMap = M.insert 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 ks v MKMap{keyMap, highestIk, valMap} =
    MKMap { keyMap = foldl' (\m k -> M.insert k next m) keyMap ks
          , highestIk = next
          , valMap = M.insert next v valMap }
    where next = succ highestIk
  