Extract SharedIdEnv into its own module

It's now named `GHC.Types.Unique.SDFM.UniqSDFM`.
The implementation is more clear about its stated goals and supported
operations.

5 files changed, 160 insertions, 143 deletions

diff --git a/compiler/GHC/HsToCore/Pmc/Ppr.hs b/compiler/GHC/HsToCore/Pmc/Ppr.hs
index ebad0757b2..bb30cd61ed 100644
--- a/compiler/GHC/HsToCore/Pmc/Ppr.hs
+++ b/compiler/GHC/HsToCore/Pmc/Ppr.hs
@@ -98,20 +98,20 @@ substitution to the vectors before printing them out (see function `pprOne' in
 
 -- | Extract and assigns pretty names to constraint variables with refutable
 -- shapes.
-prettifyRefuts :: Nabla -> DIdEnv SDoc -> DIdEnv (SDoc, [PmAltCon])
+prettifyRefuts :: Nabla -> DIdEnv (Id, SDoc) -> DIdEnv (SDoc, [PmAltCon])
 prettifyRefuts nabla = listToUDFM_Directly . map attach_refuts . udfmToList
   where
-    attach_refuts (u, sdoc) = (u, (sdoc, lookupRefuts nabla u))
+    attach_refuts (u, (x, sdoc)) = (u, (sdoc, lookupRefuts nabla x))
 
 
-type PmPprM a = RWS Nabla () (DIdEnv SDoc, [SDoc]) a
+type PmPprM a = RWS Nabla () (DIdEnv (Id, SDoc), [SDoc]) a
 
 -- Try nice names p,q,r,s,t before using the (ugly) t_i
 nameList :: [SDoc]
 nameList = map text ["p","q","r","s","t"] ++
             [ text ('t':show u) | u <- [(0 :: Int)..] ]
 
-runPmPpr :: Nabla -> PmPprM a -> (a, DIdEnv SDoc)
+runPmPpr :: Nabla -> PmPprM a -> (a, DIdEnv (Id, SDoc))
 runPmPpr nabla m = case runRWS m nabla (emptyDVarEnv, nameList) of
   (a, (renamings, _), _) -> (a, renamings)
 
@@ -122,9 +122,9 @@ getCleanName x = do
   (renamings, name_supply) <- get
   let (clean_name:name_supply') = name_supply
   case lookupDVarEnv renamings x of
-    Just nm -> pure nm
+    Just (_, nm) -> pure nm
     Nothing -> do
-      put (extendDVarEnv renamings x clean_name, name_supply')
+      put (extendDVarEnv renamings x (x, clean_name), name_supply')
       pure clean_name
 
 checkRefuts :: Id -> PmPprM (Maybe SDoc) -- the clean name if it has negative info attached
@@ -139,8 +139,8 @@ checkRefuts x = do
 -- underscores. Even with a type signature, if it's not too noisy.
 pprPmVar :: PprPrec -> Id -> PmPprM SDoc
 -- Type signature is "too noisy" by my definition if it needs to parenthesize.
--- I like           "not matched: _ :: Proxy (DIdEnv SDoc)",
--- but I don't like "not matched: (_ :: stuff) (_:_) (_ :: Proxy (DIdEnv SDoc))"
+-- I like           "not matched: _ :: Proxy (DIdEnv (Id, SDoc))",
+-- but I don't like "not matched: (_ :: stuff) (_:_) (_ :: Proxy (DIdEnv (Id, SDoc)))"
 -- The useful information in the latter case is the constructor that we missed,
 -- not the types of the wildcards in the places that aren't matched as a result.
 pprPmVar prec x = do
diff --git a/compiler/GHC/HsToCore/Pmc/Solver.hs b/compiler/GHC/HsToCore/Pmc/Solver.hs
index cfbe6c3b94..3fc1471835 100644
--- a/compiler/GHC/HsToCore/Pmc/Solver.hs
+++ b/compiler/GHC/HsToCore/Pmc/Solver.hs
@@ -45,10 +45,9 @@ import GHC.Utils.Error ( pprErrMsgBagWithLoc )
 import GHC.Utils.Misc
 import GHC.Utils.Panic
 import GHC.Data.Bag
-import GHC.Types.Unique
 import GHC.Types.Unique.Set
 import GHC.Types.Unique.DSet
-import GHC.Types.Unique.DFM
+import GHC.Types.Unique.SDFM
 import GHC.Types.Id
 import GHC.Types.Name
 import GHC.Types.Var      (EvVar)
@@ -494,7 +493,7 @@ emptyVarInfo x
 
 lookupVarInfo :: TmState -> Id -> VarInfo
 -- (lookupVarInfo tms x) tells what we know about 'x'
-lookupVarInfo (TmSt env _ _) x = fromMaybe (emptyVarInfo x) (lookupSDIE env x)
+lookupVarInfo (TmSt env _ _) x = fromMaybe (emptyVarInfo x) (lookupUSDFM env x)
 
 -- | Like @lookupVarInfo ts x@, but @lookupVarInfo ts x = (y, vi)@ also looks
 -- through newtype constructors. We have @x ~ N1 (... (Nk y))@ such that the
@@ -521,7 +520,7 @@ trvVarInfo f nabla@MkNabla{ nabla_tm_st = ts@TmSt{ts_facts = env} } x
   = set_vi <$> f (lookupVarInfo ts x)
   where
     set_vi (a, vi') =
-      (a, nabla{ nabla_tm_st = ts{ ts_facts = setEntrySDIE env (vi_id vi') vi' } })
+      (a, nabla{ nabla_tm_st = ts{ ts_facts = addToUSDFM env (vi_id vi') vi' } })
 
 {- Note [Coverage checking Newtype matches]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -554,14 +553,11 @@ where you can find the solution in a perhaps more digestible format.
 ------------------------------------------------
 -- * Exported utility functions querying 'Nabla'
 
-lookupRefuts :: Uniquable k => Nabla -> k -> [PmAltCon]
+lookupRefuts :: Nabla -> Id -> [PmAltCon]
 -- Unfortunately we need the extra bit of polymorphism and the unfortunate
 -- duplication of lookupVarInfo here.
-lookupRefuts MkNabla{ nabla_tm_st = ts@(TmSt{ts_facts = (SDIE env)}) } k =
-  case lookupUDFM_Directly env (getUnique k) of
-    Nothing -> []
-    Just (Indirect y) -> pmAltConSetElems (vi_neg (lookupVarInfo ts y))
-    Just (Entry vi)   -> pmAltConSetElems (vi_neg vi)
+lookupRefuts MkNabla{ nabla_tm_st = ts } x =
+  pmAltConSetElems $ vi_neg $ lookupVarInfo ts x
 
 isDataConSolution :: PmAltConApp -> Bool
 isDataConSolution PACA{paca_con = PmAltConLike (RealDataCon _)} = True
@@ -718,7 +714,7 @@ addBotCt nabla@MkNabla{ nabla_tm_st = ts@TmSt{ ts_facts=env } } x = do
     IsBot    -> pure nabla -- There already is x ~ ⊥. Nothing left to do
     MaybeBot -> do         -- We add x ~ ⊥
       let vi' = vi{ vi_bot = IsBot }
-      pure nabla{ nabla_tm_st = ts{ts_facts = setEntrySDIE env y vi' } }
+      pure nabla{ nabla_tm_st = ts{ts_facts = addToUSDFM env y vi' } }
 
 -- | Adds the constraint @x ~/ ⊥@ to 'Nabla'. Quite similar to 'addNotConCt',
 -- but only cares for the ⊥ "constructor".
@@ -732,7 +728,7 @@ addNotBotCt nabla@MkNabla{ nabla_tm_st = ts@TmSt{ts_facts=env} } x = do
       -- Mark dirty for a delayed inhabitation test
       let vi' = vi{ vi_bot = IsNotBot}
       pure $ markDirty y
-           $ nabla{ nabla_tm_st = ts{ ts_facts = setEntrySDIE env y vi' } }
+           $ nabla{ nabla_tm_st = ts{ ts_facts = addToUSDFM env y vi' } }
 
 -- | Record a @x ~/ K@ constraint, e.g. that a particular 'Id' @x@ can't
 -- take the shape of a 'PmAltCon' @K@ in the 'Nabla' and return @Nothing@ if
@@ -805,7 +801,7 @@ addConCt nabla@MkNabla{ nabla_tm_st = ts@TmSt{ ts_facts=env } } x alt tvs args =
     Nothing -> do
       let pos' = PACA alt tvs args : pos
       let nabla_with bot' =
-            nabla{ nabla_tm_st = ts{ts_facts = setEntrySDIE env x (vi{vi_pos = pos', vi_bot = bot'})} }
+            nabla{ nabla_tm_st = ts{ts_facts = addToUSDFM env x (vi{vi_pos = pos', vi_bot = bot'})} }
       -- Do (2) in Note [Coverage checking Newtype matches]
       case (alt, args) of
         (PmAltConLike (RealDataCon dc), [y]) | isNewDataCon dc ->
@@ -825,55 +821,27 @@ equateTys ts us =
   , not (eqType t u)
   ]
 
--- | Adds a @x ~ y@ constraint by trying to unify two 'Id's and record the
+-- | Adds a @x ~ y@ constraint by merging the two 'VarInfo's and record the
 -- gained knowledge in 'Nabla'.
 --
--- Returns @Nothing@ when there's a contradiction. Returns @Just nabla@
--- when the constraint was compatible with prior facts, in which case @nabla@
--- has integrated the knowledge from the equality constraint.
+-- Returns @Nothing@ when there's a contradiction while merging. Returns @Just
+-- nabla@ when the constraint was compatible with prior facts, in which case
+-- @nabla@ has integrated the knowledge from the equality constraint.
 --
 -- See Note [TmState invariants].
 addVarCt :: Nabla -> Id -> Id -> MaybeT DsM Nabla
-addVarCt nabla@MkNabla{ nabla_tm_st = TmSt{ ts_facts = env } } x y
-  -- It's important that we never @equate@ two variables of the same equivalence
-  -- class, otherwise we might get cyclic substitutions.
-  -- Cf. 'extendSubstAndSolve' and
-  -- @testsuite/tests/pmcheck/should_compile/CyclicSubst.hs@.
-  | sameRepresentativeSDIE env x y = pure nabla
-  | otherwise                      = equate nabla x y
-
--- | @equate ts@(TmSt env) x y@ merges the equivalence classes of @x@ and @y@ by
--- adding an indirection to the environment.
--- Makes sure that the positive and negative facts of @x@ and @y@ are
--- compatible.
--- Preconditions: @not (sameRepresentativeSDIE env x y)@
---
--- See Note [TmState invariants].
-equate :: Nabla -> Id -> Id -> MaybeT DsM Nabla
-equate nabla@MkNabla{ nabla_tm_st = ts@TmSt{ts_facts = env} } x y
-  = ASSERT( not (sameRepresentativeSDIE env x y) )
-    case (lookupSDIE env x, lookupSDIE env y) of
-      (Nothing, _) -> pure (nabla{ nabla_tm_st = ts{ ts_facts = setIndirectSDIE env x y } })
-      (_, Nothing) -> pure (nabla{ nabla_tm_st = ts{ ts_facts = setIndirectSDIE env y x } })
-      -- Merge the info we have for x into the info for y
-      (Just vi_x, Just vi_y) -> do
-        -- This assert will probably trigger at some point...
-        -- We should decide how to break the tie
-        MASSERT2( idType (vi_id vi_x) `eqType` idType (vi_id vi_y), text "Not same type" )
-        -- First assume that x and y are in the same equivalence class
-        let env_ind = setIndirectSDIE env x y
-        -- Then sum up the refinement counters
-        let env_refs = setEntrySDIE env_ind y vi_y
-        let nabla_refs = nabla{ nabla_tm_st = ts{ts_facts = env_refs} }
-        -- and then gradually merge every positive fact we have on x into y
-        let add_fact nabla (PACA cl tvs args) = addConCt nabla y cl tvs args
-        nabla_pos <- foldlM add_fact nabla_refs (vi_pos vi_x)
-        -- Do the same for negative info
-        let add_refut nabla nalt = addNotConCt nabla y nalt
-        nabla_neg <- foldlM add_refut nabla_pos (pmAltConSetElems (vi_neg vi_x))
-        -- vi_rcm will be updated in addNotConCt, so we are good to
-        -- go!
-        pure nabla_neg
+addVarCt nabla@MkNabla{ nabla_tm_st = ts@TmSt{ ts_facts = env } } x y =
+  case equateUSDFM env x y of
+    (Nothing,   env') -> pure (nabla{ nabla_tm_st = ts{ ts_facts = env' } })
+    -- Add the constraints we had for x to y
+    (Just vi_x, env') -> do
+      let nabla_equated = nabla{ nabla_tm_st = ts{ts_facts = env'} }
+      -- and then gradually merge every positive fact we have on x into y
+      let add_pos nabla (PACA cl tvs args) = addConCt nabla y cl tvs args
+      nabla_pos <- foldlM add_pos nabla_equated (vi_pos vi_x)
+      -- Do the same for negative info
+      let add_neg nabla nalt = addNotConCt nabla y nalt
+      foldlM add_neg nabla_pos (pmAltConSetElems (vi_neg vi_x))
 
 -- | Inspects a 'PmCoreCt' @let x = e@ by recording constraints for @x@ based
 -- on the shape of the 'CoreExpr' @e@. Examples:
@@ -1221,11 +1189,11 @@ traverseDirty f ts@TmSt{ts_facts = env, ts_dirty = dirty} =
     go []     env  = pure ts{ts_facts=env}
     go (x:xs) !env = do
       vi' <- f (lookupVarInfo ts x)
-      go xs (setEntrySDIE env x vi')
+      go xs (addToUSDFM env x vi')
 
 traverseAll :: Monad m => (VarInfo -> m VarInfo) -> TmState -> m TmState
 traverseAll f ts@TmSt{ts_facts = env} = do
-  env' <- traverseSDIE f env
+  env' <- traverseUSDFM f env
   pure ts{ts_facts = env'}
 
 -- | Makes sure the given 'Nabla' is still inhabited, by trying to instantiate
diff --git a/compiler/GHC/HsToCore/Pmc/Solver/Types.hs b/compiler/GHC/HsToCore/Pmc/Solver/Types.hs
index 37829b9936..f6fea1fbbf 100644
--- a/compiler/GHC/HsToCore/Pmc/Solver/Types.hs
+++ b/compiler/GHC/HsToCore/Pmc/Solver/Types.hs
@@ -14,10 +14,6 @@ module GHC.HsToCore.Pmc.Solver.Types (
         BotInfo(..), PmAltConApp(..), VarInfo(..), TmState(..), TyState(..),
         Nabla(..), Nablas(..), initNablas,
 
-        -- ** A 'DIdEnv' where entries may be shared
-        Shared(..), SharedDIdEnv(..), emptySDIE, lookupSDIE, sameRepresentativeSDIE,
-        setIndirectSDIE, setEntrySDIE, traverseSDIE, entriesSDIE,
-
         -- ** Caching residual COMPLETE sets
         ConLikeSet, ResidualCompleteMatches(..), getRcm, isRcmInitialised,
 
@@ -46,10 +42,9 @@ import GHC.Utils.Misc
 import GHC.Data.Bag
 import GHC.Data.FastString
 import GHC.Types.Id
-import GHC.Types.Var.Env
 import GHC.Types.Var.Set
 import GHC.Types.Unique.DSet
-import GHC.Types.Unique.DFM
+import GHC.Types.Unique.SDFM
 import GHC.Types.Name
 import GHC.Core.DataCon
 import GHC.Core.ConLike
@@ -135,7 +130,7 @@ initTyState = TySt 0 emptyInert
 -- See Note [TmState invariants] in "GHC.HsToCore.Pmc.Solver".
 data TmState
   = TmSt
-  { ts_facts :: !(SharedDIdEnv VarInfo)
+  { ts_facts :: !(UniqSDFM Id VarInfo)
   -- ^ Facts about term variables. Deterministic env, so that we generate
   -- deterministic error messages.
   , ts_reps  :: !(CoreMap Id)
@@ -245,75 +240,7 @@ instance Outputable VarInfo where
 
 -- | Initial state of the term oracle.
 initTmState :: TmState
-initTmState = TmSt emptySDIE emptyCoreMap emptyDVarSet
-
--- ** A 'DIdEnv' where entries may be shared
-
--- | Either @Indirect x@, meaning the value is represented by that of @x@, or
--- an @Entry@ containing containing the actual value it represents.
-data Shared a
-  = Indirect !Id
-  | Entry !a
-
--- | A 'DIdEnv' in which entries can be shared by multiple 'Id's.
--- Merge equivalence classes of two Ids by 'setIndirectSDIE' and set the entry
--- of an Id with 'setEntrySDIE'.
-newtype SharedDIdEnv a
-  = SDIE { unSDIE :: DIdEnv (Shared a) }
-
-emptySDIE :: SharedDIdEnv a
-emptySDIE = SDIE emptyDVarEnv
-
-lookupReprAndEntrySDIE :: SharedDIdEnv a -> Id -> (Id, Maybe a)
-lookupReprAndEntrySDIE sdie@(SDIE env) x = case lookupDVarEnv env x of
-  Nothing           -> (x, Nothing)
-  Just (Indirect y) -> lookupReprAndEntrySDIE sdie y
-  Just (Entry a)    -> (x, Just a)
-
--- | @lookupSDIE env x@ looks up an entry for @x@, looking through all
--- 'Indirect's until it finds a shared 'Entry'.
-lookupSDIE :: SharedDIdEnv a -> Id -> Maybe a
-lookupSDIE sdie x = snd (lookupReprAndEntrySDIE sdie x)
-
--- | Check if two variables are part of the same equivalence class.
-sameRepresentativeSDIE :: SharedDIdEnv a -> Id -> Id -> Bool
-sameRepresentativeSDIE sdie x y =
-  fst (lookupReprAndEntrySDIE sdie x) == fst (lookupReprAndEntrySDIE sdie y)
-
--- | @setIndirectSDIE env x y@ sets @x@'s 'Entry' to @Indirect y@, thereby
--- merging @x@'s equivalence class into @y@'s. This will discard all info on
--- @x@!
-setIndirectSDIE :: SharedDIdEnv a -> Id -> Id -> SharedDIdEnv a
-setIndirectSDIE sdie@(SDIE env) x y =
-  SDIE $ extendDVarEnv env (fst (lookupReprAndEntrySDIE sdie x)) (Indirect y)
-
--- | @setEntrySDIE env x a@ sets the 'Entry' @x@ is associated with to @a@,
--- thereby modifying its whole equivalence class.
-setEntrySDIE :: SharedDIdEnv a -> Id -> a -> SharedDIdEnv a
-setEntrySDIE sdie@(SDIE env) x a =
-  SDIE $ extendDVarEnv env (fst (lookupReprAndEntrySDIE sdie x)) (Entry a)
-
-entriesSDIE :: SharedDIdEnv a -> [a]
-entriesSDIE (SDIE env) = mapMaybe preview_entry (eltsUDFM env)
-  where
-    preview_entry (Entry e) = Just e
-    preview_entry _         = Nothing
-
-traverseSDIE :: forall a b f. Applicative f => (a -> f b) -> SharedDIdEnv a -> f (SharedDIdEnv b)
-traverseSDIE f = fmap (SDIE . listToUDFM_Directly) . traverse g . udfmToList . unSDIE
-  where
-    g :: (Unique, Shared a) -> f (Unique, Shared b)
-    g (u, Indirect y) = pure (u,Indirect y)
-    g (u, Entry a)    = do
-        a' <- f a
-        pure (u,Entry a')
-
-instance Outputable a => Outputable (Shared a) where
-  ppr (Indirect x) = ppr x
-  ppr (Entry a)    = ppr a
-
-instance Outputable a => Outputable (SharedDIdEnv a) where
-  ppr (SDIE env) = ppr env
+initTmState = TmSt emptyUSDFM emptyCoreMap emptyDVarSet
 
 -- | A data type that caches for the 'VarInfo' of @x@ the results of querying
 -- 'dsGetCompleteMatches' and then striking out all occurrences of @K@ for
diff --git a/compiler/GHC/Types/Unique/SDFM.hs b/compiler/GHC/Types/Unique/SDFM.hs
new file mode 100644
index 0000000000..a0871909ed
--- /dev/null
+++ b/compiler/GHC/Types/Unique/SDFM.hs
@@ -0,0 +1,121 @@
+{-# 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 { unUSDFM :: UniqDFM key (Shared key ele) }
+
+emptyUSDFM :: UniqSDFM key ele
+emptyUSDFM = USDFM emptyUDFM
+
+lookupReprAndEntryUSDFM :: Uniquable key => UniqSDFM key ele -> key -> (key, Maybe ele)
+lookupReprAndEntryUSDFM (USDFM env) = go
+  where
+    go x = case lookupUDFM env x of
+      Nothing           -> (x, Nothing)
+      Just (Indirect y) -> go y
+      Just (Entry ele)  -> (x, Just 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 usdfm x = snd (lookupReprAndEntryUSDFM usdfm 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 usdfm@(USDFM env) x y =
+  case (lu x, lu y) of
+    ((x', _)    , (y', _))
+      | getUnique x' == getUnique y' -> (Nothing, usdfm) -- nothing to do
+    ((x', _)    , (_ , Nothing))     -> (Nothing, set_indirect y x')
+    ((_ , mb_ex), (y', _))           -> (mb_ex,   set_indirect x y')
+  where
+    lu = lookupReprAndEntryUSDFM usdfm
+    set_indirect a b = USDFM $ addToUDFM env a (Indirect 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 usdfm@(USDFM env) x v =
+  USDFM $ addToUDFM env (fst (lookupReprAndEntryUSDFM usdfm x)) (Entry v)
+
+traverseUSDFM :: forall key a b f. Applicative f => (a -> f b) -> UniqSDFM key a -> f (UniqSDFM key b)
+traverseUSDFM f = fmap (USDFM . listToUDFM_Directly) . traverse g . udfmToList . unUSDFM
+  where
+    g :: (Unique, Shared key a) -> f (Unique, Shared key b)
+    g (u, Indirect y) = pure (u,Indirect y)
+    g (u, Entry a)    = do
+        a' <- f a
+        pure (u,Entry a')
+
+instance (Outputable key, Outputable ele) => Outputable (Shared key ele) where
+  ppr (Indirect x) = ppr x
+  ppr (Entry a)    = ppr a
+
+instance (Outputable key, Outputable ele) => Outputable (UniqSDFM key ele) where
+  ppr (USDFM env) = ppr env
diff --git a/compiler/ghc.cabal.in b/compiler/ghc.cabal.in
index 918c6f7629..a80b9947a4 100644
--- a/compiler/ghc.cabal.in
+++ b/compiler/ghc.cabal.in
@@ -569,6 +569,7 @@ Library
         GHC.Data.Stream
         GHC.Data.StringBuffer
         GHC.Types.Unique.DFM
+        GHC.Types.Unique.SDFM
         GHC.Types.Unique.DSet
         GHC.Types.Unique.FM
         GHC.Types.Unique.Set