{-# 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
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
showList :: [ShareMap k v] -> ShowS
$cshowList :: forall k v. (Show k, Show v) => [ShareMap k v] -> ShowS
show :: ShareMap k v -> String
$cshow :: forall k v. (Show k, Show v) => ShareMap k v -> String
showsPrec :: Int -> ShareMap k v -> ShowS
$cshowsPrec :: forall k v. (Show k, Show v) => Int -> 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
forall k. Eq k => InternalKey k -> InternalKey k -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: InternalKey k -> InternalKey k -> Bool
$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
Eq, 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
hash :: InternalKey k -> Int
$chash :: forall k. Hashable k => InternalKey k -> Int
hashWithSalt :: Int -> InternalKey k -> Int
$chashWithSalt :: forall k. Hashable k => Int -> InternalKey k -> Int
Hashable)

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

empty :: ShareMap k v
empty :: forall k v. ShareMap k v
empty = forall k v.
HashMap (InternalKey k) v
-> ReversibleMap k (InternalKey k) -> ShareMap k v
ShareMap forall k v. HashMap k v
HashMap.empty 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 =
  forall a k v. (a -> k -> v -> a) -> a -> HashMap k v -> a
HashMap.foldlWithKey' forall {k}.
Hashable k =>
HashMap k v -> k -> InternalKey k -> HashMap k v
expand forall k v. HashMap k v
HashMap.empty (forall k v. ReversibleMap k v -> HashMap k v
directMap forall a b. (a -> b) -> a -> b
$ 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' =
      forall b a. b -> (a -> b) -> Maybe a -> b
maybe HashMap k v
m (\v
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) (forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HashMap.lookup InternalKey k
k' forall a b. (a -> b) -> a -> b
$ 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 forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HashMap.lookup k
k forall a b. (a -> b) -> a -> b
$ forall k v. ReversibleMap k v -> HashMap k v
directMap forall a b. (a -> b) -> a -> b
$ 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 (InternalKey k) v
unsharedMap = forall k v.
(Eq k, Hashable k) =>
(v -> v -> v) -> k -> v -> HashMap k v -> HashMap k v
HashMap.insertWith v -> v -> v
f InternalKey k
k' v
v (forall k v. ShareMap k v -> HashMap (InternalKey k) v
unsharedMap ShareMap k v
sm)
      }
    Maybe (InternalKey k)
Nothing -> ShareMap
      { unsharedMap :: HashMap (InternalKey k) v
unsharedMap = forall k v.
(Eq k, Hashable k) =>
(v -> v -> v) -> k -> v -> HashMap k v -> HashMap k v
HashMap.insertWith v -> v -> v
f (forall k. k -> InternalKey k
InternalKey k
k) v
v (forall k v. ShareMap k v -> HashMap (InternalKey k) v
unsharedMap ShareMap k v
sm)
      , shareMap :: ReversibleMap k (InternalKey k)
shareMap = forall k v.
(Hashable k, Eq k, Hashable v, Eq v) =>
k -> v -> ReversibleMap k v -> ReversibleMap k v
insertReversibleMap k
k (forall k. k -> InternalKey k
InternalKey k
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 forall a. Eq a => a -> a -> Bool
/= k
k1 =
  case forall k v. (Hashable k, Eq k) => k -> ReversibleMap k v -> Maybe v
lookupReversibleMap k
k1 (forall k v. ShareMap k v -> ReversibleMap k (InternalKey k)
shareMap ShareMap k v
sm) of
    Just InternalKey k
k1' | forall k. k -> InternalKey k
InternalKey k
k0 forall a. Eq a => a -> a -> Bool
/= InternalKey k
k1' -> case forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HashMap.lookup InternalKey k
k1' (forall k v. ShareMap k v -> HashMap (InternalKey k) v
unsharedMap ShareMap k v
sm) of
      Just v
v1 -> case forall k v. (Hashable k, Eq k) => k -> ReversibleMap k v -> Maybe v
lookupReversibleMap k
k0 (forall k v. ShareMap k v -> ReversibleMap k (InternalKey k)
shareMap ShareMap k v
sm) of
        Just InternalKey k
k0' | InternalKey k
k0' forall a. Eq a => a -> a -> Bool
/= InternalKey k
k1' ->
          ShareMap
            { unsharedMap :: HashMap (InternalKey k) v
unsharedMap = forall k v.
(Eq k, Hashable k) =>
(v -> v -> v) -> k -> v -> HashMap k v -> HashMap k v
HashMap.insertWith (forall a b c. (a -> b -> c) -> b -> a -> c
flip v -> v -> v
f) InternalKey k
k0' v
v1 (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':
                forall a b. (a -> b -> a) -> a -> HashSet b -> a
HashSet.foldl' (\ReversibleMap k (InternalKey k)
m k
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) (forall k v. ShareMap k v -> ReversibleMap k (InternalKey k)
shareMap ShareMap k v
sm) forall a b. (a -> b) -> a -> b
$
                forall v k.
(Hashable v, Eq v) =>
v -> ReversibleMap k v -> HashSet k
reverseLookup InternalKey k
k1' forall a b. (a -> b) -> a -> b
$ forall k v. ShareMap k v -> ReversibleMap k (InternalKey k)
shareMap ShareMap k v
sm
            }
        Maybe (InternalKey k)
Nothing ->
          ShareMap k v
sm { shareMap :: ReversibleMap k (InternalKey k)
shareMap = forall k v.
(Hashable k, Eq k, Hashable v, Eq v) =>
k -> v -> ReversibleMap k v -> ReversibleMap k v
insertReversibleMap k
k0 InternalKey k
k1' (forall k v. ShareMap k v -> ReversibleMap k (InternalKey k)
shareMap ShareMap k v
sm) }
        Maybe (InternalKey k)
_ ->
          ShareMap k v
sm
      Maybe v
Nothing -> forall a. HasCallStack => String -> a
error String
"mergeKeysWith: broken invariant: unexpected missing key in unsharedMap"
    Maybe (InternalKey k)
Nothing ->
      case forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HashMap.lookup k
k0 (forall k v. ReversibleMap k v -> HashMap k v
directMap forall a b. (a -> b) -> a -> b
$ 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 :: ReversibleMap k (InternalKey k)
shareMap = forall k v.
(Hashable k, Eq k, Hashable v, Eq v) =>
k -> v -> ReversibleMap k v -> ReversibleMap k v
insertReversibleMap k
k1 InternalKey k
k0' (forall k v. ShareMap k v -> ReversibleMap k (InternalKey k)
shareMap ShareMap k v
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 (InternalKey k) b
unsharedMap = forall v1 v2 k. (v1 -> v2) -> HashMap k v1 -> HashMap k v2
HashMap.map a -> b
f (forall k v. ShareMap k v -> HashMap (InternalKey k) v
unsharedMap ShareMap k a
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
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
showList :: [ReversibleMap k v] -> ShowS
$cshowList :: forall k v. (Show k, Show v) => [ReversibleMap k v] -> ShowS
show :: ReversibleMap k v -> String
$cshow :: forall k v. (Show k, Show v) => ReversibleMap k v -> String
showsPrec :: Int -> ReversibleMap k v -> ShowS
$cshowsPrec :: forall k v. (Show k, Show v) => Int -> ReversibleMap k v -> ShowS
Show

emptyReversibleMap :: ReversibleMap k v
emptyReversibleMap :: forall k v. ReversibleMap k v
emptyReversibleMap = forall k v.
HashMap k v -> HashMap v (HashSet k) -> ReversibleMap k v
ReversibleMap forall k v. HashMap k v
HashMap.empty 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 = forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
HashMap.insert k
k v
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 forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HashMap.lookup k
k (forall k v. ReversibleMap k v -> HashMap k v
directMap ReversibleMap k v
rm) of
            Maybe v
Nothing -> forall k v. ReversibleMap k v -> HashMap v (HashSet k)
reversedMap ReversibleMap k v
rm
            Just v
oldv -> forall k v.
(Eq k, Hashable k) =>
(v -> v) -> k -> HashMap k v -> HashMap k v
HashMap.adjust (forall a. (Eq a, Hashable a) => a -> HashSet a -> HashSet a
HashSet.delete k
k) v
oldv (forall k v. ReversibleMap k v -> HashMap v (HashSet k)
reversedMap ReversibleMap k v
rm)
       in forall k v.
(Eq k, Hashable k) =>
(v -> v -> v) -> k -> v -> HashMap k v -> HashMap k v
HashMap.insertWith forall a. (Eq a, Hashable a) => HashSet a -> HashSet a -> HashSet a
HashSet.union v
v (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 = forall a. a -> Maybe a -> a
fromMaybe forall a. HashSet a
HashSet.empty forall a b. (a -> b) -> a -> b
$ forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HashMap.lookup v
v (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 = forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HashMap.lookup k
k (forall k v. ReversibleMap k v -> HashMap k v
directMap ReversibleMap k v
rm)