cps: implement shims for docgen

cps.nim

@@ -1,5 +1,5 @@
 import std/[genasts, deques]
-import cps/[spec, transform, rewrites, hooks, exprs, normalizedast, callbacks]
+import cps/[docgen, spec, transform, rewrites, hooks, exprs, normalizedast, callbacks]
 import std/macros except newStmtList, newTree
 
 export Continuation, ContinuationProc, ContinuationFn, State, Callback
@@ -80,7 +80,8 @@ macro cps*(tipe: typed, n: untyped): untyped =
   ## When applied to a procedure, rewrites the procedure into a
   ## continuation form. When applied to a procedure type definition,
   ## rewrites the type into a callback form.
-  when defined(nimdoc): return n
+  when defined(nimdoc): return cpsDocAnnotate(tipe, n)
+
   # add the application of the typed transformation pass
   n.addPragma:
     nnkExprColonExpr.newTree(bindSym"cpsTyped", tipe)
@@ -139,6 +140,7 @@ macro whelp*(call: typed): untyped =
   ## If you pass `whelp` a continuation procedure *symbol* instead, the
   ## result is a `Callback` which you can use to create many individual
   ## continuations or recover the `result` of an extant continuation.
+  when defined(nimdoc): return docWhelp(call)
   result =
     case call.kind
     of nnkSym:

cps/docgen.nim

@@ -0,0 +1,95 @@
+when defined(nimdoc):
+  import std/macros except newTree, newStmtList
+  import callbacks, normalizedast, spec
+
+  template cps(tipe: typed) {.pragma.}
+
+  proc cpsDocAnnotate*(tipe, n: NimNode): NimNode =
+    ## Annotate `n` with a `cps` pragma that will show up in doc
+    n.addPragma:
+      nnkExprColonExpr.newTree(bindSym"cps", tipe)
+    result = n
+
+  proc getCpsBase(n: NormNode, origin: NormNode = n): NormNode =
+    ## recover the symbol of the cps base type or generate error ast
+    case n.kind
+    of NormalCallNodes:
+      getCpsBase(n[0], origin)
+    of nnkSym:
+      getCpsBase(n.getImpl, origin)
+    of nnkProcDef:
+      if n.hasPragma "borrow":
+        getCpsBase(n.last)
+      elif n.hasPragma "cps":
+        pragmaArgument(n, "cps")
+      else:
+        error "procedure " & n.name.strVal & " doesn't seem to be a cps call", origin
+        normalizedast.newCall("typeOf", n)
+    else:
+      error "procedure doesn't seem to be a cps call", origin
+      ## XXX: darn ambiguous calls
+      normalizedast.newCall("typeOf", n)
+
+  proc whelpCall(call: Call): NormNode =
+    let base = getCpsBase(call)
+    result = newStmtList()
+    for param in call[1..^1]:
+      # Generate a let statement for every param.
+      # This makes sure that effects generated by those params will be recorded
+      # by docgen.
+      result.add:
+        nnkLetSection.newTree(
+          nnkIdentDefs.newTree(
+            nnkPragmaExpr.newTree(
+              nskLet.genSym"forEffectsOnly",
+              nnkPragma.newTree(ident"used")
+            ),
+            newEmptyNode(),
+            param
+          )
+        )
+
+    result.add:
+      newCall(base, newNilLit()) # Add the "supposed" whelp result
+
+  proc whelpCallback(sym: Sym): NormNode =
+    let
+      impl = sym.getImpl.ProcDef
+      base = getCpsBase(impl)
+
+    var params = nnkFormalParams.newTree(copy base)
+    for param in impl.callingParams:
+      var idefs = nnkIdentDefs.newTree()
+
+      # Add names
+      for names in param[0..^3]:
+        idefs.add desym(names.Sym)
+
+      # Add type
+      if param[^2].kind == nnkEmpty:
+        idefs.add getTypeInst(param[^1])
+      else:
+        idefs.add param[^2]
+
+      # Make sure no "default" params are here
+      idefs.add newEmptyNode()
+      params.add idefs
+
+    let factory = nnkProcTy.newTree(params, nnkPragma.newTree ident"nimcall")
+    let returnType = NormNode: copyOrVoid(NimNode impl.returnParam)
+
+    result = newCall(
+      nnkBracketExpr.newTree(
+        bindSym"Callback",
+        copy base,
+        returnType,
+        factory
+      )
+    )
+
+  proc docWhelp*(call: NimNode): NimNode =
+    ## Fake `whelp` implementation for docgen
+    if call.kind == nnkSym:
+      whelpCallback(call.Sym)
+    else:
+      whelpCall(normalizeCall call)