{-# 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'.
    forall k v. ShareMap k v -> HashMap (InternalKey k) v
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'.
    forall k v. ShareMap k v -> ReversibleMap k (InternalKey k)
shareMap :: ReversibleMap k (InternalKey k)
  }
  deriving Int -> ShareMap k v -> ShowS
[ShareMap k v] -> ShowS
ShareMap k v -> String
(Int -> ShareMap k v -> ShowS)
-> (ShareMap k v -> String)
-> ([ShareMap k v] -> ShowS)
-> Show (ShareMap k v)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall k v. (Show k, Show v) => Int -> ShareMap k v -> ShowS
forall k v. (Show k, Show v) => [ShareMap k v] -> ShowS
forall k v. (Show k, Show v) => ShareMap k v -> String
$cshowsPrec :: forall k v. (Show k, Show v) => Int -> ShareMap k v -> ShowS
showsPrec :: Int -> ShareMap k v -> ShowS
$cshow :: forall k v. (Show k, Show v) => ShareMap k v -> String
show :: ShareMap k v -> String
$cshowList :: forall k v. (Show k, Show v) => [ShareMap k v] -> ShowS
showList :: [ShareMap k v] -> ShowS
Show

-- | This are the only keys that can be used in internal maps
newtype InternalKey k = InternalKey k
  deriving (InternalKey k -> InternalKey k -> Bool
(InternalKey k -> InternalKey k -> Bool)
-> (InternalKey k -> InternalKey k -> Bool) -> Eq (InternalKey k)
forall k. Eq k => InternalKey k -> InternalKey k -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall k. Eq k => InternalKey k -> InternalKey k -> Bool
== :: InternalKey k -> InternalKey k -> Bool
$c/= :: forall k. Eq k => InternalKey k -> InternalKey k -> Bool
/= :: InternalKey k -> InternalKey k -> Bool
Eq, Eq (InternalKey k)
Eq (InternalKey k) =>
(Int -> InternalKey k -> Int)
-> (InternalKey k -> Int) -> Hashable (InternalKey k)
Int -> InternalKey k -> Int
InternalKey k -> Int
forall a. Eq a => (Int -> a -> Int) -> (a -> Int) -> Hashable a
forall k. Hashable k => Eq (InternalKey k)
forall k. Hashable k => Int -> InternalKey k -> Int
forall k. Hashable k => InternalKey k -> Int
$chashWithSalt :: forall k. Hashable k => Int -> InternalKey k -> Int
hashWithSalt :: Int -> InternalKey k -> Int
$chash :: forall k. Hashable k => InternalKey k -> Int
hash :: InternalKey k -> Int
Hashable)

instance Show k => Show (InternalKey k) where
  show :: InternalKey k -> String
show (InternalKey k
k) = k -> String
forall a. Show a => a -> String
show k
k

empty :: ShareMap k v
empty :: forall k v. ShareMap k v
empty = HashMap (InternalKey k) v
-> ReversibleMap k (InternalKey k) -> ShareMap k v
forall k v.
HashMap (InternalKey k) v
-> ReversibleMap k (InternalKey k) -> ShareMap k v
ShareMap HashMap (InternalKey k) v
forall k v. HashMap k v
HashMap.empty ReversibleMap k (InternalKey k)
forall k v. ReversibleMap k v
emptyReversibleMap

toHashMap :: (Hashable k, Eq k) => ShareMap k v -> HashMap k v
toHashMap :: forall k v. (Hashable k, Eq k) => ShareMap k v -> HashMap k v
toHashMap ShareMap k v
sm =
  (HashMap k v -> k -> InternalKey k -> HashMap k v)
