Generate graphivs graph of nim ast

ast_to_dot.nim

import macros, tables, strformat, strutils, sequtils

import strutils

proc truncate(value: string, length = 9, ellipsis = "..."): string =
  if value.len > length:
    value[0..length - 1] & ellipsis
  else:
    value

proc surround(input: string, surrWith: (string, string)): string =
  surrWith[0] & input & surrWith[1]

proc surround(input: string, surrWith: string): string =
  surrWith & input & surrWith

proc surround1(input: string, surrWith: (string, string)): string =
  if input.len == 0: ""
  else: input.surround(surrWith)

proc escapeHTML(input: string): string =
  input.multiReplace([
    (">", ">"),
    ("<", "&lt;"),
    ("&", "&amp;"),
    ("\"", "&quot;")
  ])

macro dumpDotAst(head: string, body: untyped) =
  var idx = 0
  var descrTable: Table[int, NimNode]
  var graph: seq[string]

  proc toDot(node: NimNode): (int, string) =
    inc idx
    result[0] = idx
    result[1] = "n$# -> " % $idx

    var subnodes: seq[int]
    descrTable[idx] = node

    for sub in node:
      let res = sub.toDot()
      subnodes.add(res[0])

    result[1] &= "{ $# }" % subnodes.mapIt("n" & $it).join(",")
    graph.add(result[1])

  for node in body:
    discard node.toDot()


  var resTotal: seq[string]

  for id, node in descrTable:
    var label: string = $node.kind
    var text: string =
      case node.kind:
        of nnkIdent: node.strVal.escape()
        of nnkStrLit..nnkTripleStrLit:
          $node.strVal.truncate().surround("\"")
        of nnkCharLit .. nnkUInt64Lit: $node.intVal
        else: ""

    var color: string =
      case node.kind:
        of nnkStmtList: "azure2"
        of nnkIdent: "brown2"
        of nnkStrLit: "green"
        of nnkVarSection, nnkLetSection, nnkIdentDefs: "magenta2"
        else: "cyan2"

    var res = "n$# [label = <$#<br/>$#>, fillcolor = $#, style = filled];" % [
      $id,
      label.surround1(("<i>", "</i>")),
      text
        .escapeHTML()
        .surround1(("<b>", "</b>"))
        .surround1(("<font face='courier'>", "</font>")),
      color
    ]

    res = res.replace("\n","")
    resTotal.add(res)


  let filename = head.strVal
  let resultString = """
digraph G {
rankdir = LR;
splines = ortho;
node[shape=box];
$#
$#
}
""" %
    [
      resTotal.join("\n"),
      graph.join("\n")]
  if filename.len == 0:
    let resStrNode = newStrLitNode(resultString)
    result = quote do:
      echo `resStrNode`
  else:
    head.strVal.writeFile(resultString)







# Put any valid code under macro and add name of the file in the macro
# argument. If string is empty result will be printed to stdout. Empty
# does not mean the sting is optional!
dumpDotAst "test.dot":
  result = quote do:
    const helpTable {.inject.} = getHelpTable(`bodyNodeGen`)

    block:
      var
        optParsed {.inject.}: Table[string, CmdArg]
        argParsed {.inject.}: seq[string]
        hasErrors {.inject.}: bool = false
      for arg in body:
        if arg[0] == ident"opt" and arg[1].kind == nnkStmtList:
          ifKeyStmt.add(getOptionParserBranch(arg[1]))

          foundDoubleDash: bool = false

    for kind {.inject.}, key {.inject.}, val {.inject.} in getOpt():
      if key == "" and val == "":
        foundDoubleDash = true
        continue

      if foundDoubleDash:
        let prefix =
          if kind == cmdShortOption: "-"
          elif kind == cmdLongOption: "--"
          else: ""

        argParsed.add(prefix & key & val)
        continue

      # Insert top-level case for argument kind
      case kind
        of cmdShortOption, cmdLongOption:
          `optParserCase`
        of cmdArgument:
          `argParserCase`
        of cmdEnd:
          `endParserCase`

result_example.dot

