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

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

module Reactive.Banana.Prim.Graph
  ( Graph
  , emptyGraph
  , insertEdge
  , getChildren
  , listParents
  , dfs
  ) where

import           Control.Monad
import           Data.Functor.Identity
import qualified Data.HashMap.Strict   as Map
import qualified Data.HashSet          as Set
import           Data.Hashable
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] -> ([a] -> [a]) -> Maybe [a] -> [a]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe [] [a] -> [a]
forall a. a -> a
id (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] -> ([a] -> [a]) -> Maybe [a] -> [a]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe [] [a] -> [a]
forall a. a -> a
id (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 :: [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]
dfs' [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 successors.
type GraphM m a = a -> m [a]

-- | Depth-first search. List all transitive successors of a node.
-- A node is listed *before* all its successors have been listed.
dfs :: (Eq a, Hashable a, Monad m) => a -> GraphM m a -> m [a]
dfs :: a -> GraphM m a -> m [a]
dfs a
x = [a] -> GraphM m a -> m [a]
forall a (m :: * -> *).
(Eq a, Hashable a, Monad m) =>
[a] -> GraphM m a -> m [a]
dfs' [a
x]

-- | Depth-first serach, refined version.
-- INVARIANT: None of the nodes in the initial list have a predecessor.
dfs' :: forall a m. (Eq a, Hashable a, Monad m) => [a] -> GraphM m a -> m [a]
dfs' :: [a] -> GraphM m a -> m [a]
dfs' [a]
xs GraphM m a
succs = (([a], HashSet a) -> [a]) -> m ([a], HashSet a) -> m [a]
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM ([a], HashSet a) -> [a]
forall a b. (a, b) -> a
fst (m ([a], HashSet a) -> m [a]) -> m ([a], HashSet a) -> m [a]
forall a b. (a -> b) -> a -> b
$ [a] -> [a] -> HashSet a -> m ([a], HashSet a)
go [a]
xs [] HashSet a
forall a. HashSet a
Set.empty
    where
    go :: [a] -> [a] -> Set.HashSet a -> m ([a], Set.HashSet a)
    go :: [a] -> [a] -> HashSet a -> m ([a], HashSet a)
go []     [a]
ys HashSet a
seen            = ([a], HashSet a) -> m ([a], HashSet a)
forall (m :: * -> *) a. Monad m => a -> m a
return ([a]
ys, HashSet a
seen)    -- all nodes seen
    go (a
x:[a]
xs) [a]
ys HashSet a
seen
        | a
x a -> HashSet a -> Bool
forall a. (Eq a, Hashable a) => a -> HashSet a -> Bool
`Set.member` HashSet a
seen    = [a] -> [a] -> HashSet a -> m ([a], HashSet a)
go [a]
xs [a]
ys HashSet a
seen
        | Bool
otherwise              = do
            [a]
xs' <- GraphM m a
succs a
x
            -- visit all children
            ([a]
ys', HashSet a
seen') <- [a] -> [a] -> HashSet a -> m ([a], HashSet a)
go [a]
xs' [a]
ys (a -> HashSet a -> HashSet a
forall a. (Eq a, Hashable a) => a -> HashSet a -> HashSet a
Set.insert a
x HashSet a
seen)
            -- list this node as all successors have been seen
            [a] -> [a] -> HashSet a -> m ([a], HashSet a)
go [a]
xs (a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
ys') HashSet a
seen'