Exceptions refactoring and try-finally transform (#171)

* cps/[transform, environment]: use one exception local per try

Nim has a "exception stack" system where inner except branches will
push a new exception to the stack then pop it after it's done.

To replicate this accurately in CPS, we use one local per cps `try`
instead of a single field like how it is done currently. This is
because we set the global exception to whatever the stored exception
was in our continuation leg, so to prevent inner cps `try`s from overriding
what is perceived to be the "current" exception of the outer branch, we
simply give them different locals.

Ideally we can tap right into Nim's exception `up` field, but it's not
exported.

* cps/transform: rewrite `n` instead of doing a roundabout

Not sure what I thought of when I wrote that code, but we can just
run filter on the outer node and don't have to think about the bugs
this `for` would cause us.

* move trampoline and bind to it in bootstrap

* tests/[preamble, ttry]: test exception properties again

* cps/rewrites: add a simplifying rewrite for except T as e

Fixes disruptek/cps#164

* cps/transform: rewrite cps try-except into one continuation

Instead of creating a continuation leg for each except branch, we
merge all except branches into one, then turn that into a continuation.

That way we only have one continuation for all handlers, allow for
a better implementation of cps exceptions in the future.

* cps/transform: try-finally transformation

Implements try-finally transformation for CPS by generating
finally as a continuation leg with a static generic for where
it will continues after.

Untested because the compiler broke.

Known issues:
- Early returns are not handled, which can be fixed by improving
  isScopeExit and making early termination an annotation.

Fixes #80.

* cps/transform: switch the implementation of try-finally to templating

Instead of waiting for nim-lang/Nim#18254
and any other templates bug to be fix. We take the initiative and
write our own continuation templater. It appears to work well enough
to use as an alternative until a better alternative become available.

I've also added an extra test that verify the exception re-raise property.

* cps/spec: remove cpsRecovery

We ended up not needing it for finally

* cps/environment: implement early returns as an annotation

We implement early returns as `cpsTerminate` then tie them up at the
end via `cpsResolver`. This way try-finally can capture early termination and
specialize to those.

Need @disruptek to review this stuff.

* tdefer: enable the test for defer across continuation

* tzevv: enable the defer test

* cps/transform: support except clause with multiple exceptions

* cps/transform: cosmetics for disruptek

Co-authored-by: Andy Davidoff <github@andy.disruptek.com>

cps.nim

@@ -3,8 +3,9 @@ import cps/[spec, transform, rewrites, hooks]
 export Continuation, ContinuationProc
 export cpsCall, cpsMagicCall, cpsVoodooCall, cpsMustJump
 
+# exporting some symbols that we had to bury for bindSym reasons
 from cps/returns import pass
-export pass
+export pass, trampoline
 
 # we only support arc/orc due to its eager expr evaluation qualities
 when not(defined(gcArc) or defined(gcOrc)):
@@ -48,14 +49,6 @@ template dismissed*(c: Continuation): bool =
 
 {.pop.}
 
-proc trampoline*[T: Continuation](c: T): T =
-  ## This is the basic trampoline: it will run the continuation
-  ## until the continuation is no longer in the `Running` state.
-  var c: Continuation = c
-  while c.running:
-    c = c.fn(c)
-  result = T c
-
 macro trampolineIt*[T: Continuation](supplied: T; body: untyped) =
   ## This trampoline allows the user to interact with the continuation
   ## prior to each leg of its execution.  The continuation will be

cps/environment.nim

@@ -30,7 +30,6 @@ type
     c: NimNode                      # the sym we use for the continuation
     fn: NimNode                     # the sym we use for the goto target
     rs: IdentDefs                   # the identdefs for the result
-    ex: NimNode                     # the sym we use for current exception
     mom: NimNode                    # the sym we use for parent continuation
 
   CachePair* = tuple
@@ -189,7 +188,6 @@ proc newEnv*(parent: Env; copy = off): Env =
                  c: parent.c,
                  rs: parent.rs,
                  fn: parent.fn,
-                 ex: parent.ex,
                  parent: parent)
     when cpsReparent:
       result.seen = parent.seen
