    {-# LANGUAGE BangPatterns #-}
    {-# LANGUAGE NamedFieldPuns #-}
    {-# LANGUAGE RecordWildCards #-}
    {-----------------------------------------------------------------------------
        reactive-banana
    ------------------------------------------------------------------------------}
    module Reactive.Banana.Prim.Low.GraphGC
        ( GraphGC
        , listReachableVertices
        , getSize
        , new
        , insertEdge
        , clearPredecessors

        , Step (..)
        , walkSuccessors
        , walkSuccessors_

        , removeGarbage

        -- * Debugging
        , printDot
        ) where

    import Control.Applicative
        ( (<|>) )
    import Control.Monad
        ( unless )
    import Data.IORef
        ( IORef, atomicModifyIORef', newIORef, readIORef )
    import Data.Maybe
        ( fromJust )
    import Data.Unique.Really
        ( Unique )
    import Reactive.Banana.Prim.Low.Graph
        ( Graph, Step )
    import Reactive.Banana.Prim.Low.Ref
        ( Ref, WeakRef )

    import qualified Control.Concurrent.STM as STM
    import qualified Data.HashMap.Strict as Map
    import qualified Reactive.Banana.Prim.Low.Graph as Graph
    import qualified Reactive.Banana.Prim.Low.Ref as Ref

    type Map = Map.HashMap

    {-----------------------------------------------------------------------------
        GraphGC
    ------------------------------------------------------------------------------}
    type WeakEdge v = WeakRef v

    -- Graph data
    data GraphD v = GraphD
        { forall v. GraphD v -> Graph Unique (WeakEdge v)
    graph :: !(Graph Unique (WeakEdge v))
        , forall v. GraphD v -> Map Unique (WeakEdge v)
    references :: !(Map Unique (WeakRef v))
        }

    {- | A directed graph whose edges are mutable
        and whose vertices are subject to garbage collection.

        The vertices of the graph are mutable references of type 'Ref v'.

        Generally, the vertices of the graph are not necessarily kept reachable
        by the 'GraphGC' data structure
        — they need to be kept reachable by other parts of your program.

        That said, the edges in the graph do introduce additional reachability
        between vertices:
        Specifically, when an edge (x,y) is present in the graph,
        then the head @y@ will keep the tail @x@ reachable.
        (But the liveness of @y@ needs to come from elsewhere, e.g. another edge.)
        Use 'insertEdge' to insert an edge.

        Moreover, when a vertex is removed because it is no longer reachable,
        then all edges to and from that vertex will also be removed.
        In turn, this may cause further vertices and edges to be removed.

        Concerning garbage collection:
        Note that vertices and edges will not be removed automatically
        when the Haskell garbage collector runs —
        they will be marked as garbage by the Haskell runtime,
        but the actual removal of garbage needs
        to be done explicitly by calling 'removeGarbage'.
        This procedure makes it easier to reason about the state of the 'GraphGC'
        during a call to e.g. 'walkSuccessors'.
    -}
    data GraphGC v = GraphGC
        { forall v. GraphGC v -> IORef (GraphD v)
    graphRef :: IORef (GraphD v)
        , forall v. GraphGC v -> TQueue Unique
    deletions :: STM.TQueue Unique
        }

    -- | Create a new 'GraphGC'.
    new :: IO (GraphGC v)
    new :: forall v. IO (GraphGC v)
    new = IORef (GraphD v) -> TQueue Unique -> GraphGC v
    forall v. IORef (GraphD v) -> TQueue Unique -> GraphGC v
    GraphGC (IORef (GraphD v) -> TQueue Unique -> GraphGC v)
    -> IO (IORef (GraphD v)) -> IO (TQueue Unique -> GraphGC v)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> GraphD v -> IO (IORef (GraphD v))
    forall a. a -> IO (IORef a)
    newIORef GraphD v
    forall {v}. GraphD v
    newGraphD IO (TQueue Unique -> GraphGC v)
    -> IO (TQueue Unique) -> IO (GraphGC v)
    forall a b. IO (a -> b) -> IO a -> IO b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> IO (TQueue Unique)
    forall a. IO (TQueue a)
    STM.newTQueueIO
      where
        newGraphD :: GraphD v
    newGraphD = GraphD
            { graph :: Graph Unique (WeakEdge v)
    graph = Graph Unique (WeakEdge v)
    forall v e. Graph v e
    Graph.empty
            , references :: Map Unique (WeakEdge v)
    references = Map Unique (WeakEdge v)
    forall k v. HashMap k v
    Map.empty
            }

    getSize :: GraphGC v -> IO Int
    getSize :: forall v. GraphGC v -> IO Int
    getSize GraphGC{IORef (GraphD v)
    graphRef :: forall v. GraphGC v -> IORef (GraphD v)
    graphRef :: IORef (GraphD v)
    graphRef} = Graph Unique (WeakEdge v) -> Int
    forall v e. (Eq v, Hashable v) => Graph v e -> Int
    Graph.size (Graph Unique (WeakEdge v) -> Int)
    -> (GraphD v -> Graph Unique (WeakEdge v)) -> GraphD v -> Int
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . GraphD v -> Graph Unique (WeakEdge v)
    forall v. GraphD v -> Graph Unique (WeakEdge v)
    graph (GraphD v -> Int) -> IO (GraphD v) -> IO Int
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> IORef (GraphD v) -> IO (GraphD v)
    forall a. IORef a -> IO a
    readIORef IORef (GraphD v)
    graphRef

    -- | List all vertices that are reachable and have at least
    -- one edge incident on them.
    -- TODO: Is that really what the function does?
    listReachableVertices :: GraphGC v -> IO [Ref v]
    listReachableVertices :: forall v. GraphGC v -> IO [Ref v]
    listReachableVertices GraphGC{IORef (GraphD v)
    graphRef :: forall v. GraphGC v -> IORef (GraphD v)
    graphRef :: IORef (GraphD v)
    graphRef} = do
        GraphD{references} <- IORef (GraphD v) -> IO (GraphD v)
    forall a. IORef a -> IO a
    readIORef IORef (GraphD v)
    graphRef
        concat . Map.elems <$> traverse inspect references
      where
        inspect :: Weak a -> IO [a]
    inspect Weak a
    ref = do
            mv <- Weak a -> IO (Maybe a)
    forall v. Weak v -> IO (Maybe v)
    Ref.deRefWeak Weak a
    ref
            pure $ case mv of
                Maybe a
    Nothing -> []
                Just a
    r -> [a
    r]

    -- | Insert an edge from the first vertex to the second vertex.
    insertEdge :: (Ref v, Ref v) -> GraphGC v -> IO ()
    insertEdge :: forall v. (Ref v, Ref v) -> GraphGC v -> IO ()
    insertEdge (Ref v
    x,Ref v
    y) g :: GraphGC v
    g@GraphGC{IORef (GraphD v)
    graphRef :: forall v. GraphGC v -> IORef (GraphD v)
    graphRef :: IORef (GraphD v)
    graphRef} = do
        (xKnown, yKnown) <-
            WeakEdge v -> IO (Bool, Bool)
    insertTheEdge (WeakEdge v -> IO (Bool, Bool))
    -> IO (WeakEdge v) -> IO (Bool, Bool)
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< IO (WeakEdge v)
    makeWeakPointerThatRepresentsEdge
        unless xKnown $ Ref.addFinalizer x (finalizeVertex g ux)
        unless yKnown $ Ref.addFinalizer y (finalizeVertex g uy)
      where
        ux :: Unique
    ux = Ref v -> Unique
    forall a. Ref a -> Unique
    Ref.getUnique Ref v
    x
        uy :: Unique
    uy = Ref v -> Unique
    forall a. Ref a -> Unique
    Ref.getUnique Ref v
    y

        makeWeakPointerThatRepresentsEdge :: IO (WeakEdge v)
    makeWeakPointerThatRepresentsEdge =
            Ref v -> Ref v -> Maybe (IO ()) -> IO (WeakEdge v)
    forall k v. Ref k -> v -> Maybe (IO ()) -> IO (Weak v)
    Ref.mkWeak Ref v
    y Ref v
    x Maybe (IO ())
    forall a. Maybe a
    Nothing

        insertTheEdge :: WeakEdge v -> IO (Bool, Bool)
    insertTheEdge WeakEdge v
    we = IORef (GraphD v)
    -> (GraphD v -> (GraphD v, (Bool, Bool))) -> IO (Bool, Bool)
    forall a b. IORef a -> (a -> (a, b)) -> IO b
    atomicModifyIORef' IORef (GraphD v)
    graphRef ((GraphD v -> (GraphD v, (Bool, Bool))) -> IO (Bool, Bool))
    -> (GraphD v -> (GraphD v, (Bool, Bool))) -> IO (Bool, Bool)
    forall a b. (a -> b) -> a -> b
    $
            \GraphD{Graph Unique (WeakEdge v)
    graph :: forall v. GraphD v -> Graph Unique (WeakEdge v)
    graph :: Graph Unique (WeakEdge v)
    graph,Map Unique (WeakEdge v)
    references :: forall v. GraphD v -> Map Unique (WeakEdge v)
    references :: Map Unique (WeakEdge v)
    references} ->
                ( GraphD
                    { graph :: Graph Unique (WeakEdge v)
    graph
                        = (Unique, Unique)
    -> WeakEdge v
    -> Graph Unique (WeakEdge v)
    -> Graph Unique (WeakEdge v)
    forall v e.
    (Eq v, Hashable v) =>
    (v, v) -> e -> Graph v e -> Graph v e
    Graph.insertEdge (Unique
    ux,Unique
    uy) WeakEdge v
    we
                        (Graph Unique (WeakEdge v) -> Graph Unique (WeakEdge v))
    -> Graph Unique (WeakEdge v) -> Graph Unique (WeakEdge v)
    forall a b. (a -> b) -> a -> b
    $ Graph Unique (WeakEdge v)
    graph
                    , references :: Map Unique (WeakEdge v)
    references
                        = Unique
    -> WeakEdge v -> Map Unique (WeakEdge v) -> Map Unique (WeakEdge v)
    forall k v.
    (Eq k, Hashable k) =>
    k -> v -> HashMap k v -> HashMap k v
    Map.insert Unique
    ux (Ref v -> WeakEdge v
    forall a. Ref a -> WeakRef a
    Ref.getWeakRef Ref v
    x)
                        (Map Unique (WeakEdge v) -> Map Unique (WeakEdge v))
    -> (Map Unique (WeakEdge v) -> Map Unique (WeakEdge v))
    -> Map Unique (WeakEdge v)
    -> Map Unique (WeakEdge v)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Unique
    -> WeakEdge v -> Map Unique (WeakEdge v) -> Map Unique (WeakEdge v)
    forall k v.
    (Eq k, Hashable k) =>
    k -> v -> HashMap k v -> HashMap k v
    Map.insert Unique
    uy (Ref v -> WeakEdge v
    forall a. Ref a -> WeakRef a
    Ref.getWeakRef Ref v
    y)
                        (Map Unique (WeakEdge v) -> Map Unique (WeakEdge v))
    -> Map Unique (WeakEdge v) -> Map Unique (WeakEdge v)
    forall a b. (a -> b) -> a -> b
    $ Map Unique (WeakEdge v)
    references
                    }
                ,   ( Unique
    ux Unique -> Map Unique (WeakEdge v) -> Bool
    forall k a. (Eq k, Hashable k) => k -> HashMap k a -> Bool
    `Map.member` Map Unique (WeakEdge v)
    references
                    , Unique
    uy Unique -> Map Unique (WeakEdge v) -> Bool
    forall k a. (Eq k, Hashable k) => k -> HashMap k a -> Bool
    `Map.member` Map Unique (WeakEdge v)
    references
                    )
                )

    -- | Remove all the edges that connect the vertex to its predecessors.
    clearPredecessors :: Ref v -> GraphGC v -> IO ()
    clearPredecessors :: forall v. Ref v -> GraphGC v -> IO ()
    clearPredecessors Ref v
    x GraphGC{IORef (GraphD v)
    graphRef :: forall v. GraphGC v -> IORef (GraphD v)
    graphRef :: IORef (GraphD v)
    graphRef} = do
        g <- IORef (GraphD v)
    -> (GraphD v -> (GraphD v, GraphD v)) -> IO (GraphD v)
    forall a b. IORef a -> (a -> (a, b)) -> IO b
    atomicModifyIORef' IORef (GraphD v)
    graphRef ((GraphD v -> (GraphD v, GraphD v)) -> IO (GraphD v))
    -> (GraphD v -> (GraphD v, GraphD v)) -> IO (GraphD v)
    forall a b. (a -> b) -> a -> b
    $ \GraphD v
    g -> (GraphD v -> GraphD v
    forall {v}. GraphD v -> GraphD v
    removeIncomingEdges GraphD v
    g, GraphD v
    g)
        finalizeIncomingEdges g
      where
        removeIncomingEdges :: GraphD v -> GraphD v
    removeIncomingEdges g :: GraphD v
    g@GraphD{Graph Unique (WeakEdge v)
    graph :: forall v. GraphD v -> Graph Unique (WeakEdge v)
    graph :: Graph Unique (WeakEdge v)
    graph} =
            GraphD v
    g{ graph = Graph.clearPredecessors (Ref.getUnique x) graph }
        finalizeIncomingEdges :: GraphD v -> IO ()
    finalizeIncomingEdges GraphD{Graph Unique (WeakEdge v)
    graph :: forall v. GraphD v -> Graph Unique (WeakEdge v)
    graph :: Graph Unique (WeakEdge v)
    graph} =
            ((Unique, WeakEdge v) -> IO ()) -> [(Unique, WeakEdge v)] -> IO ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    (a -> m b) -> t a -> m ()
    mapM_ (WeakEdge v -> IO ()
    forall v. WeakRef v -> IO ()
    Ref.finalize (WeakEdge v -> IO ())
    -> ((Unique, WeakEdge v) -> WeakEdge v)
    -> (Unique, WeakEdge v)
    -> IO ()
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (Unique, WeakEdge v) -> WeakEdge v
    forall a b. (a, b) -> b
    snd) ([(Unique, WeakEdge v)] -> IO ())
    -> (Unique -> [(Unique, WeakEdge v)]) -> Unique -> IO ()
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Graph Unique (WeakEdge v) -> Unique -> [(Unique, WeakEdge v)]
    forall v e. (Eq v, Hashable v) => Graph v e -> v -> [(v, e)]
    Graph.getIncoming Graph Unique (WeakEdge v)
    graph (Unique -> IO ()) -> Unique -> IO ()
    forall a b. (a -> b) -> a -> b
    $ Ref v -> Unique
    forall a. Ref a -> Unique
    Ref.getUnique Ref v
    x

    -- | Walk through all successors. See 'Graph.walkSuccessors'.
    walkSuccessors
        :: Monad m
        => [Ref v] -> (WeakRef v -> m Step) -> GraphGC v -> IO (m [WeakRef v])
    walkSuccessors :: forall (m :: * -> *) v.
    Monad m =>
    [Ref v] -> (WeakRef v -> m Step) -> GraphGC v -> IO (m [WeakRef v])
    walkSuccessors [Ref v]
    roots WeakRef v -> m Step
    step GraphGC{IORef (GraphD v)
    TQueue Unique
    graphRef :: forall v. GraphGC v -> IORef (GraphD v)
    deletions :: forall v. GraphGC v -> TQueue Unique
    graphRef :: IORef (GraphD v)
    deletions :: TQueue Unique
    ..} = do
        GraphD{graph,references} <- IORef (GraphD v) -> IO (GraphD v)
    forall a. IORef a -> IO a
    readIORef IORef (GraphD v)
    graphRef
        let rootsMap = [(Unique, WeakRef v)] -> Map Unique (WeakRef v)
    forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
    Map.fromList
                [ (Ref v -> Unique
    forall a. Ref a -> Unique
    Ref.getUnique Ref v
    r, Ref v -> WeakRef v
    forall a. Ref a -> WeakRef a
    Ref.getWeakRef Ref v
    r) | Ref v
    r <- [Ref v]
    roots ]
            fromUnique Unique
    u = Maybe (WeakRef v) -> WeakRef v
    forall a. HasCallStack => Maybe a -> a
    fromJust (Maybe (WeakRef v) -> WeakRef v) -> Maybe (WeakRef v) -> WeakRef v
    forall a b. (a -> b) -> a -> b
    $
                Unique -> Map Unique (WeakRef v) -> Maybe (WeakRef v)
    forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
    Map.lookup Unique
    u Map Unique (WeakRef v)
    references Maybe (WeakRef v) -> Maybe (WeakRef v) -> Maybe (WeakRef v)
    forall a. Maybe a -> Maybe a -> Maybe a
    forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
    <|> Unique -> Map Unique (WeakRef v) -> Maybe (WeakRef v)
    forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
    Map.lookup Unique
    u Map Unique (WeakRef v)
    rootsMap
        pure
            . fmap (map fromUnique)
            . Graph.walkSuccessors (map Ref.getUnique roots) (step . fromUnique)
            $ graph

    -- | Walk through all successors. See 'Graph.walkSuccessors_'.
    walkSuccessors_ ::
        Monad m => [Ref v] -> (WeakRef v -> m Step) -> GraphGC v -> IO (m ())
    walkSuccessors_ :: forall (m :: * -> *) v.
    Monad m =>
    [Ref v] -> (WeakRef v -> m Step) -> GraphGC v -> IO (m ())
    walkSuccessors_ [Ref v]
    roots WeakRef v -> m Step
    step GraphGC v
    g = do
        action <- [Ref v] -> (WeakRef v -> m Step) -> GraphGC v -> IO (m [WeakRef v])
    forall (m :: * -> *) v.
    Monad m =>
    [Ref v] -> (WeakRef v -> m Step) -> GraphGC v -> IO (m [WeakRef v])
    walkSuccessors [Ref v]
    roots WeakRef v -> m Step
    step GraphGC v
    g
        pure $ action >> pure ()

    {-----------------------------------------------------------------------------
        Garbage Collection
    ------------------------------------------------------------------------------}
    -- | Explicitly remove all vertices and edges that have been marked
    -- as garbage by the Haskell garbage collector.
    removeGarbage :: GraphGC v -> IO ()
    removeGarbage :: forall v. GraphGC v -> IO ()
    removeGarbage g :: GraphGC v
    g@GraphGC{TQueue Unique
    deletions :: forall v. GraphGC v -> TQueue Unique
    deletions :: TQueue Unique
    deletions} = do
        xs <- STM [Unique] -> IO [Unique]
    forall a. STM a -> IO a
    STM.atomically (STM [Unique] -> IO [Unique]) -> STM [Unique] -> IO [Unique]
    forall a b. (a -> b) -> a -> b
    $ TQueue Unique -> STM [Unique]
    forall a. TQueue a -> STM [a]
    STM.flushTQueue TQueue Unique
    deletions
        mapM_ (deleteVertex g) xs

    -- Delete all edges associated with a vertex from the 'GraphGC'.
    --
    -- TODO: Check whether using an IORef is thread-safe.
    -- I think it's fine because we have a single thread that performs deletions.
    deleteVertex :: GraphGC v -> Unique -> IO ()
    deleteVertex :: forall v. GraphGC v -> Unique -> IO ()
    deleteVertex GraphGC{IORef (GraphD v)
    graphRef :: forall v. GraphGC v -> IORef (GraphD v)
    graphRef :: IORef (GraphD v)
    graphRef} Unique
    x =
        IORef (GraphD v) -> (GraphD v -> GraphD v) -> IO ()
    forall a. IORef a -> (a -> a) -> IO ()
    atomicModifyIORef'_ IORef (GraphD v)
    graphRef ((GraphD v -> GraphD v) -> IO ())
    -> (GraphD v -> GraphD v) -> IO ()
    forall a b. (a -> b) -> a -> b
    $ \GraphD{Graph Unique (WeakEdge v)
    graph :: forall v. GraphD v -> Graph Unique (WeakEdge v)
    graph :: Graph Unique (WeakEdge v)
    graph,Map Unique (WeakEdge v)
    references :: forall v. GraphD v -> Map Unique (WeakEdge v)
    references :: Map Unique (WeakEdge v)
    references} -> GraphD
            { graph :: Graph Unique (WeakEdge v)
    graph = Unique -> Graph Unique (WeakEdge v) -> Graph Unique (WeakEdge v)
    forall v e. (Eq v, Hashable v) => v -> Graph v e -> Graph v e
    Graph.deleteVertex Unique
    x Graph Unique (WeakEdge v)
    graph
            , references :: Map Unique (WeakEdge v)
    references = Unique -> Map Unique (WeakEdge v) -> Map Unique (WeakEdge v)
    forall k v. (Eq k, Hashable k) => k -> HashMap k v -> HashMap k v
    Map.delete Unique
    x Map Unique (WeakEdge v)
    references
            }

    -- Finalize a vertex
    finalizeVertex :: GraphGC v -> Unique -> IO ()
    finalizeVertex :: forall v. GraphGC v -> Unique -> IO ()
    finalizeVertex GraphGC{TQueue Unique
    deletions :: forall v. GraphGC v -> TQueue Unique
    deletions :: TQueue Unique
    deletions} =
        STM () -> IO ()
    forall a. STM a -> IO a
    STM.atomically (STM () -> IO ()) -> (Unique -> STM ()) -> Unique -> IO ()
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . TQueue Unique -> Unique -> STM ()
    forall a. TQueue a -> a -> STM ()
    STM.writeTQueue TQueue Unique
    deletions

    {-----------------------------------------------------------------------------
        Debugging
    ------------------------------------------------------------------------------}
    -- | Show the underlying graph in @graphviz@ dot file format.
    printDot :: (Unique -> WeakRef v -> IO String) -> GraphGC v -> IO String
    printDot :: forall v.
    (Unique -> WeakRef v -> IO String) -> GraphGC v -> IO String
    printDot Unique -> WeakRef v -> IO String
    format GraphGC{IORef (GraphD v)
    graphRef :: forall v. GraphGC v -> IORef (GraphD v)
    graphRef :: IORef (GraphD v)
    graphRef} = do
        GraphD{graph,references} <- IORef (GraphD v) -> IO (GraphD v)
    forall a. IORef a -> IO a
    readIORef IORef (GraphD v)
    graphRef
        strings <- Map.traverseWithKey format references
        pure $ Graph.showDot (strings Map.!) graph

    {-----------------------------------------------------------------------------
        Helper functions
    ------------------------------------------------------------------------------}
    -- | Atomically modify an 'IORef' without returning a result.
    atomicModifyIORef'_ :: IORef a -> (a -> a) -> IO ()
    atomicModifyIORef'_ :: forall a. IORef a -> (a -> a) -> IO ()
    atomicModifyIORef'_ IORef a
    ref a -> a
    f = IORef a -> (a -> (a, ())) -> IO ()
    forall a b. IORef a -> (a -> (a, b)) -> IO b
    atomicModifyIORef' IORef a
    ref ((a -> (a, ())) -> IO ()) -> (a -> (a, ())) -> IO ()
    forall a b. (a -> b) -> a -> b
    $ \a
    x -> (a -> a
    f a
    x, ())