Graph.hs

-- | A structure representing unique vertices and their interrelations
module Mini.Data.Graph (
  -- * Type
  Graph,
  graph,

  -- * Algorithms
  bfs,
  distance,
  path,
  reachable,
  sort,

  -- * Construction
  fromList,
  singleton,

  -- * Modification
  add,
  connect,
  disconnect,
  filter,
  remove,
  transpose,

  -- * Partition
  partition,
  split,

  -- * Query
  edges,
  incoming,
  indegree,
  indegrees,
  member,
  outdegree,
  outdegrees,
  outgoing,
  sinkMax,
  sinkMin,
  sinks,
  sourceMax,
  sourceMin,
  sources,
  vertices,
) where

import Data.Foldable (
  fold,
 )
import Data.List (
  unfoldr,
 )
import Mini.Data.Map (
  Map,
 )
import qualified Mini.Data.Map as Map (
  delete,
  filter,
  foldlWith,
  foldrWith,
  insertWith,
  lookup,
  member,
  partition,
  singleton,
  split,
  unionWith,
 )
import Mini.Data.Recursion (
  bool,
  maybe,
  uncurry,
 )
import Mini.Data.Set (
  Set,
 )
import qualified Mini.Data.Set as Set (
  delete,
  difference,
  filter,
  member,
  singleton,
 )
import Mini.Hash.Class (
  Hashable,
  toBytes,
 )
import Prelude (
  Bool,
  Eq,
  Foldable,
  Int,
  Maybe (
    Just,
    Nothing
  ),
  Monoid,
  Ord,
  Semigroup,
  Show,
  any,
  compare,
  concatMap,
  const,
  flip,
  foldMap,
  foldr,
  length,
  mempty,
  null,
  show,
  ($),
  (+),
  (.),
  (<$>),
  (<>),
  (==),
 )

-- Type

-- | A graph with directed edges between vertices of type /a/
data Graph a
  = Graph
      (Map a (Set a))
      (Map a (Set a))

instance (Eq a) => Eq (Graph a) where
  (Graph _ oes1) == (Graph _ oes2) = oes1 == oes2

instance (Ord a) => Ord (Graph a) where
  compare (Graph _ oes1) (Graph _ oes2) = compare oes1 oes2

instance (Show a) => Show (Graph a) where
  show = show . edges

instance Foldable Graph where
  foldr f b = foldr f b . vertices

instance (Ord a) => Semigroup (Graph a) where
  (Graph ies1 oes1) <> (Graph ies2 oes2) =
    Graph
      (Map.unionWith (const (<>)) ies1 ies2)
      (Map.unionWith (const (<>)) oes1 oes2)

instance (Ord a) => Monoid (Graph a) where
  mempty = Graph mempty mempty

instance (Hashable a) => Hashable (Graph a) where
  toBytes = concatMap toBytes

-- | Primitive recursion on graphs (internally represented by adjacency maps)
graph
  :: (Map a (Set a) -> Map a (Set a) -> b)
  -- ^ Function applied to the adjacency maps of the graph:
  -- incoming edges, outgoing edges
  -> Graph a
  -- ^ The graph
  -> b
graph f (Graph ies oes) = f ies oes

-- Algorithms

-- | Breadth-first search for the hierarchy in a graph from a starting vertex
bfs :: (Ord a) => Graph a -> a -> [Set a]
bfs (Graph _ oes) s =
  foldMap
    ( \vs ->
        Set.singleton s
          : unfoldr
            ( \((us, es), ds) ->
                bool
                  ( Just
                      ( us
                      ,
                        ( foldr
                            ( \u b@(us', es') ->
                                maybe
                                  b
                                  ( \vs' ->
                                      ( (us' <> vs') `Set.difference` ds
                                      , Map.delete u es'
                                      )
                                  )
                                  $ Map.lookup u es
                            )
                            (mempty, es)
                            us
                        , ds <> us
                        )
                      )
                  )
                  Nothing
                  $ null us
            )
            ((vs, oes), Set.singleton s)
    )
    $ Map.lookup s oes

-- | Get the shortest distance in a graph between a vertex and another
distance :: (Ord a) => Graph a -> a -> a -> Maybe Int
distance g s t =
  foldr
    (\a b -> bool ((+ 1) <$> b) (Just 0) $ t `Set.member` a)
    Nothing
    $ bfs g s

-- | Check whether there is a path in a graph from a vertex to another
path :: (Ord a) => Graph a -> a -> a -> Bool
path g s t = any (t `Set.member`) $ bfs g s

-- | Get the reachable vertices in a graph from a starting vertex
reachable :: (Ord a) => Graph a -> a -> Set a
reachable g = fold . bfs g

-- | Topologically sort a graph (assumes acyclicity)
sort :: (Ord a) => Graph a -> [a]
sort = unfoldr $ \g -> (\u -> (u, remove u g)) <$> sourceMin g

-- Construction

-- | Make a graph from a list of edges
fromList :: (Ord a) => [(a, a)] -> Graph a
fromList = foldr (uncurry connect) mempty

-- | Make a graph with an isolated vertex
singleton :: (Ord a) => a -> Graph a
singleton u = Graph (Map.singleton u mempty) (Map.singleton u mempty)

-- Modification

-- | Add an isolated vertex to a graph unless already present
add :: (Ord a) => a -> Graph a -> Graph a
add u (Graph ies oes) =
  Graph
    (Map.insertWith (const (<>)) u mempty ies)
    (Map.insertWith (const (<>)) u mempty oes)