digraph G {
node[shape=box];
n1 [label = <<i>nnkCall</i><br/>>];
n2 [label = <<i>nnkIdent</i><br/><font face='courier'><b>&quot;shp&quot;</b></font>>];
n3 [label = <<i>nnkStmtList</i><br/>>];
n4 [label = <<i>nnkCommand</i><br/>>];
n5 [label = <<i>nnkIdent</i><br/><font face='courier'><b>&quot;sh&quot;</b></font>>];
n6 [label = <<i>nnkStrLit</i><br/><font face='courier'><b>&quot;ls -al&quot;</b></font>>];
n7 [label = <<i>nnkPrefix</i><br/>>];
n8 [label = <<i>nnkIdent</i><br/><font face='courier'><b>&quot;|&gt;&quot;</b></font>>];
n9 [label = <<i>nnkIdent</i><br/><font face='courier'><b>&quot;duplicate&quot;</b></font>>];
n10 [label = <<i>nnkPrefix</i><br/>>];
n11 [label = <<i>nnkIdent</i><br/><font face='courier'><b>&quot;|&gt;&quot;</b></font>>];
n12 [label = <<i>nnkCommand</i><br/>>];
n13 [label = <<i>nnkIdent</i><br/><font face='courier'><b>&quot;multiplex&quot;</b></font>>];
n14 [label = <<i>nnkPar</i><br/>>];
n15 [label = <<i>nnkPrefix</i><br/>>];
n16 [label = <<i>nnkIdent</i><br/><font face='courier'><b>&quot;&gt;&quot;</b></font>>];
n17 [label = <<i>nnkPar</i><br/>>];
n18 [label = <<i>nnkInfix</i><br/>>];
n19 [label = <<i>nnkIdent</i><br/><font face='courier'><b>&quot;*&quot;</b></font>>];
n20 [label = <<i>nnkIdent</i><br/><font face='courier'><b>&quot;it&quot;</b></font>>];
n21 [label = <<i>nnkIntLit</i><br/><font face='courier'><b>2</b></font>>];
n22 [label = <<i>nnkPrefix</i><br/>>];
n23 [label = <<i>nnkIdent</i><br/><font face='courier'><b>&quot;&gt;&quot;</b></font>>];
n24 [label = <<i>nnkPar</i><br/>>];
n25 [label = <<i>nnkIdent</i><br/><font face='courier'><b>&quot;it&quot;</b></font>>];
n26 [label = <<i>nnkPrefix</i><br/>>];
n27 [label = <<i>nnkIdent</i><br/><font face='courier'><b>&quot;|&gt;&quot;</b></font>>];
n28 [label = <<i>nnkIdent</i><br/><font face='courier'><b>&quot;merge&quot;</b></font>>];
n29 [label = <<i>nnkPrefix</i><br/>>];
n30 [label = <<i>nnkIdent</i><br/><font face='courier'><b>&quot;|&gt;&quot;</b></font>>];
n31 [label = <<i>nnkCommand</i><br/>>];
n32 [label = <<i>nnkIdent</i><br/><font face='courier'><b>&quot;filter&quot;</b></font>>];
n33 [label = <<i>nnkBracket</i><br/>>];
n34 [label = <<i>nnkInfix</i><br/>>];
n35 [label = <<i>nnkIdent</i><br/><font face='courier'><b>&quot;=&gt;&quot;</b></font>>];
n36 [label = <<i>nnkIdent</i><br/><font face='courier'><b>&quot;x&quot;</b></font>>];
n37 [label = <<i>nnkInfix</i><br/>>];
n38 [label = <<i>nnkIdent</i><br/><font face='courier'><b>&quot;==&quot;</b></font>>];
n39 [label = <<i>nnkInfix</i><br/>>];
n40 [label = <<i>nnkIdent</i><br/><font face='courier'><b>&quot;%&quot;</b></font>>];
n41 [label = <<i>nnkIdent</i><br/><font face='courier'><b>&quot;x&quot;</b></font>>];
n42 [label = <<i>nnkIntLit</i><br/><font face='courier'><b>2</b></font>>];
n43 [label = <<i>nnkIntLit</i><br/><font face='courier'><b>0</b></font>>];
n2 -> {  }
n5 -> {  }
n6 -> {  }
n4 -> { n5,n6 }
n8 -> {  }
n9 -> {  }
n7 -> { n8,n9 }
n11 -> {  }
n13 -> {  }
n16 -> {  }
n19 -> {  }
n20 -> {  }
n21 -> {  }
n18 -> { n19,n20,n21 }
n17 -> { n18 }
n15 -> { n16,n17 }
n23 -> {  }
n25 -> {  }
n24 -> { n25 }
n22 -> { n23,n24 }
n14 -> { n15,n22 }
n12 -> { n13,n14 }
n10 -> { n11,n12 }
n27 -> {  }
n28 -> {  }
n26 -> { n27,n28 }
n30 -> {  }
n32 -> {  }
n35 -> {  }
n36 -> {  }
n38 -> {  }
n40 -> {  }
n41 -> {  }
n42 -> {  }
n39 -> { n40,n41,n42 }
n43 -> {  }
n37 -> { n38,n39,n43 }
n34 -> { n35,n36,n37 }
n33 -> { n34 }
n31 -> { n32,n33 }
n29 -> { n30,n31 }
n3 -> { n4,n7,n10,n26,n29 }
n1 -> { n2,n3 }
}

test.png