digraph.nim

import
  intsets,
  tables

type
  GraphNode*[T] = ref GraphNodeObj[T]
  GraphNodeObj[T] = object
    val: T
    adj: IntSet
    nodesIn: int
    nodesOut: int
  DiGraphObj[T] = object
    nod*: Table[T, int]
    edg: seq[GraphNode[T]]
    label: string
  DiGraph*[T] = ref DiGraphObj[T]

proc `$`*[T](x: GraphNode[T]): string =
  result = $x.val & '['
  let len = result.len
  for i in x.adj:
    result.addSep(", ", len)
    result.add($i)
  result.add(']')

proc `$`*[T](g: DiGraph[T]): string =
  result = ""
  for u, v in g.edges:
    result.add($u & " -> " & $v & '\l')
  result.setLen(result.len - 1) # wipe last \n

proc newDiGraph*[T](label: string = ""): DiGraph[T] =
  new(result)
  result.label = label
  result.nod = initTable[T, int]()
  result.edg = @[]

{.push inline.}

proc hasNode*[T](g: DiGraph[T], n: T): bool =
  return n in g.nod
proc contains*[T](g: DiGraph[T], n: T): bool =
  return g.hasNode(n)

proc hasEdge*[T](g: DiGraph[T], a, b: T): bool =
  return a in g and b in g and g.nod[b] in g.edg[g.nod[a]].adj
proc contains*[T](g: DiGraph[T], n: tuple[u: T, v: T]): bool =
  return g.hasEdge(n.u, n.v)

proc `label=`*[T](g: DiGraph[T], label: string) =
  g.label = label

proc label*[T](g: DiGraph[T]): string =
  return g.label

proc add*[T](g: DiGraph[T], n: T) = # add node
  if n in g: return
  g.nod[n] = g.edg.len
  g.edg.add(GraphNode[T](val: n))

{.pop.}

proc add*[T](g: DiGraph[T], n: Table[T, T]) = # add edges
  for u, v in n.pairs():
    g.add(u, v)

proc add*[T](g: DiGraph[T], n: Table[T, seq[T]]) = # add edges
  for u, vv in n.pairs():
    for v in vv:
      g.add(u, v)

proc add*[T](g: DiGraph[T], n: varargs[tuple[u: T, v: T]]) = # add edges
  for i in n:
    g.add(i.u, i.v)

proc add*[T](g: DiGraph[T], a, b: T) = # add edge
  g.add(a)
  g.add(b)
  let
    keyA = g.nod[a]
    keyB = g.nod[b]
    node = g.edg[keyA]
  if node.adj.isNil:
    node.adj = initIntSet()
  node.adj.incl(keyB)
  node.nodesOut += 1
  g.edg[keyB].nodesIn += 1


proc del*[T](g: DiGraph[T], n: T) = # delete node
  try:
    let id = g.nod[n]
    g.nod.del(n)
    if not g.edg[id].adj.isNil:
      for k in g.edg[id].adj:
        g.edg[k].nodesIn -= 1
    g.edg[id] = nil
    for v in g.edg:
      if v == nil or v.adj.isNil:
        continue
      v.adj.excl(id)
      v.nodesOut -= 1
  except KeyError:
    return

proc del*[T](g: DiGraph[T], a, b: T) = # delete edge
  if (a, b) notin g:
    return
  var node = g.edg[g.nod[a]]
  node.adj.excl(g.nod[b])
  node.nodesOut -= 1

iterator nodes*[T](g: DiGraph[T]): T {.closure.}=
  for k in g.nod.keys():
    yield k

iterator edges*[T](g: DiGraph[T]): (T, T) =
  for v in g.edg:
    if v == nil or v.adj.isNil:
      continue
    for i in intsets.items(v.adj):
      # if g.edg[i] == nil: continue
      yield (v.val, g.edg[i].val)

