    {-

    Copyright (c) 2014 Joachim Breitner

    A data structure for undirected graphs of variables
    (or in plain terms: Sets of unordered pairs of numbers)

    This is very specifically tailored for the use in CallArity. In particular it
    stores the graph as a union of complete and complete bipartite graph, which
    would be very expensive to store as sets of edges or as adjanceny lists.

    It does not normalize the graphs. This means that g `unionUnVarGraph` g is
    equal to g, but twice as expensive and large.

    -}
    module GHC.Data.Graph.UnVar
        ( UnVarSet
        , emptyUnVarSet, mkUnVarSet, varEnvDom, unionUnVarSet, unionUnVarSets
        , delUnVarSet
        , elemUnVarSet, isEmptyUnVarSet
        , UnVarGraph
        , emptyUnVarGraph
        , unionUnVarGraph, unionUnVarGraphs
        , completeGraph, completeBipartiteGraph
        , neighbors
        , hasLoopAt
        , delNode
        ) where

    import GHC.Prelude

    import GHC.Types.Id
    import GHC.Types.Var.Env
    import GHC.Types.Unique.FM
    import GHC.Utils.Outputable
    import GHC.Data.Bag
    import GHC.Types.Unique

    import qualified Data.IntSet as S

    -- We need a type for sets of variables (UnVarSet).
    -- We do not use VarSet, because for that we need to have the actual variable
    -- at hand, and we do not have that when we turn the domain of a VarEnv into a UnVarSet.
    -- Therefore, use a IntSet directly (which is likely also a bit more efficient).

    -- Set of uniques, i.e. for adjancet nodes
    newtype UnVarSet = UnVarSet (S.IntSet)
        deriving UnVarSet -> UnVarSet -> Bool
    (UnVarSet -> UnVarSet -> Bool)
    -> (UnVarSet -> UnVarSet -> Bool) -> Eq UnVarSet
    forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
    /= :: UnVarSet -> UnVarSet -> Bool
    $c/= :: UnVarSet -> UnVarSet -> Bool
    == :: UnVarSet -> UnVarSet -> Bool
    $c== :: UnVarSet -> UnVarSet -> Bool
    Eq

    k :: Var -> Int
    k :: Var -> Int
    k Var
    v = Unique -> Int
    getKey (Var -> Unique
    forall a. Uniquable a => a -> Unique
    getUnique Var
    v)

    emptyUnVarSet :: UnVarSet
    emptyUnVarSet :: UnVarSet
    emptyUnVarSet = IntSet -> UnVarSet
    UnVarSet IntSet
    S.empty

    elemUnVarSet :: Var -> UnVarSet -> Bool
    elemUnVarSet :: Var -> UnVarSet -> Bool
    elemUnVarSet Var
    v (UnVarSet IntSet
    s) = Var -> Int
    k Var
    v Int -> IntSet -> Bool
    `S.member` IntSet
    s

    isEmptyUnVarSet :: UnVarSet -> Bool
    isEmptyUnVarSet :: UnVarSet -> Bool
    isEmptyUnVarSet (UnVarSet IntSet
    s) = IntSet -> Bool
    S.null IntSet
    s

    delUnVarSet :: UnVarSet -> Var -> UnVarSet
    delUnVarSet :: UnVarSet -> Var -> UnVarSet
    delUnVarSet (UnVarSet IntSet
    s) Var
    v = IntSet -> UnVarSet
    UnVarSet (IntSet -> UnVarSet) -> IntSet -> UnVarSet
    forall a b. (a -> b) -> a -> b
    $ Var -> Int
    k Var
    v Int -> IntSet -> IntSet
    `S.delete` IntSet
    s

    mkUnVarSet :: [Var] -> UnVarSet
    mkUnVarSet :: [Var] -> UnVarSet
    mkUnVarSet [Var]
    vs = IntSet -> UnVarSet
    UnVarSet (IntSet -> UnVarSet) -> IntSet -> UnVarSet
    forall a b. (a -> b) -> a -> b
    $ [Int] -> IntSet
    S.fromList ([Int] -> IntSet) -> [Int] -> IntSet
    forall a b. (a -> b) -> a -> b
    $ (Var -> Int) -> [Var] -> [Int]
    forall a b. (a -> b) -> [a] -> [b]
    map Var -> Int
    k [Var]
    vs

    varEnvDom :: VarEnv a -> UnVarSet
    varEnvDom :: forall a. VarEnv a -> UnVarSet
    varEnvDom VarEnv a
    ae = IntSet -> UnVarSet
    UnVarSet (IntSet -> UnVarSet) -> IntSet -> UnVarSet
    forall a b. (a -> b) -> a -> b
    $ VarEnv a -> IntSet
    forall key elt. UniqFM key elt -> IntSet
    ufmToSet_Directly VarEnv a
    ae

    unionUnVarSet :: UnVarSet -> UnVarSet -> UnVarSet
    unionUnVarSet :: UnVarSet -> UnVarSet -> UnVarSet
    unionUnVarSet (UnVarSet IntSet
    set1) (UnVarSet IntSet
    set2) = IntSet -> UnVarSet
    UnVarSet (IntSet
    set1 IntSet -> IntSet -> IntSet
    `S.union` IntSet
    set2)

    unionUnVarSets :: [UnVarSet] -> UnVarSet
    unionUnVarSets :: [UnVarSet] -> UnVarSet
    unionUnVarSets = (UnVarSet -> UnVarSet -> UnVarSet)
    -> UnVarSet -> [UnVarSet] -> UnVarSet
    forall (t :: * -> *) a b.
    Foldable t =>
    (a -> b -> b) -> b -> t a -> b
    foldr UnVarSet -> UnVarSet -> UnVarSet
    unionUnVarSet UnVarSet
    emptyUnVarSet

    instance Outputable UnVarSet where
        ppr :: UnVarSet -> SDoc
    ppr (UnVarSet IntSet
    s) = SDoc -> SDoc
    braces (SDoc -> SDoc) -> SDoc -> SDoc
    forall a b. (a -> b) -> a -> b
    $
            [SDoc] -> SDoc
    hcat ([SDoc] -> SDoc) -> [SDoc] -> SDoc
    forall a b. (a -> b) -> a -> b
    $ SDoc -> [SDoc] -> [SDoc]
    punctuate SDoc
    comma [ Unique -> SDoc
    forall a. Outputable a => a -> SDoc
    ppr (Int -> Unique
    forall a. Uniquable a => a -> Unique
    getUnique Int
    i) | Int
    i <- IntSet -> [Int]
    S.toList IntSet
    s]

    -- The graph type. A list of complete bipartite graphs
    data Gen = CBPG UnVarSet UnVarSet -- complete bipartite
             | CG   UnVarSet          -- complete
    newtype UnVarGraph = UnVarGraph (Bag Gen)

    emptyUnVarGraph :: UnVarGraph
    emptyUnVarGraph :: UnVarGraph
    emptyUnVarGraph = Bag Gen -> UnVarGraph
    UnVarGraph Bag Gen
    forall a. Bag a
    emptyBag

    unionUnVarGraph :: UnVarGraph -> UnVarGraph -> UnVarGraph
    {-
    Premature optimisation, it seems.
    unionUnVarGraph (UnVarGraph [CBPG s1 s2]) (UnVarGraph [CG s3, CG s4])
        | s1 == s3 && s2 == s4
        = pprTrace "unionUnVarGraph fired" empty $
          completeGraph (s1 `unionUnVarSet` s2)
    unionUnVarGraph (UnVarGraph [CBPG s1 s2]) (UnVarGraph [CG s3, CG s4])
        | s2 == s3 && s1 == s4
        = pprTrace "unionUnVarGraph fired2" empty $
          completeGraph (s1 `unionUnVarSet` s2)
    -}
    unionUnVarGraph :: UnVarGraph -> UnVarGraph -> UnVarGraph
    unionUnVarGraph (UnVarGraph Bag Gen
    g1) (UnVarGraph Bag Gen
    g2)
        = -- pprTrace "unionUnVarGraph" (ppr (length g1, length g2)) $
          Bag Gen -> UnVarGraph
    UnVarGraph (Bag Gen
    g1 Bag Gen -> Bag Gen -> Bag Gen
    forall a. Bag a -> Bag a -> Bag a
    `unionBags` Bag Gen
    g2)

    unionUnVarGraphs :: [UnVarGraph] -> UnVarGraph
    unionUnVarGraphs :: [UnVarGraph] -> UnVarGraph
    unionUnVarGraphs = (UnVarGraph -> UnVarGraph -> UnVarGraph)
    -> UnVarGraph -> [UnVarGraph] -> UnVarGraph
    forall (t :: * -> *) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    foldl' UnVarGraph -> UnVarGraph -> UnVarGraph
    unionUnVarGraph UnVarGraph
    emptyUnVarGraph

    -- completeBipartiteGraph A B = { {a,b} | a ∈ A, b ∈ B }
    completeBipartiteGraph :: UnVarSet -> UnVarSet -> UnVarGraph
    completeBipartiteGraph :: UnVarSet -> UnVarSet -> UnVarGraph
    completeBipartiteGraph UnVarSet
    s1 UnVarSet
    s2 = UnVarGraph -> UnVarGraph
    prune (UnVarGraph -> UnVarGraph) -> UnVarGraph -> UnVarGraph
    forall a b. (a -> b) -> a -> b
    $ Bag Gen -> UnVarGraph
    UnVarGraph (Bag Gen -> UnVarGraph) -> Bag Gen -> UnVarGraph
    forall a b. (a -> b) -> a -> b
    $ Gen -> Bag Gen
    forall a. a -> Bag a
    unitBag (Gen -> Bag Gen) -> Gen -> Bag Gen
    forall a b. (a -> b) -> a -> b
    $ UnVarSet -> UnVarSet -> Gen
    CBPG UnVarSet
    s1 UnVarSet
    s2

    completeGraph :: UnVarSet -> UnVarGraph
    completeGraph :: UnVarSet -> UnVarGraph
    completeGraph UnVarSet
    s = UnVarGraph -> UnVarGraph
    prune (UnVarGraph -> UnVarGraph) -> UnVarGraph -> UnVarGraph
    forall a b. (a -> b) -> a -> b
    $ Bag Gen -> UnVarGraph
    UnVarGraph (Bag Gen -> UnVarGraph) -> Bag Gen -> UnVarGraph
    forall a b. (a -> b) -> a -> b
    $ Gen -> Bag Gen
    forall a. a -> Bag a
    unitBag (Gen -> Bag Gen) -> Gen -> Bag Gen
    forall a b. (a -> b) -> a -> b
    $ UnVarSet -> Gen
    CG UnVarSet
    s

    neighbors :: UnVarGraph -> Var -> UnVarSet
    neighbors :: UnVarGraph -> Var -> UnVarSet
    neighbors (UnVarGraph Bag Gen
    g) Var
    v = [UnVarSet] -> UnVarSet
    unionUnVarSets ([UnVarSet] -> UnVarSet) -> [UnVarSet] -> UnVarSet
    forall a b. (a -> b) -> a -> b
    $ (Gen -> [UnVarSet]) -> [Gen] -> [UnVarSet]
    forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
    concatMap Gen -> [UnVarSet]
    go ([Gen] -> [UnVarSet]) -> [Gen] -> [UnVarSet]
    forall a b. (a -> b) -> a -> b
    $ Bag Gen -> [Gen]
    forall a. Bag a -> [a]
    bagToList Bag Gen
    g
      where go :: Gen -> [UnVarSet]
    go (CG UnVarSet
    s)       = (if Var
    v Var -> UnVarSet -> Bool
    `elemUnVarSet` UnVarSet
    s then [UnVarSet
    s] else [])
            go (CBPG UnVarSet
    s1 UnVarSet
    s2) = (if Var
    v Var -> UnVarSet -> Bool
    `elemUnVarSet` UnVarSet
    s1 then [UnVarSet
    s2] else []) [UnVarSet] -> [UnVarSet] -> [UnVarSet]
    forall a. [a] -> [a] -> [a]
    ++
                              (if Var
    v Var -> UnVarSet -> Bool
    `elemUnVarSet` UnVarSet
    s2 then [UnVarSet
    s1] else [])

    -- hasLoopAt G v <=> v--v ∈ G
    hasLoopAt :: UnVarGraph -> Var -> Bool
    hasLoopAt :: UnVarGraph -> Var -> Bool
    hasLoopAt (UnVarGraph Bag Gen
    g) Var
    v = (Gen -> Bool) -> [Gen] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    any Gen -> Bool
    go ([Gen] -> Bool) -> [Gen] -> Bool
    forall a b. (a -> b) -> a -> b
    $ Bag Gen -> [Gen]
    forall a. Bag a -> [a]
    bagToList Bag Gen
    g
      where go :: Gen -> Bool
    go (CG UnVarSet
    s)       = Var
    v Var -> UnVarSet -> Bool
    `elemUnVarSet` UnVarSet
    s
            go (CBPG UnVarSet
    s1 UnVarSet
    s2) = Var
    v Var -> UnVarSet -> Bool
    `elemUnVarSet` UnVarSet
    s1 Bool -> Bool -> Bool
    && Var
    v Var -> UnVarSet -> Bool
    `elemUnVarSet` UnVarSet
    s2

    delNode :: UnVarGraph -> Var -> UnVarGraph
    delNode :: UnVarGraph -> Var -> UnVarGraph
    delNode (UnVarGraph Bag Gen
    g) Var
    v = UnVarGraph -> UnVarGraph
    prune (UnVarGraph -> UnVarGraph) -> UnVarGraph -> UnVarGraph
    forall a b. (a -> b) -> a -> b
    $ Bag Gen -> UnVarGraph
    UnVarGraph (Bag Gen -> UnVarGraph) -> Bag Gen -> UnVarGraph
    forall a b. (a -> b) -> a -> b
    $ (Gen -> Gen) -> Bag Gen -> Bag Gen
    forall a b. (a -> b) -> Bag a -> Bag b
    mapBag Gen -> Gen
    go Bag Gen
    g
      where go :: Gen -> Gen
    go (CG UnVarSet
    s)       = UnVarSet -> Gen
    CG (UnVarSet
    s UnVarSet -> Var -> UnVarSet
    `delUnVarSet` Var
    v)
            go (CBPG UnVarSet
    s1 UnVarSet
    s2) = UnVarSet -> UnVarSet -> Gen
    CBPG (UnVarSet
    s1 UnVarSet -> Var -> UnVarSet
    `delUnVarSet` Var
    v) (UnVarSet
    s2 UnVarSet -> Var -> UnVarSet
    `delUnVarSet` Var
    v)

    prune :: UnVarGraph -> UnVarGraph
    prune :: UnVarGraph -> UnVarGraph
    prune (UnVarGraph Bag Gen
    g) = Bag Gen -> UnVarGraph
    UnVarGraph (Bag Gen -> UnVarGraph) -> Bag Gen -> UnVarGraph
    forall a b. (a -> b) -> a -> b
    $ (Gen -> Bool) -> Bag Gen -> Bag Gen
    forall a. (a -> Bool) -> Bag a -> Bag a
    filterBag Gen -> Bool
    go Bag Gen
    g
      where go :: Gen -> Bool
    go (CG UnVarSet
    s)       = Bool -> Bool
    not (UnVarSet -> Bool
    isEmptyUnVarSet UnVarSet
    s)
            go (CBPG UnVarSet
    s1 UnVarSet
    s2) = Bool -> Bool
    not (UnVarSet -> Bool
    isEmptyUnVarSet UnVarSet
    s1) Bool -> Bool -> Bool
    && Bool -> Bool
    not (UnVarSet -> Bool
    isEmptyUnVarSet UnVarSet
    s2)

    instance Outputable Gen where
        ppr :: Gen -> SDoc
    ppr (CG UnVarSet
    s)       = UnVarSet -> SDoc
    forall a. Outputable a => a -> SDoc
    ppr UnVarSet
    s  SDoc -> SDoc -> SDoc
    <> Char -> SDoc
    char Char
    '²'
        ppr (CBPG UnVarSet
    s1 UnVarSet
    s2) = UnVarSet -> SDoc
    forall a. Outputable a => a -> SDoc
    ppr UnVarSet
    s1 SDoc -> SDoc -> SDoc
    <+> Char -> SDoc
    char Char
    'x' SDoc -> SDoc -> SDoc
    <+> UnVarSet -> SDoc
    forall a. Outputable a => a -> SDoc
    ppr UnVarSet
    s2
    instance Outputable UnVarGraph where
        ppr :: UnVarGraph -> SDoc
    ppr (UnVarGraph Bag Gen
    g) = Bag Gen -> SDoc
    forall a. Outputable a => a -> SDoc
    ppr Bag Gen
    g