@@ -384,20 +382,18 @@ proc rewriteReturn*(e: var Env; n: NimNode): NimNode =
       result = newStmtList()
       # ignore the result symbol and create a new assignment
       result.add newAssignment(e.get, n.last.last)
-      # and just issue an empty `return`
-      result.add nnkReturnStmt.newNimNode(n).add newEmptyNode()
-    of nnkEmpty, nnkIdent:
-      # this is `return` or `return continuation`, so that's fine...
-      result = n
+      # and add the termination annotation
+      result.add newCpsTerminate()
+    of nnkEmpty:
+      # this is an empty return
+      result = newCpsTerminate()
     else:
       # okay, it's a return of some rando expr
       result = newStmtList()
       # ignore the result symbol and create a new assignment
       result.add newAssignment(e.get, n.last)
-      # signify the end of the continuation
-      result.add newAssignment(newDotExpr(e.first, e.fn), newNilLit())
-      # and return the continuation
-      result.add nnkReturnStmt.newNimNode(n).add(e.first)
+      # and add the termination annotation
+      result.add newCpsTerminate()
 
 proc rewriteSymbolsIntoEnvDotField*(e: var Env; n: NimNode): NimNode =
   ## swap symbols for those in the continuation
@@ -432,13 +428,15 @@ proc createContinuation*(e: Env; name: NimNode; goto: NimNode): NimNode =
     result.add:
       newAssignment(resultdot e.fn, goto)
 
-proc getException*(e: var Env): NimNode =
-  ## get the current exception from the env, instantiating it if necessary
-  if e.ex.isNil:
-    e.ex = genField"ex"
-    e = e.set e.ex:
-      newIdentDefVar(e.ex, nnkRefTy.newTree(bindSym"Exception"), newNilLit())
-  result = newDotExpr(e.castToChild(e.first), e.ex)
+proc genException*(e: var Env): NimNode =
+  ## generates a new symbol of type ref Exception, then put it in the env.
+  ##
+  ## returns the access to the exception symbol from the env.
+  let ex = genField("ex")
+  e = e.set ex:
+    # XXX: Should be IdentDefLet but saem haven't wrote it yet
+    newIdentDefVar(ex, nnkRefTy.newTree(bindSym"Exception"), newNilLit())
+  result = newDotExpr(e.castToChild(e.first), ex)
 
 proc createWhelp*(env: Env; n: ProcDef, goto: NimNode): ProcDef =
   ## the whelp needs to create a continuation
@@ -489,11 +487,9 @@ proc createBootstrap*(env: Env; n: ProcDef, goto: NimNode): ProcDef =
     result = desym(result, defs[0])
 
   # now the trampoline
+  let tramp = bindSym"trampoline"
   result.body.add:
-    nnkWhileStmt.newTree: [
-      newCall(ident"running", c),  # XXX: bindSym?  bleh.
-      newAssignment(c, env.castToRoot newDotExpr(c, env.fn).newCall(c))
-    ]
+    newAssignment(c, newCall(tramp, c))
 
   # do an easy static check, and then
   if env.rs.typ != result.returnParam:

cps/rewrites.nim

@@ -272,6 +272,40 @@ proc normalizingRewrites*(n: NimNode): NimNode =
       else:
         discard
 
+    proc rewriteExceptBranch(n: NimNode): NimNode =
+      ## Rewrites except branches in the form of `except T as e` into:
+      ##
+      ## ```
+      ## except T:
+      ##   let e = (ref T)(getCurrentException())
+      ## ```
+      ##
+      ## We simplify this AST so that our rewrites can capture it.
+      case n.kind
+      of nnkExceptBranch:
+        # If this except branch has exactly one exception matching clause
+        if n.len == 2:
+          # If the exception matching clause is an infix expression (T as e)
+          if n[0].kind == nnkInfix:
+            let
+              typ = n[0][1] # T in (T as e)
+              refTyp = nnkRefTy.newTree(typ) # make a (ref T) node
+              ex = n[0][2] # our `e`
+              body = n[1]
+
+            result = copyNimNode(n) # copy the `except`
+            result.add typ # add only `typ`
+            result.add:
+              newStmtList:
+                # let ex: ref T = (ref T)(getCurrentException())
+                nnkLetSection.newTree:
+                  newIdentDefs(ex, refTyp, newCall(refTyp, newCall(bindSym"getCurrentException")))
+
+            # add the rewritten body
+            result.last.add:
+              normalizingRewrites body
+      else: discard
+
     case n.kind
     of nnkIdentDefs:
       rewriteIdentDefs n