iterator outDegree*[T](g: DiGraph[T]): (T, int) =
  for k, v in g.nod.pairs():
    yield (k, g.edg[v].nodesOut)

iterator inDegree*[T](g: DiGraph[T]): (T, int) =
  for k, v in g.nod.pairs():
    yield (k, g.edg[v].nodesIn)

iterator adjancedNodes*[T](g: DiGraph[T], parent: T): T =
  for child in intsets.items(g.edg[parent].adj):
    yield g.edg[child].val

proc relabelNodes*[T](g: DiGraph[T], a: Table[T, T]) =
  for k, v in a.pairs():
    if k in g:
      if v in g:
        let old = g.edg[g.nod[v]]
        let removed = g.edg[g.nod[k]]
        if removed.adj.isNil:
          removed.adj = old.adj
        else:
          for i in intsets.items(old.adj):
            removed.adj.incl(i)
        removed.val = old.val
      g.nod[v] = g.nod[k]
      g.nod.del(k)

template toClosure*(i): auto =
  iterator j: type(i) {.closure.} =
    for x in i:
      yield x
  j

template doDfs(g, nodes, body) =
  var visited = initIntSet()
  for key in nodes:
    let start = g.nod[key]
    if intsets.contains(visited, start): continue
    var stack = @[start]
    while stack.len > 0:
      let parentId = stack.pop()
      if not intsets.contains(visited, parentId):
        let parentObj = g.edg[parentId]
        let parent {.inject.} = parentObj.val
        visited.incl(parentId)
        for childId in intsets.items(parentObj.adj):
          let child {.inject.} = g.edg[childId].val
          body
          stack.add(childId)

iterator dfs*[T](g: DiGraph[T], it: iterator: T): (T, T) {.closure.} =
  doDfs(g, it()):
    yield (parent, child)

# conflicts with varargs[T]
# iterator dfs*[T](g: DiGraph[T]): (T, T) {.closure.} =
  # for u, v in g.dfs(toClosure(g.nodes())):
    # yield (u, v)

# produces messy C code and gcc fails with error
# iterator dfs*[T](g: DiGraph[T], source: varargs[T]): (T, T) {.closure.} =
  # if source.len > 0:
    # doDfs(g, source):
      # yield (parent, child)
  # else:
    # doDfs(g, g.nodes()):
      # yield (parent, child)

# Error: illegal capture 'source'
# iterator dfs*[T](g: DiGraph[T], source: varargs[T]): (T, T) =
  # if source.len > 0:
    # for parent, child in g.dfs(toClosure(source.items())):
      # yield (parent, child)
  # else:
    # for parent, child in g.dfs(toClosure(g.nodes())):
      # yield (parent, child)

# works, but this inline iterator is huge
iterator dfs*[T](g: DiGraph[T], source: varargs[T]): (T, T) =
  if source.len > 0:
    doDfs(g, source):
      yield (parent, child)
  else:
    doDfs(g, g.nodes()):
      yield (parent, child)

iterator dfs*[T](g: DiGraph[T], source: T): (T, T) {.closure.} =
  doDfs(g, [source]):
    yield (parent, child)

when isMainModule:
  import strutils, sequtils
  var g = newDiGraph[int]()
  g.add({
    1: @[2, 3, 4],
    2: @[3, 5],
    6: @[5, 7, 8, 9]
  }.toTable)
  doAssert((6, 5) in g == true)
  doAssert((1, 7) in g == false)
  doAssert(1 in g == true)
  doAssert toSeq(g.dfs([1,2])) == @[(1, 2), (1, 3), (1, 4), (2, 3), (2, 5)]
  doAssert toSeq(g.dfs()) == @[(1, 2), (1, 3), (1, 4), (2, 3), (2, 5), (6, 5), (6, 7), (6, 8), (6, 9)]
  
  var s = "@["
  let len = s.len
  for u, v in g.dfs():
    s.addSep(", ", len)
    s.add("($1, $2)".format(u, v), )
  s.add("]")
  echo s