From 709fe9b52d8d13e21aa93bdc7d882f071c01e48d Mon Sep 17 00:00:00 2001
From: Leorize <leorize+oss@disroot.org>
Date: Wed, 26 May 2021 15:24:49 -0500
Subject: [PATCH] cps/xfrm: experiment at inlining next body

We are inlining if there is only one continuing control
---
 cps/xfrm.nim | 47 +++++++++++++++++++++++++++++++++++++++--------
 1 file changed, 39 insertions(+), 8 deletions(-)

diff --git a/cps/xfrm.nim b/cps/xfrm.nim
index ccf28e61..b243ff25 100644
--- a/cps/xfrm.nim
+++ b/cps/xfrm.nim
@@ -205,23 +205,54 @@ macro cpsMayJump(cont, n, after: typed): untyped =
   var n = normalizingRewrites:
     # we always wrap the input because there's no reason not to
     newStmtList n
+  var after = normalizingRewrites:
+    # we always wrap the input because there's no reason not to
+    newStmtList after
 
-  let
-    afterProc = makeContProc(genSym(nskProc, "done"), cont, after)
-    afterTail = tailCall(desym cont, afterProc.name)
+  proc countToTwo(n: NimNode): int =
+    ## Count the amount of cpsPending in `n`, stopping as soon
+    ## as two or more cpsPending is found
+    if n.isCpsPending:
+      inc result
+    else:
+      for child in n.items:
+        result += child.countToTwo
+        if result >= 2:
+          return
+
+  # whether a jump is needed to leave `n`
+  let jumpCount = n.countToTwo + int(n.firstReturn.isNil)
+
+  result = newStmtList()
+  var resolvedBody: NimNode
 
-    # FIXME: fix this to use makeReturn
+  if jumpCount > 1:
+    let
+      afterProc = makeContProc(genSym(nskProc, "done"), cont, after)
+      afterTail = tailCall(desym cont, afterProc.name)
+
+      # FIXME: fix this to use makeReturn
 
     resolvedBody =
       n.replace(isCpsPending):
         afterTail
 
-  if resolvedBody.firstReturn.isNil:
-    resolvedBody.add afterTail
+    if resolvedBody.firstReturn.isNil:
+      resolvedBody.add afterTail
+
+    result.add afterProc
+  else:
+    # If no jump is needed, we inline instead
+    let nextBody = after.resym(cont, desym cont)
+    if nextBody.firstReturn.isNil:
+      nextBody.add newCpsPending()
+
+    resolvedBody =
+      n.replace(isCpsPending):
+        nextBody
 
-  result = newStmtList()
-  result.add afterProc
   result.add resolvedBody
+
   result = workaroundRewrites result
 
   #debug("cpsMayJump", result, Transformed, n)