@@ -287,6 +321,8 @@ proc normalizingRewrites*(n: NimNode): NimNode =
       rewriteFormalParams n
     of CallNodes, nnkHiddenSubConv, nnkHiddenStdConv:
       rewriteHidden n
+    of nnkExceptBranch:
+      rewriteExceptBranch n
     else:
       nil
 

cps/spec.nim

@@ -25,6 +25,7 @@ template cpsContinue*() {.pragma.}      ##
 template cpsCont*() {.pragma.}          ## this is a continuation
 template cpsBootstrap*(whelp: typed) {.pragma.}  ##
 ## the symbol for creating a continuation
+template cpsTerminate*() {.pragma.}     ## this is the end of this procedure
 
 type
   Continuation* = ref object of RootObj
@@ -99,8 +100,9 @@ proc stripPragma*(n: NimNode; s: static[string]): NimNode =
     result = n
 
 proc hash*(n: NimNode): Hash =
+  ## Hash a NimNode via it's representation
   var h: Hash = 0
-  h = h !& hash($n)
+  h = h !& hash(repr n)
   result = !$h
 
 func newCpsPending*(): NimNode =
@@ -136,18 +138,6 @@ proc isCpsContinue*(n: NimNode): bool =
   ## Return whether a node is a {.cpsContinue.} annotation
   n.kind == nnkPragma and n.len == 1 and n.hasPragma("cpsContinue")
 
-when defined(cpsRecover):
-  template cpsRecover() {.pragma.}   ## the next step in finally recovery path
-
-  func newCpsRecover(n: NimNode): NimNode =
-    ## Produce a {.cpsRecover.} annotation
-    nnkPragma.newNimNode(n).add:
-      bindSym"cpsRecover"
-
-  func isCpsRecover(n: NimNode): bool =
-    ## Return whether a node is a {.cpsRecover.} annotation
-    n.kind == nnkPragma and n.len == 1 and n.hasPragma("cpsRecover")
-
 proc breakLabel*(n: NimNode): NimNode =
   ## Return the break label of a `break` statement or a `cpsBreak` annotation
   if n.isCpsBreak():
@@ -172,11 +162,18 @@ proc getContSym*(n: NimNode): NimNode =
   else:
     nil
 
+proc newCpsTerminate*(): NimNode =
+  ## Create a new node signifying early termination of the procedure
+  nnkPragma.newTree:
+    bindSym"cpsTerminate"
+
+proc isCpsTerminate*(n: NimNode): bool =
+  ## Return whether `n` is a cpsTerminate annotation
+  n.kind == nnkPragma and n.len == 1 and n.hasPragma("cpsTerminate")
+
 proc isScopeExit*(n: NimNode): bool =
-  ## Return whether the given node signify a scope exit
-  ##
-  ## TODO: Handle early exit (ie. `c.fn = nil; return`)
-  n.isCpsPending or n.isCpsBreak or n.isCpsContinue
+  ## Return whether the given node signify a CPS scope exit
+  n.isCpsPending or n.isCpsBreak or n.isCpsContinue or n.isCpsTerminate
 
 template rewriteIt*(n: typed; body: untyped): NimNode =
   var it {.inject.} = normalizingRewrites:
@@ -218,3 +215,11 @@ proc isVoodooCall*(n: NimNode): bool =
       let callee = n[0]
       if not callee.isNil and callee.kind == nnkSym:
         result = callee.getImpl.hasPragma "cpsVoodooCall"
+
+proc trampoline*[T: Continuation](c: T): T =
+  ## This is the basic trampoline: it will run the continuation
+  ## until the continuation is no longer in the `Running` state.
+  var c: Continuation = c
+  while not c.isNil and not c.fn.isNil:
+    c = c.fn(c)
+  result = T c