{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE ApplicativeDo #-}
{-# OPTIONS_GHC -Wall #-}

-- | Like a 'UniqDFM', but maintains equivalence classes of keys sharing the
-- same entry. See 'UniqSDFM'.
module GHC.Types.Unique.SDFM (
        -- * Unique-keyed, /shared/, deterministic mappings
        UniqSDFM,

        emptyUSDFM,
        lookupUSDFM,
        equateUSDFM, addToUSDFM,
        traverseUSDFM
    ) where

import GHC.Prelude

import GHC.Types.Unique
import GHC.Types.Unique.DFM
import GHC.Utils.Outputable

-- | Either @Indirect x@, meaning the value is represented by that of @x@, or
-- an @Entry@ containing containing the actual value it represents.
data Shared key ele
  = Indirect !key
  | Entry !ele

-- | A 'UniqDFM' whose domain is /sets/ of 'Unique's, each of which share a
-- common value of type @ele@.
-- Every such set (\"equivalence class\") has a distinct representative
-- 'Unique'. Supports merging the entries of multiple such sets in a union-find
-- like fashion.
--
-- An accurate model is that of @[(Set key, Maybe ele)]@: A finite mapping from
-- sets of @key@s to possibly absent entries @ele@, where the sets don't overlap.
-- Example:
-- @
--   m = [({u1,u3}, Just ele1), ({u2}, Just ele2), ({u4,u7}, Nothing)]
-- @
-- On this model we support the following main operations:
--
--   * @'lookupUSDFM' m u3 == Just ele1@, @'lookupUSDFM' m u4 == Nothing@,
--     @'lookupUSDFM' m u5 == Nothing@.
--   * @'equateUSDFM' m u1 u3@ is a no-op, but
--     @'equateUSDFM' m u1 u2@ merges @{u1,u3}@ and @{u2}@ to point to
--     @Just ele2@ and returns the old entry of @{u1,u3}@, @Just ele1@.
--   * @'addToUSDFM' m u3 ele4@ sets the entry of @{u1,u3}@ to @Just ele4@.
--
-- As well as a few means for traversal/conversion to list.
newtype UniqSDFM key ele
  = USDFM { forall key ele. UniqSDFM key ele -> UniqDFM key (Shared key ele)
unUSDFM :: UniqDFM key (Shared key ele) }

emptyUSDFM :: UniqSDFM key ele
emptyUSDFM :: forall key ele. UniqSDFM key ele
emptyUSDFM = UniqDFM key (Shared key ele) -> UniqSDFM key ele
forall key ele. UniqDFM key (Shared key ele) -> UniqSDFM key ele
USDFM UniqDFM key (Shared key ele)
forall {k} (key :: k) elt. UniqDFM key elt
emptyUDFM

lookupReprAndEntryUSDFM :: Uniquable key => UniqSDFM key ele -> key -> (key, Maybe ele)
lookupReprAndEntryUSDFM :: forall key ele.
Uniquable key =>
UniqSDFM key ele -> key -> (key, Maybe ele)
lookupReprAndEntryUSDFM (USDFM UniqDFM key (Shared key ele)
env) = key -> (key, Maybe ele)
go
  where
    go :: key -> (key, Maybe ele)
go key
x = case UniqDFM key (Shared key ele) -> key -> Maybe (Shared key ele)
forall key elt.
Uniquable key =>
UniqDFM key elt -> key -> Maybe elt
lookupUDFM UniqDFM key (Shared key ele)
env key
x of
      Maybe (Shared key ele)
Nothing           -> (key
x, Maybe ele
forall a. Maybe a
Nothing)
      Just (Indirect key
y) -> key -> (key, Maybe ele)
go key
y
      Just (Entry ele
ele)  -> (key
x, ele -> Maybe ele
forall a. a -> Maybe a
Just ele
ele)

-- | @lookupSUDFM env x@ looks up an entry for @x@, looking through all
-- 'Indirect's until it finds a shared 'Entry'.
--
-- Examples in terms of the model (see 'UniqSDFM'):
-- >>> lookupUSDFM [({u1,u3}, Just ele1), ({u2}, Just ele2)] u3 == Just ele1
-- >>> lookupUSDFM [({u1,u3}, Just ele1), ({u2}, Just ele2)] u4 == Nothing
-- >>> lookupUSDFM [({u1,u3}, Just ele1), ({u2}, Nothing)] u2 == Nothing
lookupUSDFM :: Uniquable key => UniqSDFM key ele -> key -> Maybe ele
lookupUSDFM :: forall key ele.
Uniquable key =>
UniqSDFM key ele -> key -> Maybe ele
lookupUSDFM UniqSDFM key ele
usdfm key
x = (key, Maybe ele) -> Maybe ele
forall a b. (a, b) -> b
snd (UniqSDFM key ele -> key -> (key, Maybe ele)
forall key ele.
Uniquable key =>
UniqSDFM key ele -> key -> (key, Maybe ele)
lookupReprAndEntryUSDFM UniqSDFM key ele
usdfm key
x)

-- | @equateUSDFM env x y@ makes @x@ and @y@ point to the same entry,
-- thereby merging @x@'s class with @y@'s.
-- If both @x@ and @y@ are in the domain of the map, then @y@'s entry will be
-- chosen as the new entry and @x@'s old entry will be returned.
--
-- Examples in terms of the model (see 'UniqSDFM'):
-- >>> equateUSDFM [] u1 u2 == (Nothing, [({u1,u2}, Nothing)])
-- >>> equateUSDFM [({u1,u3}, Just ele1)] u3 u4 == (Nothing, [({u1,u3,u4}, Just ele1)])
-- >>> equateUSDFM [({u1,u3}, Just ele1)] u4 u3 == (Nothing, [({u1,u3,u4}, Just ele1)])
-- >>> equateUSDFM [({u1,u3}, Just ele1), ({u2}, Just ele2)] u3 u2 == (Just ele1, [({u2,u1,u3}, Just ele2)])
equateUSDFM
  :: Uniquable key => UniqSDFM key ele -> key -> key -> (Maybe ele, UniqSDFM key ele)
equateUSDFM :: forall key ele.
Uniquable key =>
UniqSDFM key ele -> key -> key -> (Maybe ele, UniqSDFM key ele)
equateUSDFM usdfm :: UniqSDFM key ele
usdfm@(USDFM UniqDFM key (Shared key ele)
env) key
x key
y =
  case (key -> (key, Maybe ele)
lu key
x, key -> (key, Maybe ele)
lu key
y) of
    ((key
x', Maybe ele
_)    , (key
y', Maybe ele
_))
      | key -> Unique
forall a. Uniquable a => a -> Unique
getUnique key
x' Unique -> Unique -> Bool
forall a. Eq a => a -> a -> Bool
== key -> Unique
forall a. Uniquable a => a -> Unique
getUnique key
y' -> (Maybe ele
forall a. Maybe a
Nothing, UniqSDFM key ele
usdfm) -- nothing to do
    ((key
x', Maybe ele
_)    , (key
y', Maybe ele
Nothing))     -> (Maybe ele
forall a. Maybe a
Nothing, key -> key -> UniqSDFM key ele
set_indirect key
y' key
x')
    ((key
x', Maybe ele
mb_ex), (key
y', Maybe ele
_))           -> (Maybe ele
mb_ex,   key -> key -> UniqSDFM key ele
set_indirect key
x' key
y')
  where
    lu :: key -> (key, Maybe ele)
lu = UniqSDFM key ele -> key -> (key, Maybe ele)
forall key ele.
Uniquable key =>
UniqSDFM key ele -> key -> (key, Maybe ele)
lookupReprAndEntryUSDFM UniqSDFM key ele
usdfm
    set_indirect :: key -> key -> UniqSDFM key ele
set_indirect key
a key
b = UniqDFM key (Shared key ele) -> UniqSDFM key ele
forall key ele. UniqDFM key (Shared key ele) -> UniqSDFM key ele
USDFM (UniqDFM key (Shared key ele) -> UniqSDFM key ele)
-> UniqDFM key (Shared key ele) -> UniqSDFM key ele
forall a b. (a -> b) -> a -> b
$ UniqDFM key (Shared key ele)
-> key -> Shared key ele -> UniqDFM key (Shared key ele)
forall key elt.
Uniquable key =>
UniqDFM key elt -> key -> elt -> UniqDFM key elt
addToUDFM UniqDFM key (Shared key ele)
env key
a (key -> Shared key ele
forall key ele. key -> Shared key ele
Indirect key
b)

-- | @addToUSDFM env x a@ sets the entry @x@ is associated with to @a@,
-- thereby modifying its whole equivalence class.
--
-- Examples in terms of the model (see 'UniqSDFM'):
-- >>> addToUSDFM [] u1 ele1 == [({u1}, Just ele1)]
-- >>> addToUSDFM [({u1,u3}, Just ele1)] u3 ele2 == [({u1,u3}, Just ele2)]
addToUSDFM :: Uniquable key => UniqSDFM key ele -> key -> ele -> UniqSDFM key ele
addToUSDFM :: forall key ele.
Uniquable key =>
UniqSDFM key ele -> key -> ele -> UniqSDFM key ele
addToUSDFM usdfm :: UniqSDFM key ele
usdfm@(USDFM UniqDFM key (Shared key ele)
env) key
x ele
v =
  UniqDFM key (Shared key ele) -> UniqSDFM key ele
forall key ele. UniqDFM key (Shared key ele) -> UniqSDFM key ele
USDFM (UniqDFM key (Shared key ele) -> UniqSDFM key ele)
-> UniqDFM key (Shared key ele) -> UniqSDFM key ele
forall a b. (a -> b) -> a -> b
$ UniqDFM key (Shared key ele)
-> key -> Shared key ele -> UniqDFM key (Shared key ele)
forall key elt.
Uniquable key =>
UniqDFM key elt -> key -> elt -> UniqDFM key elt
addToUDFM UniqDFM key (Shared key ele)
env ((key, Maybe ele) -> key
forall a b. (a, b) -> a
fst (UniqSDFM key ele -> key -> (key, Maybe ele)
forall key ele.
Uniquable key =>
UniqSDFM key ele -> key -> (key, Maybe ele)
lookupReprAndEntryUSDFM UniqSDFM key ele
usdfm key
x)) (ele -> Shared key ele
forall key ele. ele -> Shared key ele
Entry ele
v)

traverseUSDFM :: forall key a b f. Applicative f => (a -> f b) -> UniqSDFM key a -> f (UniqSDFM key b)
traverseUSDFM :: forall key a b (f :: * -> *).
Applicative f =>
(a -> f b) -> UniqSDFM key a -> f (UniqSDFM key b)
traverseUSDFM a -> f b
f = ([(Unique, Shared key b)] -> UniqSDFM key b)
-> f [(Unique, Shared key b)] -> f (UniqSDFM key b)
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (UniqDFM key (Shared key b) -> UniqSDFM key b
forall key ele. UniqDFM key (Shared key ele) -> UniqSDFM key ele
USDFM (UniqDFM key (Shared key b) -> UniqSDFM key b)
-> ([(Unique, Shared key b)] -> UniqDFM key (Shared key b))
-> [(Unique, Shared key b)]
-> UniqSDFM key b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Unique, Shared key b)] -> UniqDFM key (Shared key b)
forall {k} elt (key :: k). [(Unique, elt)] -> UniqDFM key elt
listToUDFM_Directly) (f [(Unique, Shared key b)] -> f (UniqSDFM key b))
-> (UniqSDFM key a -> f [(Unique, Shared key b)])
-> UniqSDFM key a
-> f (UniqSDFM key b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Unique, Shared key a) -> f (Unique, Shared key b))
-> [(Unique, Shared key a)] -> f [(Unique, Shared key b)]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse (Unique, Shared key a) -> f (Unique, Shared key b)
g ([(Unique, Shared key a)] -> f [(Unique, Shared key b)])
-> (UniqSDFM key a -> [(Unique, Shared key a)])
-> UniqSDFM key a
-> f [(Unique, Shared key b)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. UniqDFM key (Shared key a) -> [(Unique, Shared key a)]
forall {k} (key :: k) elt. UniqDFM key elt -> [(Unique, elt)]
udfmToList (UniqDFM key (Shared key a) -> [(Unique, Shared key a)])
-> (UniqSDFM key a -> UniqDFM key (Shared key a))
-> UniqSDFM key a
-> [(Unique, Shared key a)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. UniqSDFM key a -> UniqDFM key (Shared key a)
forall key ele. UniqSDFM key ele -> UniqDFM key (Shared key ele)
unUSDFM
  where
    g :: (Unique, Shared key a) -> f (Unique, Shared key b)
    g :: (Unique, Shared key a) -> f (Unique, Shared key b)
g (Unique
u, Indirect key
y) = (Unique, Shared key b) -> f (Unique, Shared key b)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Unique
u,key -> Shared key b
forall key ele. key -> Shared key ele
Indirect key
y)
    g (Unique
u, Entry a
a)    = do
        b
a' <- a -> f b
f a
a
        pure (Unique
u,b -> Shared key b
forall key ele. ele -> Shared key ele
Entry b
a')

instance (Outputable key, Outputable ele) => Outputable (Shared key ele) where
  ppr :: Shared key ele -> SDoc
ppr (Indirect key
x) = key -> SDoc
forall a. Outputable a => a -> SDoc
ppr key
x
  ppr (Entry ele
a)    = ele -> SDoc
forall a. Outputable a => a -> SDoc
ppr ele
a

instance (Outputable key, Outputable ele) => Outputable (UniqSDFM key ele) where
  ppr :: UniqSDFM key ele -> SDoc
ppr (USDFM UniqDFM key (Shared key ele)
env) = UniqDFM key (Shared key ele) -> SDoc
forall a. Outputable a => a -> SDoc
ppr UniqDFM key (Shared key ele)
env