    {-# 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