-> HashMap k v -> HashMap k (InternalKey k) -> HashMap k v
forall a k v. (a -> k -> v -> a) -> a -> HashMap k v -> a
HashMap.foldlWithKey' HashMap k v -> k -> InternalKey k -> HashMap k v
forall {k}.
Hashable k =>
HashMap k v -> k -> InternalKey k -> HashMap k v
expand HashMap k v
forall k v. HashMap k v
HashMap.empty (ReversibleMap k (InternalKey k) -> HashMap k (InternalKey k)
forall k v. ReversibleMap k v -> HashMap k v
directMap (ReversibleMap k (InternalKey k) -> HashMap k (InternalKey k))
-> ReversibleMap k (InternalKey k) -> HashMap k (InternalKey k)
forall a b. (a -> b) -> a -> b
$ ShareMap k v -> ReversibleMap k (InternalKey k)
forall k v. ShareMap k v -> ReversibleMap k (InternalKey k)
shareMap ShareMap k v
sm)
  where
    expand :: HashMap k v -> k -> InternalKey k -> HashMap k v
expand HashMap k v
m k
k InternalKey k
k' =
      HashMap k v -> (v -> HashMap k v) -> Maybe v -> HashMap k v
forall b a. b -> (a -> b) -> Maybe a -> b
maybe HashMap k v
m (\v
v -> k -> v -> HashMap k v -> HashMap k v
forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
HashMap.insert k
k v
v HashMap k v
m) (InternalKey k -> HashMap (InternalKey k) v -> Maybe v
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HashMap.lookup InternalKey k
k' (HashMap (InternalKey k) v -> Maybe v)
-> HashMap (InternalKey k) v -> Maybe v
forall a b. (a -> b) -> a -> b
$ ShareMap k v -> HashMap (InternalKey k) v
forall k v. ShareMap k v -> HashMap (InternalKey k) v
unsharedMap ShareMap k v
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 :: forall k v.
(Hashable k, Eq k) =>
(v -> v -> v) -> k -> v -> ShareMap k v -> ShareMap k v
insertWith v -> v -> v
f k
k v
v ShareMap k v
sm =
  case k -> HashMap k (InternalKey k) -> Maybe (InternalKey k)
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HashMap.lookup k
k (HashMap k (InternalKey k) -> Maybe (InternalKey k))
-> HashMap k (InternalKey k) -> Maybe (InternalKey k)
forall a b. (a -> b) -> a -> b
$ ReversibleMap k (InternalKey k) -> HashMap k (InternalKey k)
forall k v. ReversibleMap k v -> HashMap k v
directMap (ReversibleMap k (InternalKey k) -> HashMap k (InternalKey k))
-> ReversibleMap k (InternalKey k) -> HashMap k (InternalKey k)
forall a b. (a -> b) -> a -> b
$ ShareMap k v -> ReversibleMap k (InternalKey k)
forall k v. ShareMap k v -> ReversibleMap k (InternalKey k)
shareMap ShareMap k v
sm of
    Just InternalKey k
k' -> ShareMap k v
sm
      { unsharedMap = HashMap.insertWith f k' v (unsharedMap sm)
      }
    Maybe (InternalKey k)
Nothing -> ShareMap
      { unsharedMap :: HashMap (InternalKey k) v
unsharedMap = (v -> v -> v)
-> InternalKey k
-> v
-> HashMap (InternalKey k) v
-> HashMap (InternalKey k) v
forall k v.
(Eq k, Hashable k) =>
(v -> v -> v) -> k -> v -> HashMap k v -> HashMap k v
HashMap.insertWith v -> v -> v
f (k -> InternalKey k
forall k. k -> InternalKey k
InternalKey k
k) v
v (ShareMap k v -> HashMap (InternalKey k) v
forall k v. ShareMap k v -> HashMap (InternalKey k) v
unsharedMap ShareMap k v
sm)
      , shareMap :: ReversibleMap k (InternalKey k)
shareMap = k
-> InternalKey k
-> ReversibleMap k (InternalKey k)
-> ReversibleMap k (InternalKey k)
forall k v.
(Hashable k, Eq k, Hashable v, Eq v) =>
k -> v -> ReversibleMap k v -> ReversibleMap k v
insertReversibleMap k
k (k -> InternalKey k
forall k. k -> InternalKey k
InternalKey k
k) (ShareMap k v -> ReversibleMap k (InternalKey k)
forall k v. ShareMap k v -> ReversibleMap k (InternalKey k)
shareMap ShareMap k v
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 :: forall k v.
(Hashable k, Eq k) =>
(v -> v -> v) -> k -> k -> ShareMap k v -> ShareMap k v
mergeKeysWith v -> v -> v
f k
k0 k
k1 ShareMap k v
sm | k
k0 k -> k -> Bool
forall a. Eq a => a -> a -> Bool
/= k
k1 =
  case k -> ReversibleMap k (InternalKey k) -> Maybe (InternalKey k)
forall k v. (Hashable k, Eq k) => k -> ReversibleMap k v -> Maybe v
lookupReversibleMap k
k1 (ShareMap k v -> ReversibleMap k (InternalKey k)
forall k v. ShareMap k v -> ReversibleMap k (InternalKey k)
shareMap ShareMap k v
sm) of
    Just InternalKey k
k1' | k -> InternalKey k
forall k. k -> InternalKey k
InternalKey k
k0 InternalKey k -> InternalKey k -> Bool
forall a. Eq a => a -> a -> Bool
/= InternalKey k
k1' -> case InternalKey k -> HashMap (InternalKey k) v -> Maybe v
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HashMap.lookup InternalKey k
k1' (ShareMap k v -> HashMap (InternalKey k) v
forall k v. ShareMap k v -> HashMap (InternalKey k) v
unsharedMap ShareMap k v
sm) of
      Just v
v1 -> case k -> ReversibleMap k (InternalKey k) -> Maybe (InternalKey k)
forall k v. (Hashable k, Eq k) => k -> ReversibleMap k v -> Maybe v
lookupReversibleMap k
k0 (ShareMap k v -> ReversibleMap k (InternalKey k)
forall k v. ShareMap k v -> ReversibleMap k (InternalKey k)
shareMap ShareMap k v
sm) of
        Just InternalKey k
k0' | InternalKey k
k0' InternalKey k -> InternalKey k -> Bool
forall a. Eq a => a -> a -> Bool
/= InternalKey k
k1' ->
          ShareMap
            { unsharedMap :: HashMap (InternalKey k) v
unsharedMap = (v -> v -> v)
-> InternalKey k
-> v
-> HashMap (InternalKey k) v
-> HashMap (InternalKey k) v
forall k v.
(Eq k, Hashable k) =>
(v -> v -> v) -> k -> v -> HashMap k v -> HashMap k v
HashMap.insertWith ((v -> v -> v) -> v -> v -> v
forall a b c. (a -> b -> c) -> b -> a -> c
flip v -> v -> v
f) InternalKey k
k0' v
v1 (ShareMap k v -> HashMap (InternalKey k) v
forall k v. ShareMap k v -> HashMap (InternalKey k) v
unsharedMap ShareMap k v
sm)
            , shareMap :: ReversibleMap k (InternalKey k)
shareMap = -- Any values pointing to k1 are redirected to k0':
                (ReversibleMap k (InternalKey k)
 -> k -> ReversibleMap k (InternalKey k))
-> ReversibleMap k (InternalKey k)
-> HashSet k
-> ReversibleMap k (InternalKey k)
forall a b. (a -> b -> a) -> a -> HashSet b -> a
HashSet.foldl' (\ReversibleMap k (InternalKey k)
m k
k -> k
-> InternalKey k
-> ReversibleMap k (InternalKey k)
-> ReversibleMap k (InternalKey k)
forall k v.
(Hashable k, Eq k, Hashable v, Eq v) =>
k -> v -> ReversibleMap k v -> ReversibleMap k v
insertReversibleMap k
k InternalKey k
k0' ReversibleMap k (InternalKey k)
m) (ShareMap k v -> ReversibleMap k (InternalKey k)
forall k v. ShareMap k v -> ReversibleMap k (InternalKey k)
shareMap ShareMap k v
sm) (HashSet k -> ReversibleMap k (InternalKey k))
-> HashSet k -> ReversibleMap k (InternalKey k)
forall a b. (a -> b) -> a -> b
$
                InternalKey k -> ReversibleMap k (InternalKey k) -> HashSet k
forall v k.
(Hashable v, Eq v) =>
v -> ReversibleMap k v -> HashSet k
reverseLookup InternalKey k
k1' (ReversibleMap k (InternalKey k) -> HashSet k)
-> ReversibleMap k (InternalKey k) -> HashSet k
forall a b. (a -> b) -> a -> b
$ ShareMap k v -> ReversibleMap k (InternalKey k)
forall k v. ShareMap k v -> ReversibleMap k (InternalKey k)
shareMap ShareMap k v
sm
            }
        Maybe (InternalKey k)
Nothing ->
          ShareMap k v
sm { shareMap = insertReversibleMap k0 k1' (shareMap sm) }
        Maybe (InternalKey k)
_ ->
          ShareMap k v
sm
      Maybe v
Nothing -> String -> ShareMap k v
forall a. HasCallStack => String -> a
error String
"mergeKeysWith: broken invariant: unexpected missing key in unsharedMap"
    Maybe (InternalKey k)
Nothing ->
      case k -> HashMap k (InternalKey k) -> Maybe (InternalKey k)
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HashMap.lookup k
k0 (ReversibleMap k (InternalKey k) -> HashMap k (InternalKey k)
forall k v. ReversibleMap k v -> HashMap k v
directMap (ReversibleMap k (InternalKey k) -> HashMap k (InternalKey k))
-> ReversibleMap k (InternalKey k) -> HashMap k (InternalKey k)
forall a b. (a -> b) -> a -> b
$ ShareMap k v -> ReversibleMap k (InternalKey k)
forall k v. ShareMap k v -> ReversibleMap k (InternalKey k)
shareMap ShareMap k v
sm) of
        Just InternalKey k
k0' ->
          ShareMap k v
sm { shareMap = insertReversibleMap k1 k0' (shareMap sm) }
        Maybe (InternalKey k)
Nothing ->
          ShareMap k v
sm
    Maybe (InternalKey k)
_ ->
      ShareMap k v
sm
mergeKeysWith v -> v -> v
_ k
_ k
_ ShareMap k v
sm = ShareMap k v
sm

map :: (a -> b) -> ShareMap k a -> ShareMap k b
map :: forall a b k. (a -> b) -> ShareMap k a -> ShareMap k b
map a -> b
f ShareMap k a
sm = ShareMap k a
sm { unsharedMap = HashMap.map f (unsharedMap sm) }

-- | A map with an efficient 'reverseLookup'
data ReversibleMap k v = ReversibleMap
  { forall k v. ReversibleMap k v -> HashMap k v
directMap :: HashMap k v
  , -- |
    -- > forall (v, ks) in reversedMap.
    -- >   forall k in ks.
    -- >     (k, v) is in directMap
    forall k v. ReversibleMap k v -> HashMap v (HashSet k)
reversedMap :: HashMap v (HashSet k)
  }
  deriving Int -> ReversibleMap k v -> ShowS
[ReversibleMap k v] -> ShowS
ReversibleMap k v -> String
(Int -> ReversibleMap k v -> ShowS)
-> (ReversibleMap k v -> String)
-> ([ReversibleMap k v] -> ShowS)
-> Show (ReversibleMap k v)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall k v. (Show k, Show v) => Int -> ReversibleMap k v -> ShowS
forall k v. (Show k, Show v) => [ReversibleMap k v] -> ShowS
forall k v. (Show k, Show v) => ReversibleMap k v -> String
$cshowsPrec :: forall k v. (Show k, Show v) => Int -> ReversibleMap k v -> ShowS
showsPrec :: Int -> ReversibleMap k v -> ShowS
$cshow :: forall k v. (Show k, Show v) => ReversibleMap k v -> String
show :: ReversibleMap k v -> String
$cshowList :: forall k v. (Show k, Show v) => [ReversibleMap k v] -> ShowS
showList :: [ReversibleMap k v] -> ShowS
Show

emptyReversibleMap :: ReversibleMap k v
emptyReversibleMap :: forall k v. ReversibleMap k v
emptyReversibleMap = HashMap k v -> HashMap v (HashSet k) -> ReversibleMap k v
forall k v.
HashMap k v -> HashMap v (HashSet k) -> ReversibleMap k v
ReversibleMap HashMap k v
forall k v. HashMap k v
HashMap.empty HashMap v (HashSet k)
forall k v. HashMap k v
HashMap.empty

insertReversibleMap
  :: (Hashable k, Eq k, Hashable v, Eq v)
  => k
  -> v
  -> ReversibleMap k v
  -> ReversibleMap k v
insertReversibleMap :: forall k v.
(Hashable k, Eq k, Hashable v, Eq v) =>
k -> v -> ReversibleMap k v -> ReversibleMap k v
insertReversibleMap k
k v
v ReversibleMap k v
rm = ReversibleMap
  { directMap :: HashMap k v
directMap = k -> v -> HashMap k v -> HashMap k v
forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
HashMap.insert k
k v
v (ReversibleMap k v -> HashMap k v
forall k v. ReversibleMap k v -> HashMap k v
directMap ReversibleMap k v
rm)
  , reversedMap :: HashMap v (HashSet k)
reversedMap =
      let m' :: HashMap v (HashSet k)
m' = case k -> HashMap k v -> Maybe v
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HashMap.lookup k
k (ReversibleMap k v -> HashMap k v
forall k v. ReversibleMap k v -> HashMap k v
directMap ReversibleMap k v
rm) of
            Maybe v
Nothing -> ReversibleMap k v -> HashMap v (HashSet k)
forall k v. ReversibleMap k v -> HashMap v (HashSet k)
reversedMap ReversibleMap k v
rm
            Just v
oldv -> (HashSet k -> HashSet k)
-> v -> HashMap v (HashSet k) -> HashMap v (HashSet k)
forall k v.
(Eq k, Hashable k) =>
(v -> v) -> k -> HashMap k v -> HashMap k v
HashMap.adjust (k -> HashSet k -> HashSet k
forall a. (Eq a, Hashable a) => a -> HashSet a -> HashSet a
HashSet.delete k
k) v
oldv (ReversibleMap k v -> HashMap v (HashSet k)
forall k v. ReversibleMap k v -> HashMap v (HashSet k)
reversedMap ReversibleMap k v
rm)
       in (HashSet k -> HashSet k -> HashSet k)
-> v -> HashSet k -> HashMap v (HashSet k) -> HashMap v (HashSet k)
forall k v.
(Eq k, Hashable k) =>
(v -> v -> v) -> k -> v -> HashMap k v -> HashMap k v
HashMap.insertWith HashSet k -> HashSet k -> HashSet k
forall a. Eq a => HashSet a -> HashSet a -> HashSet a
HashSet.union v
v (k -> HashSet k
forall a. Hashable a => a -> HashSet a
HashSet.singleton k
k) HashMap v (HashSet k)
m'
  }

reverseLookup :: (Hashable v, Eq v) => v -> ReversibleMap k v -> HashSet k
reverseLookup :: forall v k.
(Hashable v, Eq v) =>
v -> ReversibleMap k v -> HashSet k
reverseLookup v
v ReversibleMap k v
rm = HashSet k -> Maybe (HashSet k) -> HashSet k
forall a. a -> Maybe a -> a
fromMaybe HashSet k
forall a. HashSet a
HashSet.empty (Maybe (HashSet k) -> HashSet k) -> Maybe (HashSet k) -> HashSet k
forall a b. (a -> b) -> a -> b
$ v -> HashMap v (HashSet k) -> Maybe (HashSet k)
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HashMap.lookup v
v (ReversibleMap k v -> HashMap v (HashSet k)
forall k v. ReversibleMap k v -> HashMap v (HashSet k)
reversedMap ReversibleMap k v
rm)

lookupReversibleMap :: (Hashable k, Eq k) => k -> ReversibleMap k v -> Maybe v
lookupReversibleMap :: forall k v. (Hashable k, Eq k) => k -> ReversibleMap k v -> Maybe v
lookupReversibleMap k
k ReversibleMap k v
rm = k -> HashMap k v -> Maybe v
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HashMap.lookup k
k (ReversibleMap k v -> HashMap k v
forall k v. ReversibleMap k v -> HashMap k v
directMap ReversibleMap k v
rm)