From 2aaae6a201f79fa0f74f4bf9d8b991ca42e966b0 Mon Sep 17 00:00:00 2001
From: Leorize <leorize+oss@disroot.org>
Date: Fri, 23 Aug 2024 01:43:41 -0500
Subject: [PATCH] cps: implement shims for docgen

Currently CPS simply left the procedures untouched during docgen. This
causes problems when code for `whelp` gets sem-ed, breaking docgen.

This commit implements a couple docgen shims:

- `cps` now result in a `{.cps: baseType.}` pragma tagged to the
  resulting procedure and will show up in the generated documentation.
- `whelp` will generate `nil` continuation for the base type of a cps
  proc.
- `whelp` will generate a `Callback` with a factory type that yields a
  the base continuation type, not the environment type. This should work
  for most cases when docgen are concerned.

This worked well enough for `insideout` docs to be generated
successfully under NimSkull.
---
 cps.nim        |  6 ++--
 cps/docgen.nim | 95 ++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 99 insertions(+), 2 deletions(-)
 create mode 100644 cps/docgen.nim

diff --git a/cps.nim b/cps.nim
index c4893ad6..2ec247a3 100644
--- a/cps.nim
+++ b/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:
diff --git a/cps/docgen.nim b/cps/docgen.nim
new file mode 100644
index 00000000..aa48028b
--- /dev/null
+++ b/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)