-- | Add edges from a vertex to another vertex in a graph
connect :: (Ord a) => a -> a -> Graph a -> Graph a
connect u v (Graph ies oes) =
  Graph
    ( Map.insertWith (const (<>)) v (Set.singleton u) $
        Map.insertWith (const (<>)) u mempty ies
    )
    ( Map.insertWith (const (<>)) u (Set.singleton v) $
        Map.insertWith (const (<>)) v mempty oes
    )

-- | Remove an edge from a vertex to another vertex in a graph
disconnect :: (Ord a) => a -> a -> Graph a -> Graph a
disconnect u v (Graph ies oes) =
  Graph
    (Map.insertWith (const $ flip Set.difference) v (Set.singleton u) ies)
    (Map.insertWith (const $ flip Set.difference) u (Set.singleton v) oes)

-- | Keep the vertices that satisfy a predicate in a graph
filter :: (Ord a) => (a -> Bool) -> Graph a -> Graph a
filter p (Graph ies oes) =
  Graph
    (Set.filter p <$> Map.filter (const . p) ies)
    (Set.filter p <$> Map.filter (const . p) oes)

-- | Remove a vertex and its edges from a graph
remove :: (Ord a) => a -> Graph a -> Graph a
remove u (Graph ies oes) =
  Graph
    (Set.delete u <$> Map.delete u ies)
    (Set.delete u <$> Map.delete u oes)

-- | Reverse the edges of a graph
transpose :: Graph a -> Graph a
transpose (Graph ies oes) = Graph oes ies

-- Partition

-- | Partition a graph with a predicate into @(true, false)@ subgraphs
partition :: (Ord a) => (a -> Bool) -> Graph a -> (Graph a, Graph a)
partition p (Graph ies oes) =
  let (ies1, ies2) = Map.partition (const . p) ies
      (oes1, oes2) = Map.partition (const . p) oes
   in (Graph ies1 oes1, Graph ies2 oes2)

-- | Split a graph by a vertex into @(lt, eq, gt)@ subgraphs
split :: (Ord a) => a -> Graph a -> (Graph a, Maybe a, Graph a)
split v (Graph ies oes) =
  let (lti, _, gti) = Map.split v ies
      (lto, _, gto) = Map.split v oes
   in (Graph lti lto, bool Nothing (Just v) $ v `Map.member` oes, Graph gti gto)

-- Query

-- | Get the edges of a graph
edges :: Graph a -> [(a, a)]
edges (Graph _ oes) =
  Map.foldrWith
    (\u -> flip $ foldr (\v -> (:) (u, v)))
    []
    oes

-- | Get the incoming edges of each vertex in a graph
incoming :: Graph a -> Map a (Set a)
incoming (Graph ies _) = ies

-- | Get the number of incoming edges to a vertex in a graph
indegree :: (Ord a) => a -> Graph a -> Maybe Int
indegree v (Graph ies _) = length <$> Map.lookup v ies

-- | Get the number of incoming edges of each vertex in a graph
indegrees :: Graph a -> Map a Int
indegrees (Graph ies _) = length <$> ies

-- | Check whether a vertex is in a graph
member :: (Ord a) => a -> Graph a -> Bool
member u (Graph _ oes) = u `Map.member` oes

-- | Get the number of outgoing edges from a vertex in a graph
outdegree :: (Ord a) => a -> Graph a -> Maybe Int
outdegree u (Graph _ oes) = length <$> Map.lookup u oes

-- | Get the number of outgoing edges of each vertex in a graph
outdegrees :: Graph a -> Map a Int
outdegrees (Graph _ oes) = length <$> oes

-- | Get the outgoing edges of each vertex in a graph
outgoing :: Graph a -> Map a (Set a)
outgoing (Graph _ oes) = oes

-- | Get the maximum vertex with no outgoing edges from a graph
sinkMax :: Graph a -> Maybe a
sinkMax (Graph _ oes) =
  Map.foldlWith
    ( \b k ->
        bool
          b
          (Just k)
          . null
    )
    Nothing
    oes

-- | Get the minimum vertex with no outgoing edges from a graph
sinkMin :: Graph a -> Maybe a
sinkMin (Graph _ oes) =
  Map.foldrWith
    ( \k a b ->
        bool
          b
          (Just k)
          $ null a
    )
    Nothing
    oes

-- | Get the vertices with no outgoing edges from a graph
sinks :: Graph a -> [a]
sinks (Graph _ oes) =
  Map.foldrWith
    ( \k a b ->
        bool
          b
          (k : b)
          $ null a
    )
    []
    oes

-- | Get the maximum vertex with no incoming edges from a graph
sourceMax :: Graph a -> Maybe a
sourceMax (Graph ies _) =
  Map.foldlWith
    ( \b k ->
        bool
          b
          (Just k)
          . null
    )
    Nothing
    ies

-- | Get the minimum vertex with no incoming edges from a graph
sourceMin :: Graph a -> Maybe a
sourceMin (Graph ies _) =
  Map.foldrWith
    ( \k a b ->
        bool
          b
          (Just k)
          $ null a
    )
    Nothing
    ies

-- | Get the vertices with no incoming edges from a graph
sources :: Graph a -> [a]
sources (Graph ies _) =
  Map.foldrWith
    ( \k a b ->
        bool
          b
          (k : b)
          $ null a
    )
    []
    ies

-- | Get the vertices of a graph
vertices :: Graph a -> [a]
vertices (Graph _ oes) = Map.foldrWith (\k _ -> (:) k) [] oes