From 88063c0da9bd7d8e48197c074ac7420ea833fb6d Mon Sep 17 00:00:00 2001
From: Leorize <leorize+oss@disroot.org>
Date: Sat, 10 Feb 2024 13:42:58 -0600
Subject: [PATCH] experimental cps detection via scope injection

This employs the use of a "phantom" scope to inject new symbols into the
body of a cps block before they are semantically checked.

After that, we surgically extract the transformed output, lifting the
result from the phantom scope to the invocation scope.
---
 cps.nim               | 28 +++++++++++++++++++---------
 cps/normalizedast.nim |  2 +-
 cps/spec.nim          |  1 +
 cps/transform.nim     |  7 +++++++
 4 files changed, 28 insertions(+), 10 deletions(-)

diff --git a/cps.nim b/cps.nim
index 657a6e62..24cd9957 100644
--- a/cps.nim
+++ b/cps.nim
@@ -76,19 +76,29 @@ macro cpsTyped(tipe: typed, n: typed): untyped =
     else:
       result = getAst(cpsTransform(tipe, n))
 
-macro cps*(tipe: typed, n: untyped): untyped =
+macro cpsMarkExtract(n: untyped): untyped =
+  let cpsExtract = bindSym"cpsExtract"
+  result = nnkPragmaBlock.newTree(
+    nnkPragma.newTree(cpsExtract),
+    n
+  )
+
+template 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.
-  result = n
+  bind cpsTyped
+  bind cpsScalpel
+  bind cpsMarkExtract
+
   when not defined(nimdoc):
-    # add the application of the typed transformation pass
-    n.addPragma:
-      nnkExprColonExpr.newTree(bindSym"cpsTyped", tipe)
-    # let the untyped pass do what it will with this input
-    # XXX: currently disabled because it's a slipperly slope of regret
-    #result = performUntypedPass(T, n)
-    result = n
+    cpsScalpel:
+      const isInCps {.inject.} = true
+
+      cpsMarkExtract:
+        cpsTyped(tipe, n)
+  else:
+    n
 
 proc adaptArguments(sym: NormNode; args: seq[NormNode]): seq[NormNode] =
   ## convert any arguments in the list as necessary to match those of
diff --git a/cps/normalizedast.nim b/cps/normalizedast.nim
index f5abf1b2..f510797c 100644
--- a/cps/normalizedast.nim
+++ b/cps/normalizedast.nim
@@ -267,7 +267,7 @@ defineToNimNodeConverter(
 allowAutoDowngradeNormalizedNode(
     Name, TypeExpr, Call, Conv, PragmaStmt, PragmaAtom, IdentDef, RoutineDef,
     ProcDef, FormalParams, RoutineParam, VarSection, LetSection, VarLet,
-    VarLetIdentDef, VarLetTuple, DefVarLet, IdentDefLet, Sym
+    VarLetIdentDef, VarLetTuple, DefVarLet, IdentDefLet, Sym, PragmaBlock
   )
 
 # types that go from a specific type to a less specific type, "downgrade"
diff --git a/cps/spec.nim b/cps/spec.nim
index d44dd411..2db99420 100644
--- a/cps/spec.nim
+++ b/cps/spec.nim
@@ -50,6 +50,7 @@ template cpsTerminate*() {.pragma.}     ## this is the end of this procedure
 template cpsHasException*(cont, ex: typed) {.pragma.}  ##
 ## the continuation has an exception stored in `ex`, with `cont` being the
 ## continuation symbol used.
+template cpsExtract*() {.pragma.} ## extract this block from the scope
 
 const
   cpsStackFrames* {.booldefine, used.} = compileOption"stacktrace"
diff --git a/cps/transform.nim b/cps/transform.nim
index e833e5c2..a6e5d7d0 100644
--- a/cps/transform.nim
+++ b/cps/transform.nim
@@ -975,6 +975,13 @@ macro cpsFloater(n: typed): untyped =
   result = floater:
     copyNimTree n
 
+macro cpsScalpel*(n: typed): untyped =
+  ## returns the first `{.cpsExtract.}` block from `n`
+  debugAnnotation cpsScalpel, n:
+    it = it.findChildRecursive() do (n: NormNode) -> bool:
+      n.kind == nnkPragmaBlock and n.asPragmaBlock().hasPragma("cpsExtract")
+    it = it.last
+
 macro cpsManageException(name: static[string]; n: typed): untyped =
   ## rewrites all continuations in `n` containing an exception so that exception
   ## become the "current" exception of that continuation while preserving the