content: {-# LANGUAGE GeneralizedNewtypeDeriving #-}
  module Data.ShareMap
    ( ShareMap
    , empty
    , toHashMap
    , insertWith
    , map
    , mergeKeysWith
    ) where
  
  import           Data.Hashable (Hashable)
  import           Data.HashMap.Lazy (HashMap)
  import qualified Data.HashMap.Lazy as HashMap
  import           Data.HashSet (HashSet)
  import qualified Data.HashSet as HashSet
  import           Data.Maybe (fromMaybe)
  import           Prelude hiding (map)
  
  -- | A HashMap that can share the values of some entries
  --
  -- If two keys @k1@ and @k2@ are mapped to single @v@, updating
  -- the entry for @k1@ also updates the entry for @k2@ and viceversa.
  --
  -- The user of the map is responsible for indicating which keys are
  -- going to share their values.
  --
  -- Keys can be updated with 'shareMapInsertWith' and 'mergeKeysWith'.
  data ShareMap k v = ShareMap
    { -- | @(k, v)@ pairs in the map.
      --
      -- Contains at least an entry for each key in the values of
      -- of 'shareMap'.
      unsharedMap :: HashMap (InternalKey k) v
    , -- | If @k1@ is mapped to @k2@, then both keys are considered
      -- associated to the value of @k2@ in 'unsharedMap'.
      --
      -- Contains an entry for each key in the 'ShareMap'.
      shareMap :: ReversibleMap k (InternalKey k)
    }
    deriving Show
  
  -- | This are the only keys that can be used in internal maps
  newtype InternalKey k = InternalKey k
    deriving (Eq, Hashable)
  
  instance Show k => Show (InternalKey k) where
    show (InternalKey k) = show k
  
  empty :: ShareMap k v
  empty = ShareMap HashMap.empty emptyReversibleMap
  
  toHashMap :: (Hashable k, Eq k) => ShareMap k v -> HashMap k v
  toHashMap sm =
    HashMap.foldlWithKey' expand HashMap.empty (directMap $ shareMap sm)
    where
      expand m k k' =
        maybe m (\v -> HashMap.insert k v m) (HashMap.lookup k' $ unsharedMap sm)
  
  -- | @insertWith f k v m@ is the map @m@ plus key @k@ being associated to
  -- value @v@.
  --
  -- If @k@ is present in @m@, then @k@ and any other key sharing its value
  -- will be associated to @f v (m ! k)@.
  --
  insertWith
    :: (Hashable k, Eq k)
    => (v -> v -> v)
    -> k
    -> v
    -> ShareMap k v
    -> ShareMap k v
  insertWith f k v sm =
    case HashMap.lookup k $ directMap $ shareMap sm of
      Just k' -> sm
        { unsharedMap = HashMap.insertWith f k' v (unsharedMap sm)
        }
      Nothing -> ShareMap
        { unsharedMap = HashMap.insertWith f (InternalKey k) v (unsharedMap sm)
        , shareMap = insertReversibleMap k (InternalKey k) (shareMap sm)
        }
  
  -- | @mergeKeysWith f k0 k1 m@ updates the @k0@ value to @f (m ! k0) (m ! k1)@
  -- and @k1@ shares the value with @k0@.
  --
  -- If @k0@ and @k1@ are already sharing their values in @m@, or both keys are
  -- missing, this operation returns @m@ unmodified.
  --
  -- If only one of the keys is present, the other key is associated with the
  -- existing value.
  mergeKeysWith
    :: (Hashable k, Eq k)
    => (v -> v -> v)
    -> k
    -> k
    -> ShareMap k v
    -> ShareMap k v
  mergeKeysWith f k0 k1 sm | k0 /= k1 =
    case lookupReversibleMap k1 (shareMap sm) of
      Just k1' | InternalKey k0 /= k1' -> case HashMap.lookup k1' (unsharedMap sm) of
        Just v1 -> case lookupReversibleMap k0 (shareMap sm) of
          Just k0' | k0' /= k1' ->
            ShareMap
              { unsharedMap = HashMap.insertWith (flip f) k0' v1 (unsharedMap sm)
              , shareMap = -- Any values pointing to k1 are redirected to k0':
                  HashSet.foldl' (\m k -> insertReversibleMap k k0' m) (shareMap sm) $
                  reverseLookup k1' $ shareMap sm
              }
          Nothing ->
            sm { shareMap = insertReversibleMap k0 k1' (shareMap sm) }
          _ ->
            sm
        Nothing -> error "mergeKeysWith: broken invariant: unexpected missing key in unsharedMap"
      Nothing ->
        case HashMap.lookup k0 (directMap $ shareMap sm) of
          Just k0' ->
            sm { shareMap = insertReversibleMap k1 k0' (shareMap sm) }
          Nothing ->
            sm
      _ ->
        sm
  mergeKeysWith _ _ _ sm = sm
  
  map :: (a -> b) -> ShareMap k a -> ShareMap k b
  map f sm = sm { unsharedMap = HashMap.map f (unsharedMap sm) }
  
  -- | A map with an efficient 'reverseLookup'
  data ReversibleMap k v = ReversibleMap
    { directMap :: HashMap k v
    , -- |
      -- > forall (v, ks) in reversedMap.
      -- >   forall k in ks.
      -- >     (k, v) is in directMap
      reversedMap :: HashMap v (HashSet k)
    }
    deriving Show
  
  emptyReversibleMap :: ReversibleMap k v
  emptyReversibleMap = ReversibleMap HashMap.empty HashMap.empty
  
  insertReversibleMap
    :: (Hashable k, Eq k, Hashable v, Eq v)
    => k
    -> v
    -> ReversibleMap k v
    -> ReversibleMap k v
  insertReversibleMap k v rm = ReversibleMap
    { directMap = HashMap.insert k v (directMap rm)
    , reversedMap =
        let m' = case HashMap.lookup k (directMap rm) of
              Nothing -> reversedMap rm
              Just oldv -> HashMap.adjust (HashSet.delete k) oldv (reversedMap rm)
         in HashMap.insertWith HashSet.union v (HashSet.singleton k) m'
    }
  
  reverseLookup :: (Hashable v, Eq v) => v -> ReversibleMap k v -> HashSet k
  reverseLookup v rm = fromMaybe HashSet.empty $ HashMap.lookup v (reversedMap rm)
  
  lookupReversibleMap :: (Hashable k, Eq k) => k -> ReversibleMap k v -> Maybe v
  lookupReversibleMap k rm = HashMap.lookup k (directMap rm)
  