{-----------------------------------------------------------------------------
    reactive-banana

    Implementation of graph-related functionality
------------------------------------------------------------------------------}
{-# language ScopedTypeVariables#-}

module Reactive.Banana.Prim.Low.Graph
  ( Graph
  , emptyGraph
  , insertEdge
  , getChildren
  , getParents
  , listParents
  , reversePostOrder
  ) where

import           Data.Functor.Identity
import           Data.Hashable
import qualified Data.HashMap.Strict   as Map
import qualified Data.HashSet          as Set
import           Data.Maybe

{-----------------------------------------------------------------------------
    Graphs and topological sorting
------------------------------------------------------------------------------}
data Graph a = Graph
    { -- | The mapping from each node to the set of nodes reachable by an out-edge. If a node has no out-edges, it is
      -- not a member of this map.
      --
      -- Invariant: the values are non-empty lists.
      Graph a -> HashMap a [a]
children :: Map.HashMap a [a]
      -- | The Mapping from each node to the set of nodes reachable by an in-edge. If a node has no in-edges, it is not
      -- a member of this map.
      --
      -- Invariant: the values are non-empty lists.
    , Graph a -> HashMap a [a]
parents  :: Map.HashMap a [a]
      -- | The set of nodes.
      --
      -- Invariant: equals (key children `union` keys parents)
    , Graph a -> HashSet a
nodes    :: Set.HashSet a
    }

-- | The graph with no edges and no nodes.
emptyGraph :: Graph a
emptyGraph :: Graph a
emptyGraph = HashMap a [a] -> HashMap a [a] -> HashSet a -> Graph a
forall a. HashMap a [a] -> HashMap a [a] -> HashSet a -> Graph a
Graph HashMap a [a]
forall k v. HashMap k v
Map.empty HashMap a [a]
forall k v. HashMap k v
Map.empty HashSet a
forall a. HashSet a
Set.empty

-- | Insert an edge from the first node to the second node into the graph.
insertEdge :: (Eq a, Hashable a) => (a,a) -> Graph a -> Graph a
insertEdge :: (a, a) -> Graph a -> Graph a
insertEdge (a
x,a
y) Graph a
gr = Graph a
gr
    { children :: HashMap a [a]
children = ([a] -> [a] -> [a]) -> a -> [a] -> HashMap a [a] -> HashMap a [a]
forall k v.
(Eq k, Hashable k) =>
(v -> v -> v) -> k -> v -> HashMap k v -> HashMap k v
Map.insertWith (\[a]
new [a]
old -> [a]
new [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a]
old) a
x [a
y] (Graph a -> HashMap a [a]
forall a. Graph a -> HashMap a [a]
children Graph a
gr)
    , parents :: HashMap a [a]
parents  = ([a] -> [a] -> [a]) -> a -> [a] -> HashMap a [a] -> HashMap a [a]
forall k v.
(Eq k, Hashable k) =>
(v -> v -> v) -> k -> v -> HashMap k v -> HashMap k v
Map.insertWith (\[a]
new [a]
old -> [a]
new [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a]
old) a
y [a
x] (Graph a -> HashMap a [a]
forall a. Graph a -> HashMap a [a]
parents  Graph a
gr)
    , nodes :: HashSet a
nodes    = a -> HashSet a -> HashSet a
forall a. (Eq a, Hashable a) => a -> HashSet a -> HashSet a
Set.insert a
x (HashSet a -> HashSet a) -> HashSet a -> HashSet a
forall a b. (a -> b) -> a -> b
$ a -> HashSet a -> HashSet a
forall a. (Eq a, Hashable a) => a -> HashSet a -> HashSet a
Set.insert a
y (HashSet a -> HashSet a) -> HashSet a -> HashSet a
forall a b. (a -> b) -> a -> b
$ Graph a -> HashSet a
forall a. Graph a -> HashSet a
nodes Graph a
gr
    }

-- | Get all immediate children of a node in a graph.
getChildren :: (Eq a, Hashable a) => Graph a -> a -> [a]
getChildren :: Graph a -> a -> [a]
getChildren Graph a
gr a
x = [a] -> Maybe [a] -> [a]
forall a. a -> Maybe a -> a
fromMaybe [] (Maybe [a] -> [a]) -> (Graph a -> Maybe [a]) -> Graph a -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> HashMap a [a] -> Maybe [a]
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
Map.lookup a
x (HashMap a [a] -> Maybe [a])
-> (Graph a -> HashMap a [a]) -> Graph a -> Maybe [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Graph a -> HashMap a [a]
forall a. Graph a -> HashMap a [a]
children (Graph a -> [a]) -> Graph a -> [a]
forall a b. (a -> b) -> a -> b
$ Graph a
gr

-- | Get all immediate parents of a node in a graph.
getParents :: (Eq a, Hashable a) => Graph a -> a -> [a]
getParents :: Graph a -> a -> [a]
getParents Graph a
gr a
x = [a] -> Maybe [a] -> [a]
forall a. a -> Maybe a -> a
fromMaybe [] (Maybe [a] -> [a]) -> (Graph a -> Maybe [a]) -> Graph a -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> HashMap a [a] -> Maybe [a]
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
Map.lookup a
x (HashMap a [a] -> Maybe [a])
-> (Graph a -> HashMap a [a]) -> Graph a -> Maybe [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Graph a -> HashMap a [a]
forall a. Graph a -> HashMap a [a]
parents (Graph a -> [a]) -> Graph a -> [a]
forall a b. (a -> b) -> a -> b
$ Graph a
gr

-- | List all nodes such that each parent is listed before all of its children.
listParents :: forall a. (Eq a, Hashable a) => Graph a -> [a]
listParents :: Graph a -> [a]
listParents Graph a
gr = [a]
list
    where
    -- all nodes without parents
    ancestors :: [a]
    -- We can filter from `children`, because a node without incoming edges can only be in the graph if it has outgoing edges.
    ancestors :: [a]
ancestors    = [a
x | a
x <- HashMap a [a] -> [a]
forall k v. HashMap k v -> [k]
Map.keys (Graph a -> HashMap a [a]
forall a. Graph a -> HashMap a [a]
children Graph a
gr), Bool -> Bool
not (a -> Bool
hasParents a
x)]
    hasParents :: a -> Bool
hasParents a
x = a -> HashMap a [a] -> Bool
forall k a. (Eq k, Hashable k) => k -> HashMap k a -> Bool
Map.member a
x (Graph a -> HashMap a [a]
forall a. Graph a -> HashMap a [a]
parents Graph a
gr)
    -- all nodes in topological order "parents before children"
    list :: [a]
list = Identity [a] -> [a]
forall a. Identity a -> a
runIdentity (Identity [a] -> [a]) -> Identity [a] -> [a]
forall a b. (a -> b) -> a -> b
$ [a] -> GraphM Identity a -> Identity [a]
forall a (m :: * -> *).
(Eq a, Hashable a, Monad m) =>
[a] -> GraphM m a -> m [a]
reversePostOrder' [a]
ancestors ([a] -> Identity [a]
forall a. a -> Identity a
Identity ([a] -> Identity [a]) -> (a -> [a]) -> GraphM Identity a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Graph a -> a -> [a]
forall a. (Eq a, Hashable a) => Graph a -> a -> [a]
getChildren Graph a
gr)

{-----------------------------------------------------------------------------
    Graph traversal
------------------------------------------------------------------------------}
-- | Graph represented as map of immediate children.
type GraphM m a = a -> m [a]

-- | Computes the reverse post-order,
-- listing all transitive children of a node.
-- Each node is listed *before* all its children have been listed.
reversePostOrder :: (Eq a, Hashable a, Monad m) => a -> GraphM m a -> m [a]
reversePostOrder :: a -> GraphM m a -> m [a]
reversePostOrder a
x = [a] -> GraphM m a -> m [a]
forall a (m :: * -> *).
(Eq a, Hashable a, Monad m) =>
[a] -> GraphM m a -> m [a]
reversePostOrder' [a
x]

-- | Reverse post-order from multiple nodes.
-- INVARIANT: For this to be a valid topological order,
-- none of the nodes may have a parent.
reversePostOrder' :: (Eq a, Hashable a, Monad m) => [a] -> GraphM m a -> m [a]
reversePostOrder' :: [a] -> GraphM m a -> m [a]
reversePostOrder' [a]
xs GraphM m a
children = ([a], HashSet a) -> [a]
forall a b. (a, b) -> a
fst (([a], HashSet a) -> [a]) -> m ([a], HashSet a) -> m [a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [a] -> [a] -> HashSet a -> m ([a], HashSet a)
go [a]
xs [] HashSet a
forall a. HashSet a
Set.empty
    where
    go :: [a] -> [a] -> HashSet a -> m ([a], HashSet a)
go []     [a]
rpo HashSet a
visited        = ([a], HashSet a) -> m ([a], HashSet a)
forall (m :: * -> *) a. Monad m => a -> m a
return ([a]
rpo, HashSet a
visited)
    go (a
x:[a]
xs) [a]
rpo HashSet a
visited
        | a
x a -> HashSet a -> Bool
forall a. (Eq a, Hashable a) => a -> HashSet a -> Bool
`Set.member` HashSet a
visited = [a] -> [a] -> HashSet a -> m ([a], HashSet a)
go [a]
xs [a]
rpo HashSet a
visited
        | Bool
otherwise              = do
            [a]
xs' <- GraphM m a
children a
x
            -- visit all children
            ([a]
rpo', HashSet a
visited') <- [a] -> [a] -> HashSet a -> m ([a], HashSet a)
go [a]
xs' [a]
rpo (a -> HashSet a -> HashSet a
forall a. (Eq a, Hashable a) => a -> HashSet a -> HashSet a
Set.insert a
x HashSet a
visited)
            -- prepend this node as all children have been visited
            [a] -> [a] -> HashSet a -> m ([a], HashSet a)
go [a]
xs (a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
rpo') HashSet a
visited'