Immutable DAG

DAG.scala

val graph = DAG.empty[String]
val graph2 = graph.addVertex("Foo")
val graph3 = graph2.addEdge("Foo", "Bar")
val graph4 = graph3.addEdge("Baz", "Quux")
val graph5 = graph4.removeEdge("Baz", "Quux")
val graph6 = graph5.addEdge("Foo", "Baz")
val graph7 = graph6.addEdge("Bar" -> "Quux").addEdge("Baz" -> "Quux")
graph: DAG[String] =


scala> graph2: DAG[String] =

 ┌───┐
 │Foo│
 └───┘

scala> graph3: DAG[String] =

 ┌───┐
 │Foo│
 └─┬─┘
   │
   v
 ┌───┐
 │Bar│
 └───┘

scala> graph4: DAG[String] =

 ┌───┐ ┌───┐
 │Foo│ │Baz│
 └─┬─┘ └──┬┘
   │      │
   v      v
 ┌───┐ ┌────┐
 │Bar│ │Quux│
 └───┘ └────┘

scala> graph5: DAG[String] =

       ┌───┐
       │Foo│
       └─┬─┘
         │
   ┌─────┘
   │
   v
 ┌───┐ ┌───┐ ┌────┐
 │Bar│ │Baz│ │Quux│
 └───┘ └───┘ └────┘

scala> graph6: DAG[String] =

      ┌─────┐
      │ Foo │
      └─┬┬──┘
        ││
   ┌────┘│
   │     │
   v     v
 ┌───┐ ┌───┐ ┌────┐
 │Bar│ │Baz│ │Quux│
 └───┘ └───┘ └────┘

scala> graph7: DAG[String] =

  ┌─────┐
  │ Foo │
  └┬──┬─┘
   │  │
   │  └──┐
   │     │
   v     v
 ┌───┐ ┌───┐
 │Bar│ │Baz│
 └──┬┘ └─┬─┘
    │    │
    │ ┌──┘
    │ │
    v v
  ┌─────┐
  │Quux │
  └─────┘

scala> 

import com.github.mdr.ascii.graph.Graph
import com.github.mdr.ascii.layout.GraphLayout

object DAG {

  def empty[V] = DAG[V](vertices = Set(), outEdgeMap = SetMultimap.empty, inEdgeMap = SetMultimap.empty)

}

object SetMultimap {

  def empty[K, V] = SetMultimap[K, V](Map())

}

case class SetMultimap[K, V](rawMap: Map[K, Set[V]] = Map()) {

  def add(k: K, v: V) = SetMultimap(rawMap + (k -> (this(k) + v)))

  def apply(k: K) = rawMap.getOrElse(k, Set())

  def delete(k: K) = SetMultimap(rawMap - k)

  def remove(k: K, v: V) = SetMultimap(rawMap + (k -> (this(k) - v)))

}

case class DAG[V](vertices: Set[V], outEdgeMap: SetMultimap[V, V], inEdgeMap: SetMultimap[V, V])  {

  def addVertex(v: V): DAG[V] = copy(vertices = vertices + v)

  def addEdge(e: (V, V)): DAG[V] = addEdge(e._1, e._2)
  
  def addEdge(v1: V, v2: V): DAG[V] = {
    if (descendents(v2).contains(v1))
      throw new IllegalArgumentException("Cannot add edge between $v1 and $v2 as this would create a cycle")
    addVertex(v1).addVertex(v2).copy(outEdgeMap = outEdgeMap.add(v1, v2), inEdgeMap = inEdgeMap.add(v2, v1))
  }

  def removeEdge(v1: V, v2: V): DAG[V] =
    copy(outEdgeMap = outEdgeMap.remove(v1, v2), inEdgeMap = outEdgeMap.remove(v2, v1))

  def removeVertex(v: V): DAG[V] = {
    val newOutEdgeMap = getOutVertices(v).foldLeft(outEdgeMap) { case (map, v2) ⇒ map.remove(v, v2) }
    val newInEdgeMap = getInVertices(v).foldLeft(inEdgeMap) { case (map, v2) ⇒ map.remove(v, v2) }
    copy(vertices - v, inEdgeMap = newInEdgeMap, outEdgeMap = newOutEdgeMap)
  }

  def descendents(v: V): Set[V] =  getOutVertices(v).flatMap(descendents) + v

  def ancestors(v: V): Set[V] = getInVertices(v).flatMap(ancestors) + v

  def getOutVertices(v: V): Set[V] = outEdgeMap(v)

  def getInVertices(v: V): Set[V] = inEdgeMap(v)

  override def toString = {
    val edges = for {
      (v1, vs) ← outEdgeMap.rawMap.toList
      v2 ← vs
    } yield (v1, v2)
    val graph = Graph(vertices, edges)
    "\n" + GraphLayout.renderGraph(graph)
  }

}