From 050f39c2157390b96e479b416a03c84f4e7875b0 Mon Sep 17 00:00:00 2001
From: Leorize <leorize+oss@disroot.org>
Date: Sun, 2 May 2021 22:47:21 -0500
Subject: [PATCH] cps: experimental support for for-loops

For-loops are inlined using block statements, which messes with their
control flow due to #76.

Breaks "for loop with continue, break" test. On the flip side,
splitting now functions on them.

Also breaks a test in tzevv due to defer being rewritten into
try-finally, which doesn't work.
---
 cps.nim         | 28 ++++++++++++++++++++++++++--
 tests/taste.nim | 12 ++++++++++++
 2 files changed, 38 insertions(+), 2 deletions(-)

diff --git a/cps.nim b/cps.nim
index 08a4313b..1c66c321 100644
--- a/cps.nim
+++ b/cps.nim
@@ -603,7 +603,12 @@ proc normalizingRewrites(n: NimNode): NimNode =
       if n.len == 2:
         n.add newEmptyNode()
       elif n[1].isEmpty:          # add explicit type symbol
-        n[1] = getTypeInst n[2]
+        if not n[2].isEmpty:
+          n[1] = getTypeInst n[2]
+        else:
+          # get the type from the symbol as the last resort.
+          # walkaround for #48.
+          n[1] = getTypeInst n[0]
       result = n
 
   proc rewriteVarLet(n: NimNode): NimNode =
@@ -637,6 +642,15 @@ proc normalizingRewrites(n: NimNode): NimNode =
       result = n[0]
     else: discard
 
+  proc rewriteFastAsgn(n: NimNode): NimNode =
+    ## Rewrite nnkFastAsgn into nnkAsgn because sem don't like them
+    case n.kind
+    of nnkFastAsgn:
+      result = newNimNode(nnkAsgn, n)
+      for child in n.items:
+        result.add child
+    else: discard
+
   case n.kind
   of nnkIdentDefs:
     rewriteIdentDefs n
@@ -644,6 +658,8 @@ proc normalizingRewrites(n: NimNode): NimNode =
     rewriteVarLet n
   of nnkHiddenAddr, nnkHiddenDeref:
     rewriteHiddenAddrDeref n
+  of nnkFastAsgn:
+    rewriteFastAsgn n
   else:
     nil
 
@@ -854,13 +870,21 @@ proc cpsXfrm(T: NimNode, n: NimNode): NimNode =
     result = copy n
   result = workaroundRewrites(result)
 
-macro cps*(T: typed, n: typed): untyped =
+macro cps2(T: typed, n: typed): untyped =
   # I hate doing stuff inside macros, call the proc to do the work
   when defined(nimdoc):
     result = n
   else:
     result = cpsXfrm(T, n)
 
+macro cps*(T: typed, n: typed): untyped =
+  case n.kind
+  of nnkProcDef:
+    result = getImplTransformed n.name
+    result.addPragma newColonExpr(bindSym"cps2", T)
+  else:
+    result = cpsXfrm(T, n)
+
 macro cpsMagic*(n: untyped): untyped =
   ## upgrade cps primitives to generate errors out of context
   ## and take continuations as input inside {.cps.} blocks
diff --git a/tests/taste.nim b/tests/taste.nim
index 2b989626..72f2f6e0 100644
--- a/tests/taste.nim
+++ b/tests/taste.nim
@@ -656,3 +656,15 @@ suite "tasteful tests":
       check i == 4
 
     trampoline foo()
+
+  block:
+    ## for-loop with splitting
+    r = 0
+    proc foo() {.cps: Cont.} =
+      inc r
+      for i in 0 ..< 3:
+        noop()
+        inc r
+
+    trampoline foo()
+    check r == 4