
   0 | module Util.Graph
   1 | 
   2 | import Data.SortedMap
   3 | import Data.List1
   4 | 
   5 | import Control.Barbie
   6 | import Derive.Prelude
   7 | import Derive.Barbie
   8 | import Monocle
   9 | 
  10 | %language ElabReflection
  11 | %default total
  12 | 
  13 | ||| An opaque ID type for nodes in a graph
  14 | export
  15 | data NodeId : Type where
  16 |   MkNodeId : (id : Nat) -> NodeId
  17 | %name NodeId id, id_2, id_3
  18 | %runElab derive "NodeId" [Eq, Ord]
  19 | 
  20 | export
  21 | Show NodeId where show (MkNodeId id) = show id
  22 | 
  23 | inner : NodeId -> Nat
  24 | inner (MkNodeId id) = id
  25 | 
  26 | next : List NodeId -> NodeId
  27 | next [] = MkNodeId 0
  28 | next (x :: xs) =
  29 |   let ids = map inner $ x ::: xs
  30 |       max = foldl1 max ids
  31 |   in MkNodeId (S max)
  32 | 
  33 | public export
  34 | data Field = NodeWeights | EdgeWeights
  35 | 
  36 | public export
  37 | Tpe : (n, e : Type) -> Field -> Type
  38 | Tpe n e NodeWeights = n
  39 | Tpe n e EdgeWeights = e
  40 | 
  41 | %prefix_record_projections off
  42 | 
  43 | export
  44 | record Graph (f : Field -> Type) where
  45 |   constructor MkGraph
  46 |   nodes : SortedMap NodeId (f NodeWeights)
  47 |   edges : SortedMap (NodeId, NodeId) (f EdgeWeights)
  48 | 
  49 | FunctorB Field Graph where
  50 |   bmap_ f (MkGraph nodes edges) =
  51 |     let nodes = map (\x => (f NodeWeights) x) nodes
  52 |         edges = map (\x => (f EdgeWeights) x) edges
  53 |     in MkGraph nodes edges
  54 | 
  55 | public export
  56 | IGraph : (n, e : Type) -> Type
  57 | IGraph n e = Graph (Tpe n e)
  58 | 
  59 | export
  60 | Show (f NodeWeights) => Show (f EdgeWeights) => Show (Graph f) where
  61 |   show (MkGraph nodes edges) =
  62 |     let (node_size, edge_size) = (length . keys  $ nodes, length . keys $ edges)
  63 |         header = "Graph with \{show node_size} nodes and \{show edge_size} edges"
  64 |         nodes =
  65 |           joinBy ", "
  66 |           . map (\(x,y) => "\{show x}: \{show y}")
  67 |           . Data.SortedMap.toList
  68 |           $ nodes
  69 |         edges =
  70 |           joinBy ", "
  71 |           . map (\(x,y) => "\{show x}: \{show y}")
  72 |           . Data.SortedMap.toList
  73 |           $ edges
  74 |     in "\{header}\n  Nodes ->\n    \{nodes}\n  Edges ->\n    \{edges}"
  75 | 
  76 | namespace Graph
  77 |   ||| Make an empty graph
  78 |   export
  79 |   empty : Graph f
  80 |   empty = MkGraph empty empty
  81 | 
  82 |   ||| Insert a node, returning the new NodeId
  83 |   export
  84 |   insert_node :
  85 |     (weight : f NodeWeights) -> (graph : Graph f)
  86 |     -> (NodeId, Graph f)
  87 |   insert_node weight (MkGraph nodes edges) =
  88 |     let node_id = next . keys $ nodes
  89 |         nodes = insert node_id weight nodes
  90 |     in (node_id, MkGraph nodes edges)
  91 | 
  92 |   ||| Insert an edge
  93 |   export
  94 |   insert_edge :
  95 |     (edge : (NodeId, NodeId)) -> (weight : f EdgeWeights) -> (graph : Graph f)
  96 |     -> Graph f
  97 |   insert_edge edge weight (MkGraph nodes edges) =
  98 |     let edges = insert edge weight edges
  99 |     in MkGraph nodes edges
 100 | 
 101 |   ||| Get the weight of a node, if it exists
 102 |   export
 103 |   node_weight :
 104 |     (node_id : NodeId) -> (graph : Graph f)
 105 |     -> Maybe (f NodeWeights)
 106 |   node_weight node_id (MkGraph nodes edges) = lookup node_id nodes
 107 | 
 108 |   ||| Get the weight of an edge, if it exists
 109 |   export
 110 |   edge_weight :
 111 |     (pair : (NodeId, NodeId)) -> (graph : Graph f)
 112 |     -> Maybe (f EdgeWeights)
 113 |   edge_weight pair (MkGraph nodes edges) = lookup pair edges
 114 | 
 115 |   ||| Returns the NodeId if the graph contains a given node weight
 116 |   export
 117 |   contains_node_weight :
 118 |     Eq (f NodeWeights) =>
 119 |     (weight : f NodeWeights) -> (graph : Graph f)
 120 |     -> Maybe NodeId
 121 |   contains_node_weight weight (MkGraph nodes edges) =
 122 |     map fst . head' . filter ((== weight) . snd) . toList $ nodes
 123 | 
 124 |   ||| Returns the NodeId pair if the graph contains a given edge weight
 125 |   export
 126 |   contains_edge_weight :
 127 |     Eq (f EdgeWeights) =>
 128 |     (weight : f EdgeWeights) -> (graph : Graph f)
 129 |     -> Maybe (NodeId, NodeId)
 130 |   contains_edge_weight weight (MkGraph nodes edges) =
 131 |     map fst . head' . filter ((== weight) . snd) . toList $ edges
 132 | 
 133 |   ||| List of the Ids of all the nodes in the graph
 134 |   export
 135 |   nodes : (graph : Graph f) -> List NodeId
 136 |   nodes (MkGraph ns edges) = keys ns
 137 | 
 138 |   ||| List of all the edges in the graph
 139 |   export
 140 |   edges : (graph : Graph f) -> List (NodeId, NodeId)
 141 |   edges (MkGraph nodes es) = keys es
 142 | 
 143 |   ||| All the edges that point away from a specific node
 144 |   export
 145 |   outbound_edges : (node : NodeId) -> (graph : Graph f) -> Maybe (List (NodeId, NodeId))
 146 |   outbound_edges node (MkGraph nodes edges) = do
 147 |     _ <- lookup node nodes
 148 |     pure . filter ((== node) . fst) . keys $ edges
 149 | 
 150 |   ||| All the edges that point to a specific node
 151 |   export
 152 |   inbound_edges : (node : NodeId) -> (graph : Graph f) -> Maybe (List (NodeId, NodeId))
 153 |   inbound_edges node (MkGraph nodes edges) = do
 154 |     _ <- lookup node nodes
 155 |     pure . filter ((== node) . snd) . keys $ edges
 156 | 
 157 |   ||| Convert a graph to one that holds visited booleans in the weights
 158 |   export
 159 |   visited : Graph (Tpe n e) -> Graph (const Bool)
 160 |   visited graph = bmap (const False) graph
 161 |