From d1a6e5a21dd900526c373e7430d4a4e2be7657fd Mon Sep 17 00:00:00 2001
From: Timothee Cour <timothee.cour2@gmail.com>
Date: Wed, 5 May 2021 15:15:48 -0700
Subject: [PATCH] add genPNode to allow writing PNode-based compiler code
 similarly to `genAst`

---
 compiler/passes.nim | 21 ++++++-----
 compiler/vmconv.nim | 90 ++++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 99 insertions(+), 12 deletions(-)

diff --git a/compiler/passes.nim b/compiler/passes.nim
index 6376d2581c1bc..347268419d46a 100644
--- a/compiler/passes.nim
+++ b/compiler/passes.nim
@@ -14,10 +14,7 @@ import
   options, ast, llstream, msgs,
   idents,
   syntaxes, modulegraphs, reorder,
-  lineinfos, pathutils
-from parser import parseString
-from sugar import dup
-from strutils import `%`
+  lineinfos, pathutils, vmconv
 from os import splitFile
 
 type
@@ -92,14 +89,18 @@ proc processImplicits(graph: ModuleGraph; implicits: seq[string], nodeKind: TNod
   for module in items(implicits):
     # implicit imports should not lead to a module importing itself
     if m.position != resolveMod(graph.config, module, relativeTo).int32:
-      var code = ""
-      let quoted = "".dup(addQuoted(module))
-      let name = module.splitFile.name
+      var c = GenContext(cache: graph.cache, info: m.info)
+      var node: PNode
       case nodeKind
-      of nkImportStmt: code = "import $1 as $2\n{.used: $2.}" % [quoted, name]
-      of nkIncludeStmt: code = "include $1" % quoted
+      of nkImportStmt:
+        let name = graph.cache.getIdent(module.splitFile.name)
+        node = genPNode(c, module, name):
+          import module as name
+          {.used: name.}
+      of nkIncludeStmt:
+        node = genPNode(c, module):
+          include module
       else: assert false
-      let node = parseString(code, graph.cache, graph.config, filename = "", line = 0, errorHandler = nil)
       if not processTopLevelStmt(graph, node, a): break
 
 const
diff --git a/compiler/vmconv.nim b/compiler/vmconv.nim
index b82fb2ff3199d..2914198d5057f 100644
--- a/compiler/vmconv.nim
+++ b/compiler/vmconv.nim
@@ -1,4 +1,8 @@
+from std/strutils import cmpIgnoreStyle
+import std/macros
 import ast
+from idents import PIdent, IdentCache
+from lineinfos import TLineInfo
 
 template elementType*(T: typedesc): typedesc =
   typeof(block:
@@ -19,7 +23,12 @@ proc fromLit*(a: PNode, T: typedesc): auto =
 proc toLit*[T](a: T): PNode =
   ## generic type => PNode
   ## see also reverse operation `fromLit`
-  when T is string: newStrNode(nkStrLit, a)
+  # xxx also allow an optional `info` param
+  when false: discard
+  elif T is string: newStrNode(nkStrLit, a)
+  elif T is PIdent:
+    result = newNode(nkIdent)
+    result.ident = a
   elif T is Ordinal: newIntNode(nkIntLit, a.ord)
   elif T is (proc): newNode(nkNilLit)
   elif T is ref:
@@ -28,7 +37,7 @@ proc toLit*[T](a: T): PNode =
   elif T is tuple:
     result = newTree(nkTupleConstr)
     for ai in fields(a): result.add toLit(ai)
-  elif T is seq:
+  elif T is seq | array:
     result = newNode(nkBracket)
     for ai in a:
       result.add toLit(ai)
@@ -43,3 +52,80 @@ proc toLit*[T](a: T): PNode =
   else:
     static: doAssert false, "not yet supported: " & $T # add as needed
 
+type
+  GenContext* = object
+    cache*: IdentCache
+    info*: TLineInfo
+  ContextVars = seq[tuple[name: string, val: NimNode]]
+
+proc genPNodeImpl(c: NimNode, code: var NimNode, vals: ContextVars, n: NimNode): NimNode =
+  if n.kind == nnkIdent:
+    for v in vals:
+      if n.strVal.cmpIgnoreStyle(v.name) == 0:
+        return v.val
+  result = genSym(nskVar, "ret")
+  let kind2 = n.kind.ord.TNodeKind.newLit
+  code.add quote do:
+    # alternatively, get `info` from `n.info`, which requires exposing this in macros.
+    var `result` = newNodeI(`kind2`, `c`.info)
+  # keep in sync with `ast.TNode`
+  case n.kind
+  of nnkCharLit..nnkUInt64Lit:
+    let val = n.intVal.newLit
+    code.add quote do:
+      `result`.intVal = `val`
+  of nnkFloatLit..nnkFloat128Lit:
+    let val = n.floatVal.newLit
+    code.add quote do:
+      `result`.floatVal = `val`
+    let tmp = quote do:
+      `result`.floatVal = `val`
+  of nnkStrLit..nnkTripleStrLit:
+    let val = n.strVal.newLit
+    code.add quote do:
+      `result`.strVal = `val`
+  of nnkIdent:
+    let val = n.strVal.newLit
+    code.add quote do:
+      `result`.ident = getIdent(`c`.cache, `val`)
+  of nnkSym: doAssert false # not implemented, but shouldn't be needed
+  else:
+    for ni in n:
+      let reti = genPNodeImpl(c, code, vals, ni)
+      code.add quote do:
+        `result`.add `reti`
+
+macro genPNode*(c: GenContext, args: varargs[untyped]): PNode =
+  ## Converts an AST into a PNode, and works similarly to std/genasts.
+  ## This can simplify writing compiler code, avoiding to manually write `PNode` ASTs.
+  runnableExamples:
+    import idents, renderer
+    let cache = newIdentCache()
+    let a = [1,2]
+    let b = cache.getIdent("foo")
+    var c = GenContext(cache: cache)
+    let node = genPNode(c, a, b):
+      for i in 0..<3:
+        let b = @[i, 1]
+        echo (a, b, i, "abc")
+    let s = node.renderTree
+    assert s == """
+
+for i in 0 ..< 3:
+  let foo = @[i, 1]
+  echo ([1, 2], foo, i, "abc")""", s
+
+  result = newStmtList()
+  let m = args.len - 1
+  var vals: ContextVars
+  vals.setLen m
+  for i in 0..<m:
+    let ai = args[i]
+    assert ai.kind == nnkIdent, $ai.kind
+    vals[i].name = ai.repr
+    let ni = genSym(nskVar, vals[i].name)
+    vals[i].val = ni
+    result.add quote do:
+      var `ni` = toLit(`ai`)
+  let ret = genPNodeImpl(c, result, vals, args[^1])
+  result.add ret