From d71f69ab50f079c03860f244f6c64b555ca403b6 Mon Sep 17 00:00:00 2001
From: Yuriy Glukhov <yuriy.glukhov@gmail.com>
Date: Wed, 25 Apr 2018 19:38:59 +0300
Subject: [PATCH 01/14] Closure iter transformation

---
 compiler/ccgexprs.nim                |   2 +-
 compiler/ccgstmts.nim                |  10 +-
 compiler/closureiters.nim            | 632 +++++++++++++++++++++++++++
 compiler/lambdalifting.nim           |  60 ++-
 compiler/options.nim                 |   3 +-
 compiler/renderer.nim                |  14 +-
 compiler/seminst.nim                 |   2 +-
 compiler/semstmts.nim                |   6 +-
 compiler/transf.nim                  |  25 +-
 tests/async/tasync_in_seq_constr.nim |   3 +-
 10 files changed, 713 insertions(+), 44 deletions(-)
 create mode 100644 compiler/closureiters.nim

diff --git a/compiler/ccgexprs.nim b/compiler/ccgexprs.nim
index 96f9265f17..5b3f6c3d21 100644
--- a/compiler/ccgexprs.nim
+++ b/compiler/ccgexprs.nim
@@ -2326,7 +2326,7 @@ proc expr(p: BProc, n: PNode, d: var TLoc) =
   of nkParForStmt: genParForStmt(p, n)
   of nkState: genState(p, n)
   of nkGotoState: genGotoState(p, n)
-  of nkBreakState: genBreakState(p, n)
+  of nkBreakState: genBreakState(p, n, d)
   else: internalError(n.info, "expr(" & $n.kind & "); unknown node kind")
 
 proc genNamedConstExpr(p: BProc, n: PNode): Rope =
diff --git a/compiler/ccgstmts.nim b/compiler/ccgstmts.nim
index cb3d6dbe6e..96f5b53a77 100644
--- a/compiler/ccgstmts.nim
+++ b/compiler/ccgstmts.nim
@@ -177,17 +177,17 @@ proc genGotoState(p: BProc, n: PNode) =
     lineF(p, cpsStmts, "case $2: goto $1$2;$n", [prefix, rope(i)])
   lineF(p, cpsStmts, "}$n", [])
 
-proc genBreakState(p: BProc, n: PNode) =
+proc genBreakState(p: BProc, n: PNode, d: var TLoc) =
   var a: TLoc
+  initLoc(d, locExpr, n, OnUnknown)
+
   if n.sons[0].kind == nkClosure:
-    # XXX this produces quite inefficient code!
     initLocExpr(p, n.sons[0].sons[1], a)
-    lineF(p, cpsStmts, "if (((NI*) $1)[1] < 0) break;$n", [rdLoc(a)])
+    d.r = "(((NI*) $1)[1] < 0)" % [rdLoc(a)]
   else:
     initLocExpr(p, n.sons[0], a)
     # the environment is guaranteed to contain the 'state' field at offset 1:
-    lineF(p, cpsStmts, "if ((((NI*) $1.ClE_0)[1]) < 0) break;$n", [rdLoc(a)])
-  #  lineF(p, cpsStmts, "if (($1) < 0) break;$n", [rdLoc(a)])
+    d.r = "((((NI*) $1.ClE_0)[1]) < 0)" % [rdLoc(a)]
 
 proc genGotoVar(p: BProc; value: PNode) =
   if value.kind notin {nkCharLit..nkUInt64Lit}:
diff --git a/compiler/closureiters.nim b/compiler/closureiters.nim
new file mode 100644
index 0000000000..02795ab470
--- /dev/null
+++ b/compiler/closureiters.nim
@@ -0,0 +1,632 @@
+#
+#
+#           The Nim Compiler
+#        (c) Copyright 2018 Andreas Rumpf
+#
+#    See the file "copying.txt", included in this
+#    distribution, for details about the copyright.
+#
+
+# This file implements closure iterator transformations.
+# The main idea is to split the closure iterator body to top level statements.
+# The body is split by yield statement.
+#
+# Example:
+#  while a > 0:
+#    echo "hi"
+#    yield a
+#    dec a
+#
+# Should be transformed to:
+#  STATE0:
+#    if a > 0:
+#      echo "hi"
+#      :state = 1 # Next state
+#      return a # yield
+#    else:
+#      :state = 2 # Next state
+#      break :stateLoop # Proceed to the next state
+#  STATE1:
+#    dec a
+#    :state = 0 # Next state
+#    break :stateLoop # Proceed to the next state
+#  STATE2:
+#    :state = -1 # End of execution
+
+# The transformation should play well with lambdalifting, however depending
+# on situation, it can be called either before or after lambdalifting
+# transformation. As such we behave slightly differently, when accessing
+# iterator state, or using temp variables. If lambdalifting did not happen,
+# we just create local variables, so that they will be lifted further on.
+# Otherwise, we utilize existing env, created by lambdalifting.
+
+# Lambdalifting treats :state variable specially, it should always end up
+# as the first field in env. Currently C codegen depends on this behavior.
+
+# One special subtransformation is nkStmtListExpr lowering.
+# Example:
+#   template foo(): int =
+#     yield 1
+#     2
+#
+#   iterator it(): int {.closure.} =
+#     if foo() == 2:
+#       yield 3
+#
+# If a nkStmtListExpr has yield inside, it has first to be lowered to:
+#   yield 1
+#   :tmpSlLower = 2
+#   if :tmpSlLower == 2:
+#     yield 3
+
+
+import
+  intsets, strutils, options, ast, astalgo, trees, treetab, msgs, os, options,
+  idents, renderer, types, magicsys, rodread, lowerings, tables, sequtils,
+  lambdalifting
+
+type ClosureIteratorTransformationContext = object
+  fn: PSym
+  stateVarSym: PSym # :state variable. nil if env already introduced by lambdalifting
+  states: seq[PNode] # The resulting states. Every state is an nkState node.
+  blockLevel: int # Temp used to transform break and continue stmts
+  stateLoopLabel: PSym # Label to break on, when jumping between states.
+  exitStateIdx: int # index of the last state
+  tempVarId: int # unique name counter
+  tempVars: PNode # Temp var decls, nkVarSection
+  loweredStmtListExpr: PNode # Temporary used for nkStmtListExpr lowering
+
+proc newStateAssgn(ctx: var ClosureIteratorTransformationContext, stateNo: int = -2): PNode =
+  # Creates state assignmen:
+  #   :state = stateNo
+
+  result = newNode(nkAsgn)
+  if ctx.stateVarSym.isNil:
+    let state = getStateField(ctx.fn)
+    assert state != nil
+    result.add(rawIndirectAccess(newSymNode(getEnvParam(ctx.fn)),
+                      state, result.info))
+  else:
+    result.add(newSymNode(ctx.stateVarSym))
+  result.add(newIntTypeNode(nkIntLit, stateNo, getSysType(tyInt)))
+
+proc setStateInAssgn(stateAssgn: PNode, stateNo: int) =
+  assert stateAssgn.kind == nkAsgn
+  assert stateAssgn[1].kind == nkIntLit
+  stateAssgn[1].intVal = stateNo
+
+proc newState(ctx: var ClosureIteratorTransformationContext, n, gotoOut: PNode): int =
+  # Creates a new state, adds it to the context fills out `gotoOut` so that it
+  # will goto this state.
+  # Returns index of the newly created state
+
+  result = ctx.states.len
+  let resLit = newIntLit(result)
+  let s = newNodeI(nkState, n.info)
+  s.add(resLit)
+  s.add(n)
+  ctx.states.add(s)
+  if not gotoOut.isNil:
+    assert(gotoOut.len == 0)
+    gotoOut.add(newIntLit(result))
+
+proc toStmtList(n: PNode): PNode =
+  result = n
+  if result.kind notin {nkStmtList, nkStmtListExpr}:
+    result = newNodeI(nkStmtList, n.info)
+    result.add(n)
+
+proc addGotoOut(n: PNode, gotoOut: PNode): PNode =
+  # Make sure `n` is a stmtlist, and ends with `gotoOut`
+
+  result = toStmtList(n)
+  if result.len != 0 and result.sons[^1].kind != nkGotoState:
+    result.add(gotoOut)
+
+proc newTempVarAccess(ctx: var ClosureIteratorTransformationContext, typ: PType, i: TLineInfo): PNode =
+  if not ctx.stateVarSym.isNil:
+    # We haven't gone through labmda lifting yet, so just create a local var,
+    # it will be lifted later
+    let s = newSym(skVar, getIdent(":tmpSlLower" & $ctx.tempVarId), ctx.fn, i)
+    s.typ = typ
+
+    if ctx.tempVars.isNil:
+      ctx.tempVars = newNode(nkVarSection)
+      addVar(ctx.tempVars, newSymNode(s))
+
+    result = newSymNode(s)
+  else:
+    # Lambda lifting is done, insert temp var to env.
+    let s = newSym(skVar, getIdent(":tmpSlLower" & $ctx.tempVarId), ctx.fn, i)
+    s.typ = typ
+    result = freshVarForClosureIter(s, ctx.fn)
+
+  inc ctx.tempVarId
+
+proc hasYields(n: PNode): bool =
+  # TODO: This is very inefficient. It traverses the node, looking for nkYieldStmt.
+  case n.kind
+  of nkYieldStmt:
+    result = true
+  of nkCharLit..nkUInt64Lit, nkFloatLit..nkFloat128Lit, nkStrLit..nkTripleStrLit,
+      nkSym, nkIdent, procDefs, nkTemplateDef:
+    discard
+  else:
+    for c in n:
+      if c.hasYields:
+        result = true
+        break
+
+proc transformBreaksAndContinuesInWhile(ctx: var ClosureIteratorTransformationContext, n: PNode, before, after: PNode): PNode =
+  result = n
+  case n.kind
+  of nkCharLit..nkUInt64Lit, nkFloatLit..nkFloat128Lit, nkStrLit..nkTripleStrLit,
+      nkSym, nkIdent, procDefs, nkTemplateDef:
+    discard
+  of nkWhileStmt: discard # Do not recurse into nested whiles
+  of nkContinueStmt:
+    result = before
+  of nkBlockStmt:
+    inc ctx.blockLevel
+    result[1] = ctx.transformBreaksAndContinuesInWhile(result[1], before, after)
+    dec ctx.blockLevel
+  of nkBreakStmt:
+    if ctx.blockLevel == 0:
+      result = after
+  else:
+    for i in 0 ..< n.len:
+      n[i] = ctx.transformBreaksAndContinuesInWhile(n[i], before, after)
+
+proc transformBreaksInBlock(ctx: var ClosureIteratorTransformationContext, n: PNode, label, after: PNode): PNode =
+  result = n
+  case n.kind
+  of nkCharLit..nkUInt64Lit, nkFloatLit..nkFloat128Lit, nkStrLit..nkTripleStrLit,
+      nkSym, nkIdent, procDefs, nkTemplateDef:
+    discard
+  of nkBlockStmt, nkWhileStmt:
+    inc ctx.blockLevel
+    result[1] = ctx.transformBreaksInBlock(result[1], label, after)
+    dec ctx.blockLevel
+  of nkBreakStmt:
+    if n[0].kind == nkEmpty:
+      if ctx.blockLevel == 0:
+        result = after
+    else:
+      if label.kind == nkSym and n[0].sym == label.sym:
+        result = after
+  else:
+    for i in 0 ..< n.len:
+      n[i] = ctx.transformBreaksInBlock(n[i], label, after)
+
+proc collectExceptState(n: PNode): PNode =
+  var ifStmt = newNode(nkIfStmt)
+  for c in n:
+    if c.kind == nkExceptBranch:
+      var ifBranch: PNode
+      var branchBody: PNode
+
+      if c[0].kind == nkType:
+        assert(c.len == 2)
+        ifBranch = newNode(nkElifBranch)
+        let expression = newNodeI(nkCall, n.info)
+        expression.add(callCodegenProc("getCurrentException", emptyNode))
+        expression.add(c[0])
+        ifBranch.add(expression)
+        branchBody = c[1]
+      else:
+        assert(c.len == 1)
+        if ifStmt.len == 0:
+          ifStmt = newNode(nkStmtList)
+          ifBranch = newNode(nkStmtList)
+        else:
+          ifBranch = newNode(nkElse)
+        branchBody = c[0]
+
+      ifBranch.add(branchBody)
+      ifStmt.add(ifBranch)
+
+  if ifStmt.len != 0:
+    result = newNode(nkStmtList)
+    result.add(ifStmt)
+  else:
+    result = emptyNode
+
+proc getFinallyNode(n: PNode): PNode =
+  result = n[^1]
+  if result.kind == nkFinally:
+    result = result[0]
+  else:
+    result = emptyNode
+
+proc hasYieldsInExpressions(n: PNode): bool =
+  case n.kind
+  of nkCharLit..nkUInt64Lit, nkFloatLit..nkFloat128Lit, nkStrLit..nkTripleStrLit,
+      nkSym, nkIdent, procDefs, nkTemplateDef:
+    discard
+  of nkStmtListExpr:
+    result = n.hasYields
+  of nkStmtList, nkWhileStmt, nkCaseStmt, nkIfStmt:
+    discard
+  else:
+    for c in n:
+      if c.hasYieldsInExpressions:
+        return true
+
+proc lowerStmtListExpr(ctx: var ClosureIteratorTransformationContext, n: PNode): PNode =
+  result = n
+  case n.kind
+  of nkCharLit..nkUInt64Lit, nkFloatLit..nkFloat128Lit, nkStrLit..nkTripleStrLit,
+      nkSym, nkIdent, procDefs, nkTemplateDef:
+    discard
+  of nkStmtListExpr:
+    if n.hasYields:
+      for i in 0 .. n.len - 2:
+        ctx.loweredStmtListExpr.add(n[i])
+
+      let tv = ctx.newTempVarAccess(n.typ, n[^1].info)
+      let asgn = newNode(nkAsgn)
+      asgn.add(tv)
+      asgn.add(n[^1])
+      ctx.loweredStmtListExpr.add(asgn)
+      result = tv
+
+  else:
+    for i in 0 ..< n.len:
+      n[i] = ctx.lowerStmtListExpr(n[i])
+
+proc transformClosureIteratorBody(ctx: var ClosureIteratorTransformationContext, n: PNode, gotoOut: PNode): PNode =
+  result = n
+  case n.kind:
+    of nkCharLit..nkUInt64Lit, nkFloatLit..nkFloat128Lit, nkStrLit..nkTripleStrLit,
+        nkSym, nkIdent, procDefs, nkTemplateDef:
+      discard
+
+    of nkStmtList:
+      result = addGotoOut(result, gotoOut)
+      for i in 0 ..< n.len:
+        if n[i].hasYieldsInExpressions:
+          # Lower nkStmtListExpr nodes inside `n[i]` first
+          assert(ctx.loweredStmtListExpr.isNil)
+          ctx.loweredStmtListExpr = newNodeI(nkStmtList, n.info)
+          n[i] = ctx.lowerStmtListExpr(n[i])
+          ctx.loweredStmtListExpr.add(n[i])
+          n[i] = ctx.loweredStmtListExpr
+          ctx.loweredStmtListExpr = nil
+
+        if n[i].hasYields:
+          # Create a new split
+          let go = newNode(nkGotoState)
+          n[i] = ctx.transformClosureIteratorBody(n[i], go)
+
+          let s = newNode(nkStmtList)
+          for j in i + 1 ..< n.len:
+            s.add(n[j])
+
+          n.sons.setLen(i + 1)
+          discard ctx.newState(s, go)
+          discard ctx.transformClosureIteratorBody(s, gotoOut)
+          break
+
+    of nkStmtListExpr:
+      assert(false, "nkStmtListExpr not lowered")
+
+    of nkYieldStmt:
+      # echo "YIELD!"
+      result = newNodeI(nkStmtList, n.info)
+      result.add(n)
+      result.add(gotoOut)
+
+    of nkElse, nkElseExpr:
+      result[0] = addGotoOut(result[0], gotoOut)
+      result[0] = ctx.transformClosureIteratorBody(result[0], gotoOut)
+
+    of nkElifBranch, nkElifExpr, nkOfBranch:
+      result[1] = addGotoOut(result[1], gotoOut)
+      result[1] = ctx.transformClosureIteratorBody(result[1], gotoOut)
+
+    of nkIfStmt, nkCaseStmt:
+      for i in 0 ..< n.len:
+        n[i] = ctx.transformClosureIteratorBody(n[i], gotoOut)
+      if n[^1].kind != nkElse:
+        # We don't have an else branch, but every possible branch has to end with
+        # gotoOut, so add else here.
+        let elseBranch = newNode(nkElse)
+        elseBranch.add(gotoOut)
+        n.add(elseBranch)
+
+    of nkWhileStmt:
+      # while e:
+      #   s
+      # ->
+      # BEGIN_STATE:
+      #   if e:
+      #     s
+      #     goto BEGIN_STATE
+      #   else:
+      #     goto OUT
+
+      result = newNodeI(nkGotoState, n.info)
+
+      let s = newNodeI(nkStmtList, n.info)
+      discard ctx.newState(s, result)
+      let ifNode = newNodeI(nkIfStmt, n.info)
+      let elifBranch = newNodeI(nkElifBranch, n.info)
+      elifBranch.add(n[0])
+
+      var body = addGotoOut(n[1], result)
+
+      body = ctx.transformBreaksAndContinuesInWhile(body, result, gotoOut)
+      body = ctx.transformClosureIteratorBody(body, result)
+
+      elifBranch.add(body)
+      ifNode.add(elifBranch)
+
+      let elseBranch = newNode(nkElse)
+      elseBranch.add(gotoOut)
+      ifNode.add(elseBranch)
+      s.add(ifNode)
+
+    of nkBlockStmt:
+      result[1] = addGotoOut(result[1], gotoOut)
+      result[1] = ctx.transformBreaksInBlock(result[1], result[0], gotoOut)
+      result[1] = ctx.transformClosureIteratorBody(result[1], gotoOut)
+
+    of nkTryStmt:
+      var tryBody = toStmtList(n[0])
+
+      # let popTry = newNode(nkPar)
+      # popTry.add(newIdentNode(getIdent("popTry"), n.info))
+      var finallyBody = newNode(nkStmtList)
+      # finallyBody.add(popTry)
+      finallyBody.add(getFinallyNode(n))
+
+      var tryCatchOut = newNode(nkGotoState)
+
+      tryBody = ctx.transformClosureIteratorBody(tryBody, tryCatchOut)
+      var exceptBody = collectExceptState(n)
+
+      var exceptIdx = -1
+      if exceptBody.kind != nkEmpty:
+        exceptBody = ctx.transformClosureIteratorBody(exceptBody, tryCatchOut)
+        exceptIdx = ctx.newState(exceptBody, nil)
+
+      finallyBody = ctx.transformClosureIteratorBody(finallyBody, gotoOut)
+      let finallyIdx = ctx.newState(finallyBody, tryCatchOut)
+
+      # let pushTry = newNode(nkPar) #newCall(newSym("pushTry"), newIntLit(exceptIdx))
+      # pushTry.add(newIdentNode(getIdent("pushTry"), n.info))
+      # pushTry.add(newIntLit(exceptIdx))
+      # pushTry.add(newIntLit(finallyIdx))
+      # tryBody.sons.insert(pushTry, 0)
+
+      result = tryBody
+
+    of nkGotoState, nkForStmt:
+      internalError("closure iter " & $n.kind)
+
+    else:
+      for i in 0 ..< n.len:
+        n[i] = ctx.transformClosureIteratorBody(n[i], gotoOut)
+
+proc stateFromGotoState(n: PNode): int =
+  assert(n.kind == nkGotoState)
+  result = n[0].intVal.int
+
+proc tranformStateAssignments(ctx: var ClosureIteratorTransformationContext, n: PNode): PNode =
+  # This transforms 3 patterns:
+  ########################## 1
+  # yield e
+  # goto STATE
+  # ->
+  # :state = STATE
+  # return e
+  ########################## 2
+  # goto STATE
+  # ->
+  # :state = STATE
+  # break :stateLoop
+  ########################## 3
+  # return e
+  # ->
+  # :state = -1
+  # return e
+  #
+  result = n
+  case n.kind
+  of nkStmtList, nkStmtListExpr:
+    if n.len != 0 and n[0].kind == nkYieldStmt:
+      assert(n.len == 2)
+      assert(n[1].kind == nkGotoState)
+
+      result = newNodeI(nkStmtList, n.info)
+      result.add(ctx.newStateAssgn(stateFromGotoState(n[1])))
+
+      var retStmt = newNodeI(nkReturnStmt, n.info)
+      if n[0].sons[0].kind != nkEmpty:
+        var a = newNodeI(nkAsgn, n[0].sons[0].info)
+        var retVal = n[0].sons[0] #liftCapturedVars(n.sons[0], owner, d, c)
+        addSon(a, newSymNode(getClosureIterResult(ctx.fn)))
+        addSon(a, retVal)
+        retStmt.add(a)
+      else:
+        retStmt.add(emptyNode)
+
+      result.add(retStmt)
+    else:
+      for i in 0 ..< n.len:
+        n[i] = ctx.tranformStateAssignments(n[i])
+
+  of nkCharLit..nkUInt64Lit, nkFloatLit..nkFloat128Lit, nkStrLit..nkTripleStrLit,
+      nkSym, nkIdent, procDefs, nkTemplateDef:
+    discard
+
+  of nkReturnStmt:
+
+    result = newNodeI(nkStmtList, n.info)
+    result.add(ctx.newStateAssgn(-1))
+    result.add(n)
+
+  of nkGotoState:
+    result = newNodeI(nkStmtList, n.info)
+    result.add(ctx.newStateAssgn(stateFromGotoState(n)))
+
+    let breakState = newNodeI(nkBreakStmt, n.info)
+    breakState.add(newSymNode(ctx.stateLoopLabel))
+    result.add(breakState)
+
+  else:
+    for i in 0 ..< n.len:
+      n[i] = ctx.tranformStateAssignments(n[i])
+
+proc skipStmtList(n: PNode): PNode =
+  result = n
+  while result.kind in {nkStmtList}:
+    if result.len == 0: return emptyNode
+    result = result[0]
+
+proc skipThroughEmptyStates(ctx: var ClosureIteratorTransformationContext, n: PNode): PNode =
+  result = n
+  case n.kind
+  of nkCharLit..nkUInt64Lit, nkFloatLit..nkFloat128Lit, nkStrLit..nkTripleStrLit,
+      nkSym, nkIdent, procDefs, nkTemplateDef:
+    discard
+  of nkGotoState:
+    var maxJumps = ctx.states.len # maxJumps used only for debugging purposes.
+    result = copyTree(n)
+    while true:
+      let label = result[0].intVal.int
+      if label == ctx.exitStateIdx: break
+      var newLabel = label
+      if label == -1:
+        newLabel = ctx.exitStateIdx
+      else:
+        let fs = ctx.states[label][1].skipStmtList()
+        if fs.kind == nkGotoState:
+          newLabel = fs[0].intVal.int
+      if label == newLabel: break
+      result[0].intVal = newLabel
+      dec maxJumps
+      if maxJumps == 0:
+        assert(false, "Internal error")
+
+    let label = result[0].intVal.int
+    result[0].intVal = ctx.states[label][0].intVal
+  else:
+    for i in 0 ..< n.len:
+      n[i] = ctx.skipThroughEmptyStates(n[i])
+
+proc wrapIntoStateLoop(ctx: var ClosureIteratorTransformationContext, n: PNode): PNode =
+  result = newNode(nkWhileStmt)
+  result.add(newSymNode(getSysSym("true")))
+
+  let loopBody = newNodeI(nkStmtList, n.info)
+  result.add(loopBody)
+
+  if not ctx.stateVarSym.isNil:
+    let varSect = newNodeI(nkVarSection, n.info)
+    addVar(varSect, newSymNode(ctx.stateVarSym))
+    loopBody.add(varSect)
+
+    if not ctx.tempVars.isNil:
+      loopBody.add(ctx.tempVars)
+
+  let blockStmt = newNodeI(nkBlockStmt, n.info)
+  blockStmt.add(newSymNode(ctx.stateLoopLabel))
+
+  let blockBody = newNodeI(nkStmtList, n.info)
+  blockStmt.add(blockBody)
+
+  let gs = newNodeI(nkGotoState, n.info)
+  if ctx.stateVarSym.isNil:
+    gs.add(rawIndirectAccess(newSymNode(getEnvParam(ctx.fn)), getStateField(ctx.fn), n.info))
+  else:
+    gs.add(newSymNode(ctx.stateVarSym))
+
+  gs.add(newIntLit(ctx.states.len - 1))
+  blockBody.add(gs)
+  blockBody.add(n)
+  # gs.add(rawIndirectAccess(newSymNode(ctx.fn.getHiddenParam), getStateField(ctx.fn), n.info))
+
+  loopBody.add(blockStmt)
+
+proc deleteEmptyStates(ctx: var ClosureIteratorTransformationContext) =
+  let goOut = newNode(nkGotoState)
+  goOut.add(newIntLit(-1))
+
+  ctx.exitStateIdx = ctx.newState(goOut, nil)
+
+  # Apply new state indexes and mark unused states with -1
+  var iValid = 0
+  for i, s in ctx.states:
+    let body = s[1].skipStmtList()
+    if body.kind == nkGotoState and i != ctx.states.len - 1:
+      # This is an empty state. Mark with -1.
+      s[0].intVal = -1
+    else:
+      s[0].intVal = iValid
+      inc iValid
+
+  for i, s in ctx.states:
+    let body = s[1].skipStmtList()
+    if body.kind != nkGotoState:
+      discard ctx.skipThroughEmptyStates(s)
+
+  var i = 0
+  while i < ctx.states.len - 1:
+    let fs = ctx.states[i][1].skipStmtList()
+    if fs.kind == nkGotoState:
+      ctx.states.delete(i)
+    else:
+      inc i
+
+proc transformClosureIterator*(fn: PSym, n: PNode): PNode =
+  var ctx: ClosureIteratorTransformationContext
+  ctx.fn = fn
+
+  if getEnvParam(fn).isNil:
+    # Lambda lifting was not done yet. Use temporary :state sym, which
+    # be handled specially by lambda lifting. Local temp vars (if needed)
+    # should folllow the same logic.
+    ctx.stateVarSym = newSym(skVar, getIdent(":state"), fn, fn.info)
+    ctx.stateVarSym.typ = createClosureIterStateType(fn)
+
+  ctx.states = @[]
+  ctx.stateLoopLabel = newSym(skLabel, getIdent(":stateLoop"), fn, fn.info)
+  let n = n.toStmtList
+
+  discard ctx.newState(n, nil)
+  let gotoOut = newNode(nkGotoState)
+  gotoOut.add(newIntLit(-1))
+
+  # Splitting transformation
+  discard ctx.transformClosureIteratorBody(n, gotoOut)
+
+  # Optimize empty states away
+  ctx.deleteEmptyStates()
+
+  # Make new body by concating the list of states
+  result = newNode(nkStmtList)
+  for i, s in ctx.states:
+    # result.add(s)
+    let body = s[1]
+    s.sons.del(1)
+    result.add(s)
+    result.add(body)
+
+  result = ctx.tranformStateAssignments(result)
+
+  # Add excpetion handling
+  var hasExceptions = false
+  if hasExceptions:
+    discard # TODO:
+    # result = wrapIntoTryCatch(result)
+
+  # while true:
+  #   block :stateLoop:
+  #     gotoState
+  #     body
+  result = ctx.wrapIntoStateLoop(result)
+
+  # echo "TRANSFORM TO STATES2: "
+  # debug(result)
+  # echo renderTree(result)
diff --git a/compiler/lambdalifting.nim b/compiler/lambdalifting.nim
index 7757484250..a118edf00f 100644
--- a/compiler/lambdalifting.nim
+++ b/compiler/lambdalifting.nim
@@ -7,11 +7,11 @@
 #    distribution, for details about the copyright.
 #
 
-# This include file implements lambda lifting for the transformator.
+# This file implements lambda lifting for the transformator.
 
 import
-  intsets, strutils, options, ast, astalgo, trees, treetab, msgs, os,
-  idents, renderer, types, magicsys, rodread, lowerings, tables
+  intsets, strutils, options, ast, astalgo, trees, treetab, msgs, os, options,
+  idents, renderer, types, magicsys, rodread, lowerings, tables, sequtils
 
 discard """
   The basic approach is that captured vars need to be put on the heap and
@@ -125,7 +125,7 @@ proc newCall(a: PSym, b: PNode): PNode =
   result.add newSymNode(a)
   result.add b
 
-proc createStateType(iter: PSym): PType =
+proc createClosureIterStateType*(iter: PSym): PType =
   var n = newNodeI(nkRange, iter.info)
   addSon(n, newIntNode(nkIntLit, -1))
   addSon(n, newIntNode(nkIntLit, 0))
@@ -137,7 +137,7 @@ proc createStateType(iter: PSym): PType =
 
 proc createStateField(iter: PSym): PSym =
   result = newSym(skField, getIdent(":state"), iter, iter.info)
-  result.typ = createStateType(iter)
+  result.typ = createClosureIterStateType(iter)
 
 proc createEnvObj(owner: PSym; info: TLineInfo): PType =
   # YYY meh, just add the state field for every closure for now, it's too
@@ -145,7 +145,7 @@ proc createEnvObj(owner: PSym; info: TLineInfo): PType =
   result = createObj(owner, info, final=false)
   rawAddField(result, createStateField(owner))
 
-proc getIterResult(iter: PSym): PSym =
+proc getClosureIterResult*(iter: PSym): PSym =
   if resultPos < iter.ast.len:
     result = iter.ast.sons[resultPos].sym
   else:
@@ -397,7 +397,11 @@ proc detectCapturedVars(n: PNode; owner: PSym; c: var DetectionPass) =
           if not c.capturedVars.containsOrIncl(s.id):
             let obj = getHiddenParam(owner).typ.lastSon
             #let obj = c.getEnvTypeForOwner(s.owner).lastSon
-            addField(obj, s)
+
+            if s.name == getIdent(":state"):
+              obj.n[0].sym.id = -s.id
+            else:
+              addField(obj, s)
       # but always return because the rest of the proc is only relevant when
       # ow != owner:
       return
@@ -461,6 +465,7 @@ type
     processed: IntSet
     envVars: Table[int, PNode]
     inContainer: int
+    features: set[Feature]
 
 proc initLiftingPass(fn: PSym): LiftingPass =
   result.processed = initIntSet()
@@ -595,7 +600,7 @@ proc accessViaEnvVar(n: PNode; owner: PSym; d: DetectionPass;
     localError(n.info, "internal error: not part of closure object type")
     result = n
 
-proc getStateField(owner: PSym): PSym =
+proc getStateField*(owner: PSym): PSym =
   getHiddenParam(owner).typ.sons[0].n.sons[0].sym
 
 proc liftCapturedVars(n: PNode; owner: PSym; d: DetectionPass;
@@ -621,7 +626,7 @@ proc transformYield(n: PNode; owner: PSym; d: DetectionPass;
   if n.sons[0].kind != nkEmpty:
     var a = newNodeI(nkAsgn, n.sons[0].info)
     var retVal = liftCapturedVars(n.sons[0], owner, d, c)
-    addSon(a, newSymNode(getIterResult(owner)))
+    addSon(a, newSymNode(getClosureIterResult(owner)))
     addSon(a, retVal)
     retStmt.add(a)
   else:
@@ -713,7 +718,9 @@ proc liftCapturedVars(n: PNode; owner: PSym; d: DetectionPass;
         #  echo renderTree(s.getBody, {renderIds})
         let oldInContainer = c.inContainer
         c.inContainer = 0
-        let body = wrapIterBody(liftCapturedVars(s.getBody, s, d, c), s)
+        var body = liftCapturedVars(s.getBody, s, d, c)
+        if oldIterTransf in c.features:
+          body = wrapIterBody(body, s)
         if c.envvars.getOrDefault(s.id).isNil:
           s.ast.sons[bodyPos] = body
         else:
@@ -756,9 +763,9 @@ proc liftCapturedVars(n: PNode; owner: PSym; d: DetectionPass;
       if n[1].kind == nkClosure: result = n[1]
   else:
     if owner.isIterator:
-      if n.kind == nkYieldStmt:
+      if oldIterTransf in c.features and n.kind == nkYieldStmt:
         return transformYield(n, owner, d, c)
-      elif n.kind == nkReturnStmt:
+      elif oldIterTransf in c.features and n.kind == nkReturnStmt:
         return transformReturn(n, owner, d, c)
       elif nfLL in n.flags:
         # special case 'when nimVm' due to bug #3636:
@@ -805,7 +812,7 @@ proc liftIterToProc*(fn: PSym; body: PNode; ptrType: PType): PNode =
   fn.kind = oldKind
   fn.typ.callConv = oldCC
 
-proc liftLambdas*(fn: PSym, body: PNode; tooEarly: var bool): PNode =
+proc liftLambdas*(features: set[Feature], fn: PSym, body: PNode; tooEarly: var bool): PNode =
   # XXX gCmd == cmdCompileToJS does not suffice! The compiletime stuff needs
   # the transformation even when compiling to JS ...
 
@@ -815,6 +822,7 @@ proc liftLambdas*(fn: PSym, body: PNode; tooEarly: var bool): PNode =
   if body.kind == nkEmpty or (
       gCmd == cmdCompileToJS and not isCompileTime) or
       fn.skipGenericOwner.kind != skModule:
+
     # ignore forward declaration:
     result = body
     tooEarly = true
@@ -826,10 +834,13 @@ proc liftLambdas*(fn: PSym, body: PNode; tooEarly: var bool): PNode =
       d.somethingToDo = true
     if d.somethingToDo:
       var c = initLiftingPass(fn)
-      var newBody = liftCapturedVars(body, fn, d, c)
+      c.features = features
+      result = liftCapturedVars(body, fn, d, c)
       if c.envvars.getOrDefault(fn.id) != nil:
-        newBody = newTree(nkStmtList, rawClosureCreation(fn, d, c), newBody)
-      result = wrapIterBody(newBody, fn)
+        result = newTree(nkStmtList, rawClosureCreation(fn, d, c), result)
+
+      if oldIterTransf in features:
+        result = wrapIterBody(result, fn)
     else:
       result = body
     #if fn.name.s == "get2":
@@ -870,7 +881,9 @@ proc liftForLoop*(body: PNode; owner: PSym): PNode =
       cl = createClosure()
       while true:
         let i = foo(cl)
-        nkBreakState(cl.state)
+        if cl.state < 0:
+          break
+        # nkBreakState(cl.state)
         ...
     """
   if liftingHarmful(owner): return body
@@ -930,5 +943,16 @@ proc liftForLoop*(body: PNode; owner: PSym): PNode =
   loopBody.sons[0] = v2
   var bs = newNodeI(nkBreakState, body.info)
   bs.addSon(call.sons[0])
-  loopBody.sons[1] = bs
+
+  let ibs = newNode(nkIfStmt)
+  let elifBranch = newNode(nkElifBranch)
+  elifBranch.add(bs)
+
+  let br = newNode(nkBreakStmt)
+  br.add(emptyNode)
+
+  elifBranch.add(br)
+  ibs.add(elifBranch)
+
+  loopBody.sons[1] = ibs
   loopBody.sons[2] = body[L-1]
diff --git a/compiler/options.nim b/compiler/options.nim
index f8cb735ae6..0ce2f95ce9 100644
--- a/compiler/options.nim
+++ b/compiler/options.nim
@@ -110,7 +110,8 @@ type
     callOperator,
     parallel,
     destructor,
-    notnil
+    notnil,
+    oldIterTransf
 
   ConfigRef* = ref object ## eventually all global configuration should be moved here
     linesCompiled*: int  # all lines that have been compiled
diff --git a/compiler/renderer.nim b/compiler/renderer.nim
index 996168412b..0c861bdb88 100644
--- a/compiler/renderer.nim
+++ b/compiler/renderer.nim
@@ -1411,11 +1411,21 @@ proc gsub(g: var TSrcGen, n: PNode, c: TContext) =
     put(g, tkParLe, "(ComesFrom|")
     gsub(g, n, 0)
     put(g, tkParRi, ")")
-  of nkGotoState, nkState:
+  of nkGotoState:
     var c: TContext
     initContext c
-    putWithSpace g, tkSymbol, if n.kind == nkState: "state" else: "goto"
+    putWithSpace g, tkSymbol, "goto"
     gsons(g, n, c)
+  of nkState:
+    var c: TContext
+    initContext c
+    putWithSpace g, tkSymbol, "state"
+    gsub(g, n[0], c)
+    putWithSpace(g, tkColon, ":")
+    indentNL(g)
+    gsons(g, n, c, 1)
+    dedent(g)
+
   of nkBreakState:
     put(g, tkTuple, "breakstate")
   of nkTypeClassTy:
diff --git a/compiler/seminst.nim b/compiler/seminst.nim
index 32b3853089..9ec7d87985 100644
--- a/compiler/seminst.nim
+++ b/compiler/seminst.nim
@@ -145,7 +145,7 @@ proc instantiateBody(c: PContext, n, params: PNode, result, orig: PSym) =
     freshGenSyms(b, result, orig, symMap)
     b = semProcBody(c, b)
     b = hloBody(c, b)
-    n.sons[bodyPos] = transformBody(c.module, b, result)
+    n.sons[bodyPos] = transformBody(c, b, result)
     #echo "code instantiated ", result.name.s
     excl(result.flags, sfForward)
     dec c.inGenericInst
diff --git a/compiler/semstmts.nim b/compiler/semstmts.nim
index f3cf4196f2..dfee20a99a 100644
--- a/compiler/semstmts.nim
+++ b/compiler/semstmts.nim
@@ -1236,7 +1236,7 @@ proc semLambda(c: PContext, n: PNode, flags: TExprFlags): PNode =
       addResult(c, s.typ.sons[0], n.info, skProc)
       addResultNode(c, n)
       let semBody = hloBody(c, semProcBody(c, n.sons[bodyPos]))
-      n.sons[bodyPos] = transformBody(c.module, semBody, s)
+      n.sons[bodyPos] = transformBody(c, semBody, s)
       popProcCon(c)
     elif efOperand notin flags:
       localError(n.info, errGenericLambdaNotAllowed)
@@ -1277,7 +1277,7 @@ proc semInferredLambda(c: PContext, pt: TIdTable, n: PNode): PNode =
   addResult(c, n.typ.sons[0], n.info, skProc)
   addResultNode(c, n)
   let semBody = hloBody(c, semProcBody(c, n.sons[bodyPos]))
-  n.sons[bodyPos] = transformBody(c.module, semBody, s)
+  n.sons[bodyPos] = transformBody(c, semBody, s)
   popProcCon(c)
   popOwner(c)
   closeScope(c)
@@ -1590,7 +1590,7 @@ proc semProcAux(c: PContext, n: PNode, kind: TSymKind,
           let semBody = hloBody(c, semProcBody(c, n.sons[bodyPos]))
           # unfortunately we cannot skip this step when in 'system.compiles'
           # context as it may even be evaluated in 'system.compiles':
-          n.sons[bodyPos] = transformBody(c.module, semBody, s)
+          n.sons[bodyPos] = transformBody(c, semBody, s)
       else:
         if s.typ.sons[0] != nil and kind != skIterator:
           addDecl(c, newSym(skUnknown, getIdent"result", nil, n.info))
diff --git a/compiler/transf.nim b/compiler/transf.nim
index f7ec6c97f2..c0f5e5e327 100644
--- a/compiler/transf.nim
+++ b/compiler/transf.nim
@@ -19,9 +19,10 @@
 # * transforms 'defer' into a 'try finally' statement
 
 import
-  intsets, strutils, options, ast, astalgo, trees, treetab, msgs, os,
-  idents, renderer, types, passes, semfold, magicsys, cgmeth, rodread,
-  lambdalifting, sempass2, lowerings, lookups, destroyer, liftlocals
+  intsets, strutils, options, ast, astalgo, trees, treetab, msgs, os, lookups,
+  idents, renderer, types, passes, semfold, magicsys, cgmeth, rodread, semdata,
+  lambdalifting, sempass2, lowerings, destroyer, liftlocals, closureiters
+
 
 type
   PTransNode* = distinct PNode
@@ -967,20 +968,22 @@ template liftDefer(c, root) =
   if c.deferDetected:
     liftDeferAux(root)
 
-proc transformBody*(module: PSym, n: PNode, prc: PSym): PNode =
-  if nfTransf in n.flags or prc.kind in {skTemplate}:
-    result = n
-  else:
-    var c = openTransf(module, "")
-    result = liftLambdas(prc, n, c.tooEarly)
-    #result = n
+proc transformBody*(ctx: PContext, n: PNode, prc: PSym): PNode =
+  result = n
+  if nfTransf notin n.flags and prc.kind notin {skTemplate}:
+    var c = openTransf(ctx.module, "")
+    result = liftLambdas(ctx.features, prc, result, c.tooEarly)
     result = processTransf(c, result, prc)
     liftDefer(c, result)
-    #result = liftLambdas(prc, result)
+
     when useEffectSystem: trackProc(prc, result)
     result = liftLocalsIfRequested(prc, result)
     if c.needsDestroyPass: #and newDestructors:
       result = injectDestructorCalls(prc, result)
+
+    if prc.isIterator and oldIterTransf notin ctx.features:
+      result = transformClosureIterator(prc, result)
+
     incl(result.flags, nfTransf)
       #if prc.name.s == "testbody":
     #  echo renderTree(result)
diff --git a/tests/async/tasync_in_seq_constr.nim b/tests/async/tasync_in_seq_constr.nim
index 46ad744512..cf9bb54516 100644
--- a/tests/async/tasync_in_seq_constr.nim
+++ b/tests/async/tasync_in_seq_constr.nim
@@ -1,6 +1,5 @@
 discard """
-  errormsg: "invalid control flow: 'yield' within a constructor"
-  line: 16
+  output: "@[1, 2, 3, 4]"
 """
 
 # bug #5314, bug #6626

From 7d38db284ba2655bf19cb9c0785240616074ee44 Mon Sep 17 00:00:00 2001
From: Yuriy Glukhov <yuriy.glukhov@gmail.com>
Date: Mon, 30 Apr 2018 21:39:41 +0300
Subject: [PATCH 02/14] Extended tasync_in_seq_constr test

---
 tests/async/tasync_in_seq_constr.nim | 20 ++++++++++++++------
 1 file changed, 14 insertions(+), 6 deletions(-)

diff --git a/tests/async/tasync_in_seq_constr.nim b/tests/async/tasync_in_seq_constr.nim
index cf9bb54516..3d6dae2457 100644
--- a/tests/async/tasync_in_seq_constr.nim
+++ b/tests/async/tasync_in_seq_constr.nim
@@ -1,17 +1,25 @@
 discard """
-  output: "@[1, 2, 3, 4]"
+  output: '''
+@[1, 2, 3, 4]
+123
+'''
 """
 
 # bug #5314, bug #6626
 
 import asyncdispatch
 
-proc bar(): Future[int] {.async.} =
-    await sleepAsync(500)
-    result = 3
+proc bar(i: int): Future[int] {.async.} =
+    await sleepAsync(2)
+    result = i
 
 proc foo(): Future[seq[int]] {.async.} =
-    await sleepAsync(500)
-    result = @[1, 2, await bar(), 4] # <--- The bug is here
+    await sleepAsync(2)
+    result = @[1, 2, await bar(3), 4] # <--- The bug is here
+
+proc foo2() {.async.} =
+    await sleepAsync(2)
+    echo(await bar(1), await bar(2), await bar(3))
 
 echo waitFor foo()
+waitFor foo2()

From 13167c85f6192917ebfbf224b0cc53c11ea75710 Mon Sep 17 00:00:00 2001
From: Yuriy Glukhov <yuriy.glukhov@gmail.com>
Date: Tue, 1 May 2018 10:04:52 +0300
Subject: [PATCH 03/14] Cosmetics

---
 compiler/closureiters.nim | 45 ++++++++++++++++++++-------------------
 1 file changed, 23 insertions(+), 22 deletions(-)

diff --git a/compiler/closureiters.nim b/compiler/closureiters.nim
index 02795ab470..c94d90bd14 100644
--- a/compiler/closureiters.nim
+++ b/compiler/closureiters.nim
@@ -65,18 +65,19 @@ import
   idents, renderer, types, magicsys, rodread, lowerings, tables, sequtils,
   lambdalifting
 
-type ClosureIteratorTransformationContext = object
-  fn: PSym
-  stateVarSym: PSym # :state variable. nil if env already introduced by lambdalifting
-  states: seq[PNode] # The resulting states. Every state is an nkState node.
-  blockLevel: int # Temp used to transform break and continue stmts
-  stateLoopLabel: PSym # Label to break on, when jumping between states.
-  exitStateIdx: int # index of the last state
-  tempVarId: int # unique name counter
-  tempVars: PNode # Temp var decls, nkVarSection
-  loweredStmtListExpr: PNode # Temporary used for nkStmtListExpr lowering
+type
+  Ctx = object
+    fn: PSym
+    stateVarSym: PSym # :state variable. nil if env already introduced by lambdalifting
+    states: seq[PNode] # The resulting states. Every state is an nkState node.
+    blockLevel: int # Temp used to transform break and continue stmts
+    stateLoopLabel: PSym # Label to break on, when jumping between states.
+    exitStateIdx: int # index of the last state
+    tempVarId: int # unique name counter
+    tempVars: PNode # Temp var decls, nkVarSection
+    loweredStmtListExpr: PNode # Temporary used for nkStmtListExpr lowering
 
-proc newStateAssgn(ctx: var ClosureIteratorTransformationContext, stateNo: int = -2): PNode =
+proc newStateAssgn(ctx: var Ctx, stateNo: int = -2): PNode =
   # Creates state assignmen:
   #   :state = stateNo
 
@@ -95,7 +96,7 @@ proc setStateInAssgn(stateAssgn: PNode, stateNo: int) =
   assert stateAssgn[1].kind == nkIntLit
   stateAssgn[1].intVal = stateNo
 
-proc newState(ctx: var ClosureIteratorTransformationContext, n, gotoOut: PNode): int =
+proc newState(ctx: var Ctx, n, gotoOut: PNode): int =
   # Creates a new state, adds it to the context fills out `gotoOut` so that it
   # will goto this state.
   # Returns index of the newly created state
@@ -123,7 +124,7 @@ proc addGotoOut(n: PNode, gotoOut: PNode): PNode =
   if result.len != 0 and result.sons[^1].kind != nkGotoState:
     result.add(gotoOut)
 
-proc newTempVarAccess(ctx: var ClosureIteratorTransformationContext, typ: PType, i: TLineInfo): PNode =
+proc newTempVarAccess(ctx: var Ctx, typ: PType, i: TLineInfo): PNode =
   if not ctx.stateVarSym.isNil:
     # We haven't gone through labmda lifting yet, so just create a local var,
     # it will be lifted later
@@ -157,7 +158,7 @@ proc hasYields(n: PNode): bool =
         result = true
         break
 
-proc transformBreaksAndContinuesInWhile(ctx: var ClosureIteratorTransformationContext, n: PNode, before, after: PNode): PNode =
+proc transformBreaksAndContinuesInWhile(ctx: var Ctx, n: PNode, before, after: PNode): PNode =
   result = n
   case n.kind
   of nkCharLit..nkUInt64Lit, nkFloatLit..nkFloat128Lit, nkStrLit..nkTripleStrLit,
@@ -177,7 +178,7 @@ proc transformBreaksAndContinuesInWhile(ctx: var ClosureIteratorTransformationCo
     for i in 0 ..< n.len:
       n[i] = ctx.transformBreaksAndContinuesInWhile(n[i], before, after)
 
-proc transformBreaksInBlock(ctx: var ClosureIteratorTransformationContext, n: PNode, label, after: PNode): PNode =
+proc transformBreaksInBlock(ctx: var Ctx, n: PNode, label, after: PNode): PNode =
   result = n
   case n.kind
   of nkCharLit..nkUInt64Lit, nkFloatLit..nkFloat128Lit, nkStrLit..nkTripleStrLit,
@@ -252,7 +253,7 @@ proc hasYieldsInExpressions(n: PNode): bool =
       if c.hasYieldsInExpressions:
         return true
 
-proc lowerStmtListExpr(ctx: var ClosureIteratorTransformationContext, n: PNode): PNode =
+proc lowerStmtListExpr(ctx: var Ctx, n: PNode): PNode =
   result = n
   case n.kind
   of nkCharLit..nkUInt64Lit, nkFloatLit..nkFloat128Lit, nkStrLit..nkTripleStrLit,
@@ -274,7 +275,7 @@ proc lowerStmtListExpr(ctx: var ClosureIteratorTransformationContext, n: PNode):
     for i in 0 ..< n.len:
       n[i] = ctx.lowerStmtListExpr(n[i])
 
-proc transformClosureIteratorBody(ctx: var ClosureIteratorTransformationContext, n: PNode, gotoOut: PNode): PNode =
+proc transformClosureIteratorBody(ctx: var Ctx, n: PNode, gotoOut: PNode): PNode =
   result = n
   case n.kind:
     of nkCharLit..nkUInt64Lit, nkFloatLit..nkFloat128Lit, nkStrLit..nkTripleStrLit,
@@ -412,7 +413,7 @@ proc stateFromGotoState(n: PNode): int =
   assert(n.kind == nkGotoState)
   result = n[0].intVal.int
 
-proc tranformStateAssignments(ctx: var ClosureIteratorTransformationContext, n: PNode): PNode =
+proc tranformStateAssignments(ctx: var Ctx, n: PNode): PNode =
   # This transforms 3 patterns:
   ########################## 1
   # yield e
@@ -484,7 +485,7 @@ proc skipStmtList(n: PNode): PNode =
     if result.len == 0: return emptyNode
     result = result[0]
 
-proc skipThroughEmptyStates(ctx: var ClosureIteratorTransformationContext, n: PNode): PNode =
+proc skipThroughEmptyStates(ctx: var Ctx, n: PNode): PNode =
   result = n
   case n.kind
   of nkCharLit..nkUInt64Lit, nkFloatLit..nkFloat128Lit, nkStrLit..nkTripleStrLit,
@@ -515,7 +516,7 @@ proc skipThroughEmptyStates(ctx: var ClosureIteratorTransformationContext, n: PN
     for i in 0 ..< n.len:
       n[i] = ctx.skipThroughEmptyStates(n[i])
 
-proc wrapIntoStateLoop(ctx: var ClosureIteratorTransformationContext, n: PNode): PNode =
+proc wrapIntoStateLoop(ctx: var Ctx, n: PNode): PNode =
   result = newNode(nkWhileStmt)
   result.add(newSymNode(getSysSym("true")))
 
@@ -549,7 +550,7 @@ proc wrapIntoStateLoop(ctx: var ClosureIteratorTransformationContext, n: PNode):
 
   loopBody.add(blockStmt)
 
-proc deleteEmptyStates(ctx: var ClosureIteratorTransformationContext) =
+proc deleteEmptyStates(ctx: var Ctx) =
   let goOut = newNode(nkGotoState)
   goOut.add(newIntLit(-1))
 
@@ -580,7 +581,7 @@ proc deleteEmptyStates(ctx: var ClosureIteratorTransformationContext) =
       inc i
 
 proc transformClosureIterator*(fn: PSym, n: PNode): PNode =
-  var ctx: ClosureIteratorTransformationContext
+  var ctx: Ctx
   ctx.fn = fn
 
   if getEnvParam(fn).isNil:

From 48d8e215d53b5b231483b208def8b3ed8e5fdf85 Mon Sep 17 00:00:00 2001
From: Yuriy Glukhov <yuriy.glukhov@gmail.com>
Date: Tue, 1 May 2018 10:31:49 +0300
Subject: [PATCH 04/14] Don't leak sem PContext into transf

---
 compiler/seminst.nim  |  2 +-
 compiler/semstmts.nim |  6 +++---
 compiler/transf.nim   | 11 +++++------
 3 files changed, 9 insertions(+), 10 deletions(-)

diff --git a/compiler/seminst.nim b/compiler/seminst.nim
index 9ec7d87985..0513e23956 100644
--- a/compiler/seminst.nim
+++ b/compiler/seminst.nim
@@ -145,7 +145,7 @@ proc instantiateBody(c: PContext, n, params: PNode, result, orig: PSym) =
     freshGenSyms(b, result, orig, symMap)
     b = semProcBody(c, b)
     b = hloBody(c, b)
-    n.sons[bodyPos] = transformBody(c, b, result)
+    n.sons[bodyPos] = transformBody(c.module, c.features, b, result)
     #echo "code instantiated ", result.name.s
     excl(result.flags, sfForward)
     dec c.inGenericInst
diff --git a/compiler/semstmts.nim b/compiler/semstmts.nim
index dfee20a99a..9f00e877f2 100644
--- a/compiler/semstmts.nim
+++ b/compiler/semstmts.nim
@@ -1236,7 +1236,7 @@ proc semLambda(c: PContext, n: PNode, flags: TExprFlags): PNode =
       addResult(c, s.typ.sons[0], n.info, skProc)
       addResultNode(c, n)
       let semBody = hloBody(c, semProcBody(c, n.sons[bodyPos]))
-      n.sons[bodyPos] = transformBody(c, semBody, s)
+      n.sons[bodyPos] = transformBody(c.module, c.features, semBody, s)
       popProcCon(c)
     elif efOperand notin flags:
       localError(n.info, errGenericLambdaNotAllowed)
@@ -1277,7 +1277,7 @@ proc semInferredLambda(c: PContext, pt: TIdTable, n: PNode): PNode =
   addResult(c, n.typ.sons[0], n.info, skProc)
   addResultNode(c, n)
   let semBody = hloBody(c, semProcBody(c, n.sons[bodyPos]))
-  n.sons[bodyPos] = transformBody(c, semBody, s)
+  n.sons[bodyPos] = transformBody(c.module, c.features, semBody, s)
   popProcCon(c)
   popOwner(c)
   closeScope(c)
@@ -1590,7 +1590,7 @@ proc semProcAux(c: PContext, n: PNode, kind: TSymKind,
           let semBody = hloBody(c, semProcBody(c, n.sons[bodyPos]))
           # unfortunately we cannot skip this step when in 'system.compiles'
           # context as it may even be evaluated in 'system.compiles':
-          n.sons[bodyPos] = transformBody(c, semBody, s)
+          n.sons[bodyPos] = transformBody(c.module, c.features, semBody, s)
       else:
         if s.typ.sons[0] != nil and kind != skIterator:
           addDecl(c, newSym(skUnknown, getIdent"result", nil, n.info))
diff --git a/compiler/transf.nim b/compiler/transf.nim
index c0f5e5e327..75c7f2b6c3 100644
--- a/compiler/transf.nim
+++ b/compiler/transf.nim
@@ -20,10 +20,9 @@
 
 import
   intsets, strutils, options, ast, astalgo, trees, treetab, msgs, os, lookups,
-  idents, renderer, types, passes, semfold, magicsys, cgmeth, rodread, semdata,
+  idents, renderer, types, passes, semfold, magicsys, cgmeth, rodread,
   lambdalifting, sempass2, lowerings, destroyer, liftlocals, closureiters
 
-
 type
   PTransNode* = distinct PNode
 
@@ -968,11 +967,11 @@ template liftDefer(c, root) =
   if c.deferDetected:
     liftDeferAux(root)
 
-proc transformBody*(ctx: PContext, n: PNode, prc: PSym): PNode =
+proc transformBody*(module: PSym, features: set[Feature], n: PNode, prc: PSym): PNode =
   result = n
   if nfTransf notin n.flags and prc.kind notin {skTemplate}:
-    var c = openTransf(ctx.module, "")
-    result = liftLambdas(ctx.features, prc, result, c.tooEarly)
+    var c = openTransf(module, "")
+    result = liftLambdas(features, prc, result, c.tooEarly)
     result = processTransf(c, result, prc)
     liftDefer(c, result)
 
@@ -981,7 +980,7 @@ proc transformBody*(ctx: PContext, n: PNode, prc: PSym): PNode =
     if c.needsDestroyPass: #and newDestructors:
       result = injectDestructorCalls(prc, result)
 
-    if prc.isIterator and oldIterTransf notin ctx.features:
+    if prc.isIterator and oldIterTransf notin features:
       result = transformClosureIterator(prc, result)
 
     incl(result.flags, nfTransf)

From 0ed6c3e476e421827d081e9ab0d9fcb0d3de5eb2 Mon Sep 17 00:00:00 2001
From: Yuriy Glukhov <yuriy.glukhov@gmail.com>
Date: Tue, 1 May 2018 13:19:01 +0300
Subject: [PATCH 05/14] Minor dry up

---
 compiler/closureiters.nim | 8 +++-----
 1 file changed, 3 insertions(+), 5 deletions(-)

diff --git a/compiler/closureiters.nim b/compiler/closureiters.nim
index c94d90bd14..7653176de2 100644
--- a/compiler/closureiters.nim
+++ b/compiler/closureiters.nim
@@ -125,12 +125,12 @@ proc addGotoOut(n: PNode, gotoOut: PNode): PNode =
     result.add(gotoOut)
 
 proc newTempVarAccess(ctx: var Ctx, typ: PType, i: TLineInfo): PNode =
+  let s = newSym(skVar, getIdent(":tmpSlLower" & $ctx.tempVarId), ctx.fn, i)
+  s.typ = typ
+
   if not ctx.stateVarSym.isNil:
     # We haven't gone through labmda lifting yet, so just create a local var,
     # it will be lifted later
-    let s = newSym(skVar, getIdent(":tmpSlLower" & $ctx.tempVarId), ctx.fn, i)
-    s.typ = typ
-
     if ctx.tempVars.isNil:
       ctx.tempVars = newNode(nkVarSection)
       addVar(ctx.tempVars, newSymNode(s))
@@ -138,8 +138,6 @@ proc newTempVarAccess(ctx: var Ctx, typ: PType, i: TLineInfo): PNode =
     result = newSymNode(s)
   else:
     # Lambda lifting is done, insert temp var to env.
-    let s = newSym(skVar, getIdent(":tmpSlLower" & $ctx.tempVarId), ctx.fn, i)
-    s.typ = typ
     result = freshVarForClosureIter(s, ctx.fn)
 
   inc ctx.tempVarId

From ce634909281ffc8efbc7d192f557ffe38f49e740 Mon Sep 17 00:00:00 2001
From: Yuriy Glukhov <yuriy.glukhov@gmail.com>
Date: Fri, 4 May 2018 15:23:47 +0300
Subject: [PATCH 06/14] Yield in try

---
 compiler/closureiters.nim  | 590 +++++++++++++++++++++++++++++++------
 compiler/semexprs.nim      |   2 +-
 lib/system/embedded.nim    |   3 +
 lib/system/excpt.nim       |   4 +
 tests/async/tasynctry2.nim |   4 +-
 tests/iter/tyieldintry.nim | 201 +++++++++++++
 6 files changed, 711 insertions(+), 93 deletions(-)
 create mode 100644 tests/iter/tyieldintry.nim

diff --git a/compiler/closureiters.nim b/compiler/closureiters.nim
index 7653176de2..504f70347d 100644
--- a/compiler/closureiters.nim
+++ b/compiler/closureiters.nim
@@ -59,6 +59,80 @@
 #   if :tmpSlLower == 2:
 #     yield 3
 
+# nkTryStmt Transformations:
+# If the iter has an nkTryStmt with a yield inside
+#  - the closure iter is promoted to have exceptions (ctx.hasExceptions = true)
+#  - exception table is created. This is a const array, where
+#    `abs(exceptionTable[i])` is a state idx to which we should jump from state
+#    `i` should exception be raised in state `i`. For all states in `try` block
+#    the target state is `except` block. For all states in `except` block
+#    the target state is `finally` block. For all other states there is no
+#    target state (0, as the first block can never be neither except nor finally).
+#    `exceptionTable[i]` is < 0 if `abs(exceptionTable[i])` is except block,
+#    and > 0, for finally block.
+#  - local variable :curExc is created
+#  - the iter body is wrapped into a
+#      try:
+#       closureIterSetupExc(:curExc)
+#       ...body...
+#      catch:
+#        :state = exceptionTable[:state]
+#        if :state == 0: raise # No state that could handle exception
+#        :unrollFinally = :state > 0 # Target state is finally
+#        if :state < 0:
+#           :state = -:state
+#        :curExc = getCurrentException()
+#
+# nkReturnStmt within a try/except/finally now has to behave differently as we
+# want the nearest finally block to be executed before the return, thus it is
+# transformed to:
+#  :tmpResult = returnValue (if return doesn't have a value, this is skipped)
+#  :unrollFinally = true
+#  goto nearestFinally (or -1 if not exists)
+#
+# Every finally block calls closureIterEndFinally() upon its successful
+# completion.
+#
+# Example:
+#
+# try:
+#  yield 0
+#  raise ...
+# except:
+#  yield 1
+#  return 3
+# finally:
+#  yield 2
+#
+# Is transformed to (yields are left in place for example simplicity,
+#    in reality the code is subdivided even more, as described above):
+#
+# STATE0: # Try
+#   yield 0
+#   raise ...
+#   :state = 2 # What would happen should we not raise
+#   break :stateLoop
+# STATE1: # Except
+#   yield 1
+#   :tmpResult = 3           # Return
+#   :unrollFinally = true # Return
+#   :state = 2 # Goto Finally
+#   break :stateLoop
+#   :state = 2 # What would happen should we not return
+#   break :stateLoop
+# STATE2: # Finally
+#   yield 2
+#   if :unrollFinally: # This node is created by `newEndFinallyNode`
+#     when nearestFinally == 0: # Pseudocode. The `when` is not emitted in reality
+#       if :curExc.isNil:
+#         return :tmpResult
+#       else:
+#         raise
+#     else:
+#       :state = nearestFinally
+#       break :stateLoop
+#   state = -1 # Goto next state. In this case we just exit
+#   break :stateLoop
 
 import
   intsets, strutils, options, ast, astalgo, trees, treetab, msgs, os, options,
@@ -69,6 +143,10 @@ type
   Ctx = object
     fn: PSym
     stateVarSym: PSym # :state variable. nil if env already introduced by lambdalifting
+    tmpResultSym: PSym # Used when we return, but finally has to interfere
+    unrollFinallySym: PSym # Indicates that we're unrolling finally states (either exception happened or premature return)
+    curExcSym: PSym # Current exception
+
     states: seq[PNode] # The resulting states. Every state is an nkState node.
     blockLevel: int # Temp used to transform break and continue stmts
     stateLoopLabel: PSym # Label to break on, when jumping between states.
@@ -76,20 +154,66 @@ type
     tempVarId: int # unique name counter
     tempVars: PNode # Temp var decls, nkVarSection
     loweredStmtListExpr: PNode # Temporary used for nkStmtListExpr lowering
+    exceptionTable: seq[int] # For state `i` jump to state `exceptionTable[i]` if exception is raised
+    hasExceptions: bool # Does closure have yield in try?
+    curExcHandlingState: int # Negative for except, positive for finally
+    nearestFinally: int # Index of the nearest finally block. For try/except it
+                    # is their finally. For finally it is parent finally. Otherwise -1
+
+proc newStateAccess(ctx: var Ctx): PNode =
+  if ctx.stateVarSym.isNil:
+    result = rawIndirectAccess(newSymNode(getEnvParam(ctx.fn)), getStateField(ctx.fn), ctx.fn.info)
+  else:
+    result = newSymNode(ctx.stateVarSym)
+
+proc newStateAssgn(ctx: var Ctx, toValue: PNode): PNode =
+  # Creates state assignment:
+  #   :state = toValue
+  result = newNode(nkAsgn)
+  result.add(ctx.newStateAccess())
+  result.add(toValue)
 
 proc newStateAssgn(ctx: var Ctx, stateNo: int = -2): PNode =
-  # Creates state assignmen:
+  # Creates state assignment:
   #   :state = stateNo
+  ctx.newStateAssgn(newIntTypeNode(nkIntLit, stateNo, getSysType(tyInt)))
 
-  result = newNode(nkAsgn)
-  if ctx.stateVarSym.isNil:
-    let state = getStateField(ctx.fn)
-    assert state != nil
-    result.add(rawIndirectAccess(newSymNode(getEnvParam(ctx.fn)),
-                      state, result.info))
+proc newEnvVar(ctx: var Ctx, name: string, typ: PType): PSym =
+  result = newSym(skVar, getIdent(name), ctx.fn, ctx.fn.info)
+  result.typ = typ
+
+  if not ctx.stateVarSym.isNil:
+    # We haven't gone through labmda lifting yet, so just create a local var,
+    # it will be lifted later
+    if ctx.tempVars.isNil:
+      ctx.tempVars = newNode(nkVarSection)
+      addVar(ctx.tempVars, newSymNode(result))
   else:
-    result.add(newSymNode(ctx.stateVarSym))
-  result.add(newIntTypeNode(nkIntLit, stateNo, getSysType(tyInt)))
+    let envParam = getEnvParam(ctx.fn)
+    # let obj = envParam.typ.lastSon
+    result = addUniqueField(envParam.typ.lastSon, result)
+
+proc newEnvVarAccess(ctx: Ctx, s: PSym): PNode =
+  if ctx.stateVarSym.isNil:
+    result = rawIndirectAccess(newSymNode(getEnvParam(ctx.fn)), s, ctx.fn.info)
+  else:
+    result = newSymNode(s)
+
+proc newTmpResultAccess(ctx: var Ctx): PNode =
+  if ctx.tmpResultSym.isNil:
+    debug(ctx.fn.typ)
+    ctx.tmpResultSym = ctx.newEnvVar(":tmpResult", ctx.fn.typ[0])
+  ctx.newEnvVarAccess(ctx.tmpResultSym)
+
+proc newUnrollFinallyAccess(ctx: var Ctx): PNode =
+  if ctx.unrollFinallySym.isNil:
+    ctx.unrollFinallySym = ctx.newEnvVar(":unrollFinally", getSysType(tyBool))
+  ctx.newEnvVarAccess(ctx.unrollFinallySym)
+
+proc newCurExcAccess(ctx: var Ctx): PNode =
+  if ctx.curExcSym.isNil:
+    ctx.curExcSym = ctx.newEnvVar(":curExc", callCodegenProc("getCurrentException", emptyNode).typ)
+  ctx.newEnvVarAccess(ctx.curExcSym)
 
 proc setStateInAssgn(stateAssgn: PNode, stateNo: int) =
   assert stateAssgn.kind == nkAsgn
@@ -107,6 +231,8 @@ proc newState(ctx: var Ctx, n, gotoOut: PNode): int =
   s.add(resLit)
   s.add(n)
   ctx.states.add(s)
+  ctx.exceptionTable.add(ctx.curExcHandlingState)
+
   if not gotoOut.isNil:
     assert(gotoOut.len == 0)
     gotoOut.add(newIntLit(result))
@@ -119,27 +245,13 @@ proc toStmtList(n: PNode): PNode =
 
 proc addGotoOut(n: PNode, gotoOut: PNode): PNode =
   # Make sure `n` is a stmtlist, and ends with `gotoOut`
-
   result = toStmtList(n)
   if result.len != 0 and result.sons[^1].kind != nkGotoState:
     result.add(gotoOut)
 
 proc newTempVarAccess(ctx: var Ctx, typ: PType, i: TLineInfo): PNode =
-  let s = newSym(skVar, getIdent(":tmpSlLower" & $ctx.tempVarId), ctx.fn, i)
-  s.typ = typ
-
-  if not ctx.stateVarSym.isNil:
-    # We haven't gone through labmda lifting yet, so just create a local var,
-    # it will be lifted later
-    if ctx.tempVars.isNil:
-      ctx.tempVars = newNode(nkVarSection)
-      addVar(ctx.tempVars, newSymNode(s))
-
-    result = newSymNode(s)
-  else:
-    # Lambda lifting is done, insert temp var to env.
-    result = freshVarForClosureIter(s, ctx.fn)
-
+  let s = ctx.newEnvVar(":tmpSlLower" & $ctx.tempVarId, typ)
+  result = ctx.newEnvVarAccess(s)
   inc ctx.tempVarId
 
 proc hasYields(n: PNode): bool =
@@ -197,7 +309,17 @@ proc transformBreaksInBlock(ctx: var Ctx, n: PNode, label, after: PNode): PNode
     for i in 0 ..< n.len:
       n[i] = ctx.transformBreaksInBlock(n[i], label, after)
 
-proc collectExceptState(n: PNode): PNode =
+proc newNullifyCurExc(ctx: var Ctx): PNode =
+  # :curEcx = nil
+  result = newNode(nkAsgn)
+  let curExc = ctx.newCurExcAccess()
+  result.add(curExc)
+
+  let nilnode = newNode(nkNilLit)
+  nilnode.typ = curExc.typ
+  result.add(nilnode)
+
+proc collectExceptState(ctx: var Ctx, n: PNode): PNode =
   var ifStmt = newNode(nkIfStmt)
   for c in n:
     if c.kind == nkExceptBranch:
@@ -208,8 +330,10 @@ proc collectExceptState(n: PNode): PNode =
         assert(c.len == 2)
         ifBranch = newNode(nkElifBranch)
         let expression = newNodeI(nkCall, n.info)
+        expression.add(newSymNode(getSysMagic("of", mOf)))
         expression.add(callCodegenProc("getCurrentException", emptyNode))
         expression.add(c[0])
+        expression.typ = getSysType(tyBool)
         ifBranch.add(expression)
         branchBody = c[1]
       else:
@@ -226,10 +350,37 @@ proc collectExceptState(n: PNode): PNode =
 
   if ifStmt.len != 0:
     result = newNode(nkStmtList)
+    result.add(ctx.newNullifyCurExc())
     result.add(ifStmt)
   else:
     result = emptyNode
 
+proc addElseToExcept(ctx: var Ctx, n: PNode) =
+  if n.kind == nkStmtList and n[1].kind == nkIfStmt and n[1][^1].kind != nkElse:
+    # Not all cases are covered
+    let elseBranch = newNode(nkElse)
+    let branchBody = newNode(nkStmtList)
+
+    block: # :unrollFinally = true
+      let asgn = newNode(nkAsgn)
+      asgn.add(ctx.newUnrollFinallyAccess())
+      asgn.add(newIntTypeNode(nkIntLit, 1, getSysType(tyBool)))
+      branchBody.add(asgn)
+
+    block: # :curExc = getCurrentException()
+      let asgn = newNode(nkAsgn)
+      asgn.add(ctx.newCurExcAccess)
+      asgn.add(callCodegenProc("getCurrentException", emptyNode))
+      branchBody.add(asgn)
+
+    block: # goto nearestFinally
+      let goto = newNode(nkGotoState)
+      goto.add(newIntLit(ctx.nearestFinally))
+      branchBody.add(goto)
+
+    elseBranch.add(branchBody)
+    n[1].add(elseBranch)
+
 proc getFinallyNode(n: PNode): PNode =
   result = n[^1]
   if result.kind == nkFinally:
@@ -273,6 +424,101 @@ proc lowerStmtListExpr(ctx: var Ctx, n: PNode): PNode =
     for i in 0 ..< n.len:
       n[i] = ctx.lowerStmtListExpr(n[i])
 
+proc newEndFinallyNode(ctx: var Ctx): PNode =
+  # Generate the following code:
+  #   if :unrollFinally:
+  #     when nearestFinally == 0: # Pseudocode. The `when` is not emitted in reality
+  #       if :curExc.isNil:
+  #         return :tmpResult
+  #       else:
+  #         raise
+  #     else:
+  #       goto nearestFinally
+  #       :state = nearestFinally
+  #       break :stateLoop
+
+  result = newNode(nkIfStmt)
+
+  let elifBranch = newNode(nkElifBranch)
+  elifBranch.add(ctx.newUnrollFinallyAccess())
+  result.add(elifBranch)
+
+  var ifBody: PNode
+
+  if ctx.nearestFinally == 0 or true:
+    ifBody = newNode(nkIfStmt)
+    let branch = newNode(nkElifBranch)
+
+    let cmp = newNode(nkCall)
+    cmp.add(getSysMagic("==", mEqRef).newSymNode)
+    let curExc = ctx.newCurExcAccess()
+    let nilnode = newNode(nkNilLit)
+    nilnode.typ = curExc.typ
+    cmp.add(curExc)
+    cmp.add(nilnode)
+    cmp.typ = getSysType(tyBool)
+    branch.add(cmp)
+
+    var retStmt = newNode(nkReturnStmt)
+    if true:
+      var a = newNode(nkAsgn)
+      addSon(a, newSymNode(getClosureIterResult(ctx.fn)))
+      addSon(a, ctx.newTmpResultAccess())
+      retStmt.add(a)
+    else:
+      retStmt.add(emptyNode)
+    branch.add(retStmt)
+
+    let elseBranch = newNode(nkElse)
+    let raiseStmt = newNode(nkRaiseStmt)
+
+    # The C++ backend requires `getCurrentException` here.
+    raiseStmt.add(callCodegenProc("getCurrentException", emptyNode))
+    elseBranch.add(raiseStmt)
+
+    ifBody.add(branch)
+    ifBody.add(elseBranch)
+  else:
+    ifBody = newNode(nkGotoState)
+    ifBody.add(newIntLit(ctx.nearestFinally))
+
+  elifBranch.add(ifBody)
+
+proc transformReturnsInTry(ctx: var Ctx, n: PNode): PNode =
+  result = n
+  # TODO: This is very inefficient. It traverses the node, looking for nkYieldStmt.
+  case n.kind
+  of nkReturnStmt:
+    # We're somewhere in try, transform to finally unrolling
+    assert(ctx.nearestFinally != 0)
+
+    result = newNodeI(nkStmtList, n.info)
+
+    block: # :unrollFinally = true
+      let asgn = newNodeI(nkAsgn, n.info)
+      asgn.add(ctx.newUnrollFinallyAccess())
+      asgn.add(newIntTypeNode(nkIntLit, 1, getSysType(tyBool)))
+      result.add(asgn)
+
+    if n[0].kind != nkEmpty: # TODO: And not void!
+      let asgnTmpResult = newNodeI(nkAsgn, n.info)
+      asgnTmpResult.add(ctx.newTmpResultAccess())
+      asgnTmpResult.add(n[0])
+      result.add(asgnTmpResult)
+
+    result.add(ctx.newNullifyCurExc())
+
+    let goto = newNodeI(nkGotoState, n.info)
+    goto.add(newIntLit(ctx.nearestFinally))
+    result.add(goto)
+
+  of nkCharLit..nkUInt64Lit, nkFloatLit..nkFloat128Lit, nkStrLit..nkTripleStrLit,
+      nkSym, nkIdent, procDefs, nkTemplateDef:
+    discard
+  else:
+    for i in 0 ..< n.len:
+      n[i] = ctx.transformReturnsInTry(n[i])
+
 proc transformClosureIteratorBody(ctx: var Ctx, n: PNode, gotoOut: PNode): PNode =
   result = n
   case n.kind:
@@ -310,7 +556,6 @@ proc transformClosureIteratorBody(ctx: var Ctx, n: PNode, gotoOut: PNode): PNode
       assert(false, "nkStmtListExpr not lowered")
 
     of nkYieldStmt:
-      # echo "YIELD!"
       result = newNodeI(nkStmtList, n.info)
       result.add(n)
       result.add(gotoOut)
@@ -371,34 +616,64 @@ proc transformClosureIteratorBody(ctx: var Ctx, n: PNode, gotoOut: PNode): PNode
       result[1] = ctx.transformClosureIteratorBody(result[1], gotoOut)
 
     of nkTryStmt:
+      # See explanation above about how this works
+      ctx.hasExceptions = true
+
+      result = newNode(nkGotoState)
       var tryBody = toStmtList(n[0])
-
-      # let popTry = newNode(nkPar)
-      # popTry.add(newIdentNode(getIdent("popTry"), n.info))
+      var exceptBody = ctx.collectExceptState(n)
       var finallyBody = newNode(nkStmtList)
-      # finallyBody.add(popTry)
       finallyBody.add(getFinallyNode(n))
+      finallyBody = ctx.transformReturnsInTry(finallyBody)
+      finallyBody.add(ctx.newEndFinallyNode())
 
-      var tryCatchOut = newNode(nkGotoState)
-
-      tryBody = ctx.transformClosureIteratorBody(tryBody, tryCatchOut)
-      var exceptBody = collectExceptState(n)
-
-      var exceptIdx = -1
+      # The following index calculation is based on the knowledge how state
+      # indexes are assigned
+      let tryIdx = ctx.states.len
+      var exceptIdx, finallyIdx: int
       if exceptBody.kind != nkEmpty:
-        exceptBody = ctx.transformClosureIteratorBody(exceptBody, tryCatchOut)
-        exceptIdx = ctx.newState(exceptBody, nil)
+        exceptIdx = -(tryIdx + 1)
+        finallyIdx = tryIdx + 2
+      else:
+        exceptIdx = tryIdx + 1
+        finallyIdx = tryIdx + 1
 
-      finallyBody = ctx.transformClosureIteratorBody(finallyBody, gotoOut)
-      let finallyIdx = ctx.newState(finallyBody, tryCatchOut)
+      let outToFinally = newNode(nkGotoState)
 
-      # let pushTry = newNode(nkPar) #newCall(newSym("pushTry"), newIntLit(exceptIdx))
-      # pushTry.add(newIdentNode(getIdent("pushTry"), n.info))
-      # pushTry.add(newIntLit(exceptIdx))
-      # pushTry.add(newIntLit(finallyIdx))
-      # tryBody.sons.insert(pushTry, 0)
+      block: # Create initial states.
+        let oldExcHandlingState = ctx.curExcHandlingState
+        ctx.curExcHandlingState = exceptIdx
+        let realTryIdx = ctx.newState(tryBody, result)
+        assert(realTryIdx == tryIdx)
 
-      result = tryBody
+        if exceptBody.kind != nkEmpty:
+          ctx.curExcHandlingState = finallyIdx
+          let realExceptIdx = ctx.newState(exceptBody, nil)
+          assert(realExceptIdx == -exceptIdx)
+
+        ctx.curExcHandlingState = oldExcHandlingState
+        let realFinallyIdx = ctx.newState(finallyBody, outToFinally)
+        assert(realFinallyIdx == finallyIdx)
+
+      block: # Subdivide the states
+        let oldNearestFinally = ctx.nearestFinally
+        ctx.nearestFinally = finallyIdx
+
+        let oldExcHandlingState = ctx.curExcHandlingState
+
+        ctx.curExcHandlingState = exceptIdx
+
+        discard ctx.transformReturnsInTry(tryBody)
+        discard ctx.transformClosureIteratorBody(tryBody, outToFinally)
+
+        ctx.curExcHandlingState = finallyIdx
+        ctx.addElseToExcept(exceptBody)
+        discard ctx.transformReturnsInTry(exceptBody)
+        discard ctx.transformClosureIteratorBody(exceptBody, outToFinally)
+
+        ctx.curExcHandlingState = oldExcHandlingState
+        ctx.nearestFinally = oldNearestFinally
+        discard ctx.transformClosureIteratorBody(finallyBody, gotoOut)
 
     of nkGotoState, nkForStmt:
       internalError("closure iter " & $n.kind)
@@ -460,7 +735,6 @@ proc tranformStateAssignments(ctx: var Ctx, n: PNode): PNode =
     discard
 
   of nkReturnStmt:
-
     result = newNodeI(nkStmtList, n.info)
     result.add(ctx.newStateAssgn(-1))
     result.add(n)
@@ -483,6 +757,29 @@ proc skipStmtList(n: PNode): PNode =
     if result.len == 0: return emptyNode
     result = result[0]
 
+proc skipEmptyStates(ctx: Ctx, stateIdx: int): int =
+  # Returns first non-empty state idx for `stateIdx`. Returns `stateIdx` if
+  # it is not empty
+  var maxJumps = ctx.states.len # maxJumps used only for debugging purposes.
+  var stateIdx = stateIdx
+  while true:
+    let label = stateIdx
+    if label == ctx.exitStateIdx: break
+    var newLabel = label
+    if label == -1:
+      newLabel = ctx.exitStateIdx
+    else:
+      let fs = ctx.states[label][1].skipStmtList()
+      if fs.kind == nkGotoState:
+        newLabel = fs[0].intVal.int
+    if label == newLabel: break
+    stateIdx = newLabel
+    dec maxJumps
+    if maxJumps == 0:
+      assert(false, "Internal error")
+
+  result = ctx.states[stateIdx][0].intVal.int
+
 proc skipThroughEmptyStates(ctx: var Ctx, n: PNode): PNode =
   result = n
   case n.kind
@@ -490,31 +787,143 @@ proc skipThroughEmptyStates(ctx: var Ctx, n: PNode): PNode =
       nkSym, nkIdent, procDefs, nkTemplateDef:
     discard
   of nkGotoState:
-    var maxJumps = ctx.states.len # maxJumps used only for debugging purposes.
     result = copyTree(n)
-    while true:
-      let label = result[0].intVal.int
-      if label == ctx.exitStateIdx: break
-      var newLabel = label
-      if label == -1:
-        newLabel = ctx.exitStateIdx
-      else:
-        let fs = ctx.states[label][1].skipStmtList()
-        if fs.kind == nkGotoState:
-          newLabel = fs[0].intVal.int
-      if label == newLabel: break
-      result[0].intVal = newLabel
-      dec maxJumps
-      if maxJumps == 0:
-        assert(false, "Internal error")
-
-    let label = result[0].intVal.int
-    result[0].intVal = ctx.states[label][0].intVal
+    result[0].intVal = ctx.skipEmptyStates(result[0].intVal.int)
   else:
     for i in 0 ..< n.len:
       n[i] = ctx.skipThroughEmptyStates(n[i])
 
+proc newArrayType(n: int, t: PType, owner: PSym): PType =
+  result = newType(tyArray, owner)
+
+  let rng = newType(tyRange, owner)
+  rng.n = newNode(nkRange)
+  rng.n.add(newIntLit(0))
+  rng.n.add(newIntLit(n))
+  rng.rawAddSon(t)
+
+  result.rawAddSon(rng)
+  result.rawAddSon(t)
+
+proc createExceptionTable(ctx: var Ctx): PNode =
+  result = newNode(nkBracket)
+  result.typ = newArrayType(ctx.exceptionTable.len, getSysType(tyInt16), ctx.fn)
+
+  for i in ctx.exceptionTable:
+    let elem = newIntNode(nkIntLit, i)
+    elem.typ = getSysType(tyInt16)
+    result.add(elem)
+
+proc newCatchBody(ctx: var Ctx): PNode {.inline.} =
+  # Generates the code:
+  # :state = exceptionTable[:state]
+  # if :state == 0: raise
+  # :unrollFinally = :state > 0
+  # if :state < 0:
+  #   :state = -:state
+  # :curExc = getCurrentException()
+
+  result = newNode(nkStmtList)
+
+  # :state = exceptionTable[:state]
+  block:
+
+    # exceptionTable[:state]
+    let getNextState = newNode(nkBracketExpr)
+    getNextState.add(ctx.createExceptionTable)
+    getNextState.add(ctx.newStateAccess())
+    getNextState.typ = getSysType(tyInt)
+
+    # :state = exceptionTable[:state]
+    result.add(ctx.newStateAssgn(getNextState))
+
+  # if :state == 0: raise
+  block:
+    let ifStmt = newNode(nkIfStmt)
+    let ifBranch = newNode(nkElifBranch)
+    let cond = newNode(nkCall)
+    cond.add(getSysMagic("==", mEqI).newSymNode)
+    cond.add(ctx.newStateAccess())
+    cond.add(newIntTypeNode(nkIntLit, 0, getSysType(tyInt)))
+    cond.typ = getSysType(tyBool)
+    ifBranch.add(cond)
+
+    let raiseStmt = newNode(nkRaiseStmt)
+    raiseStmt.add(emptyNode)
+
+    ifBranch.add(raiseStmt)
+    ifStmt.add(ifBranch)
+    result.add(ifStmt)
+
+  # :unrollFinally = :state > 0
+  block:
+    let asgn = newNode(nkAsgn)
+    asgn.add(ctx.newUnrollFinallyAccess())
+
+    let cond = newNode(nkCall)
+    cond.add(getSysMagic("<", mLtI).newSymNode)
+    cond.add(newIntTypeNode(nkIntLit, 0, getSysType(tyInt)))
+    cond.add(ctx.newStateAccess())
+    cond.typ = getSysType(tyBool)
+    asgn.add(cond)
+    result.add(asgn)
+
+  # if :state < 0: :state = -:state
+  block:
+    let ifStmt = newNode(nkIfStmt)
+    let ifBranch = newNode(nkElifBranch)
+    let cond = newNode(nkCall)
+    cond.add(getSysMagic("<", mLtI).newSymNode)
+    cond.add(ctx.newStateAccess())
+    cond.add(newIntTypeNode(nkIntLit, 0, getSysType(tyInt)))
+    cond.typ = getSysType(tyBool)
+    ifBranch.add(cond)
+
+    let negateState = newNode(nkCall)
+    negateState.add(getSysMagic("-", mUnaryMinusI).newSymNode)
+    negateState.add(ctx.newStateAccess())
+    negateState.typ = getSysType(tyInt)
+
+    ifBranch.add(ctx.newStateAssgn(negateState))
+    ifStmt.add(ifBranch)
+    result.add(ifStmt)
+
+  # :curExc = getCurrentException()
+  block:
+    let getCurExc = callCodegenProc("getCurrentException", emptyNode)
+    let asgn = newNode(nkAsgn)
+    asgn.add(ctx.newCurExcAccess())
+    asgn.add(getCurExc)
+    result.add(asgn)
+
+proc wrapIntoTryExcept(ctx: var Ctx, n: PNode): PNode =
+  result = newNode(nkTryStmt)
+
+  let tryBody = newNode(nkStmtList)
+
+  let setupExc = newNode(nkCall)
+  setupExc.add(newSymNode(getCompilerProc("closureIterSetupExc")))
+
+  tryBody.add(setupExc)
+
+  tryBody.add(n)
+  result.add(tryBody)
+
+  let catchNode = newNode(nkExceptBranch)
+  result.add(catchNode)
+
+  let catchBody = newNode(nkStmtList)
+  catchBody.add(ctx.newCatchBody())
+  catchNode.add(catchBody)
+
+  setupExc.add(ctx.newCurExcAccess())
+
 proc wrapIntoStateLoop(ctx: var Ctx, n: PNode): PNode =
+  # while true:
+  #   block :stateLoop:
+  #     gotoState :state
+  #     body # Might get wrapped in try-except
+
   result = newNode(nkWhileStmt)
   result.add(newSymNode(getSysSym("true")))
 
@@ -532,19 +941,19 @@ proc wrapIntoStateLoop(ctx: var Ctx, n: PNode): PNode =
   let blockStmt = newNodeI(nkBlockStmt, n.info)
   blockStmt.add(newSymNode(ctx.stateLoopLabel))
 
-  let blockBody = newNodeI(nkStmtList, n.info)
-  blockStmt.add(blockBody)
+  var blockBody = newNodeI(nkStmtList, n.info)
 
   let gs = newNodeI(nkGotoState, n.info)
-  if ctx.stateVarSym.isNil:
-    gs.add(rawIndirectAccess(newSymNode(getEnvParam(ctx.fn)), getStateField(ctx.fn), n.info))
-  else:
-    gs.add(newSymNode(ctx.stateVarSym))
-
+  gs.add(ctx.newStateAccess())
   gs.add(newIntLit(ctx.states.len - 1))
+
   blockBody.add(gs)
   blockBody.add(n)
-  # gs.add(rawIndirectAccess(newSymNode(ctx.fn.getHiddenParam), getStateField(ctx.fn), n.info))
+
+  if ctx.hasExceptions:
+    blockBody = ctx.wrapIntoTryExcept(blockBody)
+
+  blockStmt.add(blockBody)
 
   loopBody.add(blockStmt)
 
@@ -558,7 +967,7 @@ proc deleteEmptyStates(ctx: var Ctx) =
   var iValid = 0
   for i, s in ctx.states:
     let body = s[1].skipStmtList()
-    if body.kind == nkGotoState and i != ctx.states.len - 1:
+    if body.kind == nkGotoState and i != ctx.states.len - 1 and i != 0:
       # This is an empty state. Mark with -1.
       s[0].intVal = -1
     else:
@@ -567,14 +976,20 @@ proc deleteEmptyStates(ctx: var Ctx) =
 
   for i, s in ctx.states:
     let body = s[1].skipStmtList()
-    if body.kind != nkGotoState:
+    if body.kind != nkGotoState or i == 0:
       discard ctx.skipThroughEmptyStates(s)
+      let excHandlState = ctx.exceptionTable[i]
+      if excHandlState < 0:
+        ctx.exceptionTable[i] = -ctx.skipEmptyStates(-excHandlState)
+      elif excHandlState != 0:
+        ctx.exceptionTable[i] = ctx.skipEmptyStates(excHandlState)
 
   var i = 0
   while i < ctx.states.len - 1:
     let fs = ctx.states[i][1].skipStmtList()
-    if fs.kind == nkGotoState:
+    if fs.kind == nkGotoState and i != 0:
       ctx.states.delete(i)
+      ctx.exceptionTable.delete(i)
     else:
       inc i
 
@@ -591,6 +1006,7 @@ proc transformClosureIterator*(fn: PSym, n: PNode): PNode =
 
   ctx.states = @[]
   ctx.stateLoopLabel = newSym(skLabel, getIdent(":stateLoop"), fn, fn.info)
+  ctx.exceptionTable = @[]
   let n = n.toStmtList
 
   discard ctx.newState(n, nil)
@@ -613,19 +1029,11 @@ proc transformClosureIterator*(fn: PSym, n: PNode): PNode =
     result.add(body)
 
   result = ctx.tranformStateAssignments(result)
-
-  # Add excpetion handling
-  var hasExceptions = false
-  if hasExceptions:
-    discard # TODO:
-    # result = wrapIntoTryCatch(result)
-
-  # while true:
-  #   block :stateLoop:
-  #     gotoState
-  #     body
   result = ctx.wrapIntoStateLoop(result)
 
-  # echo "TRANSFORM TO STATES2: "
-  # debug(result)
+  # echo "TRANSFORM TO STATES: "
   # echo renderTree(result)
+
+  # echo "exception table:"
+  # for i, e in ctx.exceptionTable:
+  #   echo i, " -> ", e
diff --git a/compiler/semexprs.nim b/compiler/semexprs.nim
index 1ef284a779..79010bfde9 100644
--- a/compiler/semexprs.nim
+++ b/compiler/semexprs.nim
@@ -1544,7 +1544,7 @@ proc semYield(c: PContext, n: PNode): PNode =
   checkSonsLen(n, 1)
   if c.p.owner == nil or c.p.owner.kind != skIterator:
     localError(n.info, errYieldNotAllowedHere)
-  elif c.p.inTryStmt > 0 and c.p.owner.typ.callConv != ccInline:
+  elif oldIterTransf in c.features and c.p.inTryStmt > 0 and c.p.owner.typ.callConv != ccInline:
     localError(n.info, errYieldNotAllowedInTryStmt)
   elif n.sons[0].kind != nkEmpty:
     n.sons[0] = semExprWithType(c, n.sons[0]) # check for type compatibility:
diff --git a/lib/system/embedded.nim b/lib/system/embedded.nim
index 46e84e0569..4d453fcca4 100644
--- a/lib/system/embedded.nim
+++ b/lib/system/embedded.nim
@@ -41,3 +41,6 @@ proc reraiseException() {.compilerRtl.} =
 proc writeStackTrace() = discard
 
 proc setControlCHook(hook: proc () {.noconv.}) = discard
+
+proc closureIterSetupExc(e: ref Exception) {.compilerproc, inline.} =
+  sysFatal(ReraiseError, "exception handling is not available")
diff --git a/lib/system/excpt.nim b/lib/system/excpt.nim
index fb38948f7b..dabfe010ea 100644
--- a/lib/system/excpt.nim
+++ b/lib/system/excpt.nim
@@ -131,6 +131,10 @@ proc popCurrentExceptionEx(id: uint) {.compilerRtl.} =
       quitOrDebug()
     prev.up = cur.up
 
+proc closureIterSetupExc(e: ref Exception) {.compilerproc, inline.} =
+  if not e.isNil:
+    currException = e
+
 # some platforms have native support for stack traces:
 const
   nativeStackTraceSupported* = (defined(macosx) or defined(linux)) and
diff --git a/tests/async/tasynctry2.nim b/tests/async/tasynctry2.nim
index 444a058beb..f82b6cfe0d 100644
--- a/tests/async/tasynctry2.nim
+++ b/tests/async/tasynctry2.nim
@@ -1,10 +1,12 @@
 discard """
   file: "tasynctry2.nim"
   errormsg: "\'yield\' cannot be used within \'try\' in a non-inlined iterator"
-  line: 15
+  line: 17
 """
 import asyncdispatch
 
+{.experimental: "oldIterTransf".}
+
 proc foo(): Future[bool] {.async.} = discard
 
 proc test5(): Future[int] {.async.} =
diff --git a/tests/iter/tyieldintry.nim b/tests/iter/tyieldintry.nim
new file mode 100644
index 0000000000..9cb199c5b8
--- /dev/null
+++ b/tests/iter/tyieldintry.nim
@@ -0,0 +1,201 @@
+discard """
+targets: "c cpp"
+output: "ok"
+"""
+var closureIterResult = newSeq[int]()
+
+proc checkpoint(arg: int) =
+    closureIterResult.add(arg)
+
+type
+    TestException = object of Exception
+    AnotherException = object of Exception
+
+proc testClosureIterAux(it: iterator(): int, exceptionExpected: bool, expectedResults: varargs[int]) =
+    closureIterResult.setLen(0)
+
+    var exceptionCaught = false
+
+    try:
+        for i in it():
+            closureIterResult.add(i)
+    except TestException:
+        exceptionCaught = true
+
+    if closureIterResult != @expectedResults or exceptionCaught != exceptionExpected:
+        if closureIterResult != @expectedResults:
+            echo "Expected: ", @expectedResults
+            echo "Actual: ", closureIterResult
+        if exceptionCaught != exceptionExpected:
+            echo "Expected exception: ", exceptionExpected
+            echo "Got exception: ", exceptionCaught
+        doAssert(false)
+
+proc test(it: iterator(): int, expectedResults: varargs[int]) =
+    testClosureIterAux(it, false, expectedResults)
+
+proc testExc(it: iterator(): int, expectedResults: varargs[int]) =
+    testClosureIterAux(it, true, expectedResults)
+
+proc raiseException() =
+    raise newException(TestException, "Test exception!")
+
+block:
+    iterator it(): int {.closure.} =
+        var i = 5
+        while i != 0:
+            yield i
+            if i == 3:
+                yield 123
+            dec i
+
+    test(it, 5, 4, 3, 123, 2, 1)
+
+block:
+    iterator it(): int {.closure.} =
+        yield 0
+        try:
+            checkpoint(1)
+            raiseException()
+        except TestException:
+            checkpoint(2)
+            yield 3
+            checkpoint(4)
+        finally:
+            checkpoint(5)
+
+        checkpoint(6)
+
+    test(it, 0, 1, 2, 3, 4, 5, 6)
+
+block:
+    iterator it(): int {.closure.} =
+        yield 0
+        try:
+            yield 1
+            checkpoint(2)
+        finally:
+            checkpoint(3)
+            yield 4
+            checkpoint(5)
+            yield 6
+
+    test(it, 0, 1, 2, 3, 4, 5, 6)
+
+block:
+    iterator it(): int {.closure.} =
+        yield 0
+        try:
+            yield 1
+            raiseException()
+            yield 2
+        finally:
+            checkpoint(3)
+            yield 4
+            checkpoint(5)
+            yield 6
+
+    testExc(it, 0, 1, 3, 4, 5, 6)
+
+block:
+    iterator it(): int {.closure.} =
+        try:
+            try:
+                raiseException()
+            except AnotherException:
+                yield 123
+            finally:
+                checkpoint(3)
+        finally:
+            checkpoint(4)
+
+    testExc(it, 3, 4)
+
+block:
+    iterator it(): int {.closure.} =
+        try:
+            yield 1
+            raiseException()
+        except AnotherException:
+            checkpoint(123)
+        finally:
+            checkpoint(2)
+        checkpoint(3)
+
+    testExc(it, 1, 2)
+
+block:
+    iterator it(): int {.closure.} =
+        try:
+            yield 0
+            try:
+                yield 1
+                try:
+                    yield 2
+                    raiseException()
+                except AnotherException:
+                    yield 123
+                finally:
+                    yield 3
+            except AnotherException:
+                yield 124
+            finally:
+                yield 4
+            checkpoint(1234)
+        except:
+            yield 5
+            checkpoint(6)
+        finally:
+            checkpoint(7)
+            yield 8
+        checkpoint(9)
+
+    test(it, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9)
+
+block:
+    iterator it(): int {.closure.} =
+        try:
+            yield 0
+            return 2
+        finally:
+            checkpoint(1)
+        checkpoint(123)
+
+    test(it, 0, 1)
+
+block:
+    iterator it(): int {.closure.} =
+        try:
+            try:
+                yield 0
+                raiseException()
+            finally:
+                checkpoint(1)
+        except TestException:
+            yield 2
+            return
+        finally:
+            yield 3
+
+        checkpoint(123)
+
+    test(it, 0, 1, 2, 3)
+
+block:
+    iterator it(): int {.closure.} =
+        try:
+            try:
+                yield 0
+                raiseException()
+            finally:
+                return # Return in finally should stop exception propagation
+        except AnotherException:
+            yield 2
+            return
+        finally:
+            yield 3
+        checkpoint(123)
+
+    test(it, 0, 3)
+
+echo "ok"

From ac86b8ce615cbd55074dfd27f42ed0368d84b1fd Mon Sep 17 00:00:00 2001
From: Yuriy Glukhov <yuriy.glukhov@gmail.com>
Date: Fri, 4 May 2018 16:39:59 +0300
Subject: [PATCH 07/14] Cleanup

---
 compiler/closureiters.nim | 72 +++++++++++++++------------------------
 1 file changed, 27 insertions(+), 45 deletions(-)

diff --git a/compiler/closureiters.nim b/compiler/closureiters.nim
index 504f70347d..a8e7e02747 100644
--- a/compiler/closureiters.nim
+++ b/compiler/closureiters.nim
@@ -123,14 +123,10 @@
 # STATE2: # Finally
 #   yield 2
 #   if :unrollFinally: # This node is created by `newEndFinallyNode`
-#     when nearestFinally == 0: # Pseudocode. The `when` is not emitted in reality
-#       if :curExc.isNil:
-#         return :tmpResult
-#       else:
-#         raise
+#     if :curExc.isNil:
+#       return :tmpResult
 #     else:
-#       :state = nearestFinally
-#       break :stateLoop
+#       raise
 #   state = -1 # Goto next state. In this case we just exit
 #   break :stateLoop
 
@@ -427,15 +423,10 @@ proc lowerStmtListExpr(ctx: var Ctx, n: PNode): PNode =
 proc newEndFinallyNode(ctx: var Ctx): PNode =
   # Generate the following code:
   #   if :unrollFinally:
-  #     when nearestFinally == 0: # Pseudocode. The `when` is not emitted in reality
   #       if :curExc.isNil:
   #         return :tmpResult
   #       else:
   #         raise
-  #     else:
-  #       goto nearestFinally
-  #       :state = nearestFinally
-  #       break :stateLoop
 
   result = newNode(nkIfStmt)
 
@@ -443,44 +434,35 @@ proc newEndFinallyNode(ctx: var Ctx): PNode =
   elifBranch.add(ctx.newUnrollFinallyAccess())
   result.add(elifBranch)
 
-  var ifBody: PNode
+  let ifBody = newNode(nkIfStmt)
+  let branch = newNode(nkElifBranch)
 
-  if ctx.nearestFinally == 0 or true:
-    ifBody = newNode(nkIfStmt)
-    let branch = newNode(nkElifBranch)
+  let cmp = newNode(nkCall)
+  cmp.add(getSysMagic("==", mEqRef).newSymNode)
+  let curExc = ctx.newCurExcAccess()
+  let nilnode = newNode(nkNilLit)
+  nilnode.typ = curExc.typ
+  cmp.add(curExc)
+  cmp.add(nilnode)
+  cmp.typ = getSysType(tyBool)
+  branch.add(cmp)
 
-    let cmp = newNode(nkCall)
-    cmp.add(getSysMagic("==", mEqRef).newSymNode)
-    let curExc = ctx.newCurExcAccess()
-    let nilnode = newNode(nkNilLit)
-    nilnode.typ = curExc.typ
-    cmp.add(curExc)
-    cmp.add(nilnode)
-    cmp.typ = getSysType(tyBool)
-    branch.add(cmp)
+  let retStmt = newNode(nkReturnStmt)
+  let asgn = newNode(nkAsgn)
+  addSon(asgn, newSymNode(getClosureIterResult(ctx.fn)))
+  addSon(asgn, ctx.newTmpResultAccess())
+  retStmt.add(asgn)
+  branch.add(retStmt)
 
-    var retStmt = newNode(nkReturnStmt)
-    if true:
-      var a = newNode(nkAsgn)
-      addSon(a, newSymNode(getClosureIterResult(ctx.fn)))
-      addSon(a, ctx.newTmpResultAccess())
-      retStmt.add(a)
-    else:
-      retStmt.add(emptyNode)
-    branch.add(retStmt)
+  let elseBranch = newNode(nkElse)
+  let raiseStmt = newNode(nkRaiseStmt)
 
-    let elseBranch = newNode(nkElse)
-    let raiseStmt = newNode(nkRaiseStmt)
+  # The C++ backend requires `getCurrentException` here.
+  raiseStmt.add(callCodegenProc("getCurrentException", emptyNode))
+  elseBranch.add(raiseStmt)
 
-    # The C++ backend requires `getCurrentException` here.
-    raiseStmt.add(callCodegenProc("getCurrentException", emptyNode))
-    elseBranch.add(raiseStmt)
-
-    ifBody.add(branch)
-    ifBody.add(elseBranch)
-  else:
-    ifBody = newNode(nkGotoState)
-    ifBody.add(newIntLit(ctx.nearestFinally))
+  ifBody.add(branch)
+  ifBody.add(elseBranch)
 
   elifBranch.add(ifBody)
 

From 14ca79fe1f1fafb8e3aff2e4c27bcb94c0595792 Mon Sep 17 00:00:00 2001
From: Yuriy Glukhov <yuriy.glukhov@gmail.com>
Date: Tue, 8 May 2018 01:25:08 +0300
Subject: [PATCH 08/14] More elaborate nkStmtListExpr lowering

---
 compiler/closureiters.nim  | 404 ++++++++++++++++++++++++++++---
 tests/iter/tyieldintry.nim | 477 +++++++++++++++++++++++++------------
 2 files changed, 696 insertions(+), 185 deletions(-)

diff --git a/compiler/closureiters.nim b/compiler/closureiters.nim
index a8e7e02747..a30b4e10ee 100644
--- a/compiler/closureiters.nim
+++ b/compiler/closureiters.nim
@@ -149,7 +149,6 @@ type
     exitStateIdx: int # index of the last state
     tempVarId: int # unique name counter
     tempVars: PNode # Temp var decls, nkVarSection
-    loweredStmtListExpr: PNode # Temporary used for nkStmtListExpr lowering
     exceptionTable: seq[int] # For state `i` jump to state `exceptionTable[i]` if exception is raised
     hasExceptions: bool # Does closure have yield in try?
     curExcHandlingState: int # Negative for except, positive for finally
@@ -177,6 +176,7 @@ proc newStateAssgn(ctx: var Ctx, stateNo: int = -2): PNode =
 proc newEnvVar(ctx: var Ctx, name: string, typ: PType): PSym =
   result = newSym(skVar, getIdent(name), ctx.fn, ctx.fn.info)
   result.typ = typ
+  assert(not typ.isNil)
 
   if not ctx.stateVarSym.isNil:
     # We haven't gone through labmda lifting yet, so just create a local var,
@@ -197,7 +197,6 @@ proc newEnvVarAccess(ctx: Ctx, s: PSym): PNode =
 
 proc newTmpResultAccess(ctx: var Ctx): PNode =
   if ctx.tmpResultSym.isNil:
-    debug(ctx.fn.typ)
     ctx.tmpResultSym = ctx.newEnvVar(":tmpResult", ctx.fn.typ[0])
   ctx.newEnvVarAccess(ctx.tmpResultSym)
 
@@ -245,9 +244,8 @@ proc addGotoOut(n: PNode, gotoOut: PNode): PNode =
   if result.len != 0 and result.sons[^1].kind != nkGotoState:
     result.add(gotoOut)
 
-proc newTempVarAccess(ctx: var Ctx, typ: PType, i: TLineInfo): PNode =
-  let s = ctx.newEnvVar(":tmpSlLower" & $ctx.tempVarId, typ)
-  result = ctx.newEnvVarAccess(s)
+proc newTempVar(ctx: var Ctx, typ: PType): PSym =
+  result = ctx.newEnvVar(":tmpSlLower" & $ctx.tempVarId, typ)
   inc ctx.tempVarId
 
 proc hasYields(n: PNode): bool =
@@ -390,35 +388,382 @@ proc hasYieldsInExpressions(n: PNode): bool =
       nkSym, nkIdent, procDefs, nkTemplateDef:
     discard
   of nkStmtListExpr:
-    result = n.hasYields
-  of nkStmtList, nkWhileStmt, nkCaseStmt, nkIfStmt:
-    discard
+    if isEmptyType(n.typ):
+      for c in n:
+        if c.hasYieldsInExpressions:
+          return true
+    else:
+      result = n.hasYields
   else:
     for c in n:
       if c.hasYieldsInExpressions:
         return true
 
-proc lowerStmtListExpr(ctx: var Ctx, n: PNode): PNode =
+proc exprToStmtList(n: PNode): tuple[s, res: PNode] =
+  assert(n.kind == nkStmtListExpr)
+
+  var parent = n
+  var lastSon = n[^1]
+
+  while lastSon.kind == nkStmtListExpr:
+    parent = lastSon
+    lastSon = lastSon[^1]
+
+  result.s = newNodeI(nkStmtList, n.info)
+  result.s.sons = parent.sons
+  result.s.sons.setLen(result.s.sons.len - 1) # delete last son
+  result.res = lastSon
+
+proc newEnvVarAsgn(ctx: Ctx, s: PSym, v: PNode): PNode =
+  result = newNode(nkFastAsgn)
+  result.add(ctx.newEnvVarAccess(s))
+  result.add(v)
+
+proc addExprAssgn(ctx: Ctx, output, input: PNode, sym: PSym) =
+  if input.kind == nkStmtListExpr:
+    let (st, res) = exprToStmtList(input)
+    output.add(st)
+    output.add(ctx.newEnvVarAsgn(sym, res))
+  else:
+    output.add(ctx.newEnvVarAsgn(sym, input))
+
+proc convertExprBodyToAsgn(ctx: Ctx, exprBody: PNode, res: PSym): PNode =
+  result = newNode(nkStmtList)
+  ctx.addExprAssgn(result, exprBody, res)
+
+proc newNotCall(e: PNode): PNode =
+  result = newNode(nkCall)
+  result.add(newSymNode(getSysMagic("not", mNot)))
+  result.add(e)
+  result.typ = getSysType(tyBool)
+
+proc lowerStmtListExprs(ctx: var Ctx, n: PNode, needsSplit: var bool): PNode =
   result = n
   case n.kind
   of nkCharLit..nkUInt64Lit, nkFloatLit..nkFloat128Lit, nkStrLit..nkTripleStrLit,
       nkSym, nkIdent, procDefs, nkTemplateDef:
     discard
-  of nkStmtListExpr:
-    if n.hasYields:
-      for i in 0 .. n.len - 2:
-        ctx.loweredStmtListExpr.add(n[i])
 
-      let tv = ctx.newTempVarAccess(n.typ, n[^1].info)
-      let asgn = newNode(nkAsgn)
-      asgn.add(tv)
-      asgn.add(n[^1])
-      ctx.loweredStmtListExpr.add(asgn)
-      result = tv
+  of nkYieldStmt:
+    var ns = false
+    for i in 0 ..< n.len:
+      n[i] = ctx.lowerStmtListExprs(n[i], ns)
 
+    if ns:
+      assert(n[0].kind == nkStmtListExpr)
+      result = newNodeI(nkStmtList, n.info)
+      let (st, ex) = exprToStmtList(n[0])
+      result.add(st)
+      n[0] = ex
+      result.add(n)
+
+    needsSplit = true
+
+  of nkPar, nkObjConstr, nkTupleConstr, nkBracket, nkArgList:
+    var ns = false
+    for i in 0 ..< n.len:
+      n[i] = ctx.lowerStmtListExprs(n[i], ns)
+
+    if ns:
+      needsSplit = true
+
+      result = newNodeI(nkStmtListExpr, n.info)
+      if n.typ.isNil: internalError("lowerStmtListExprs: constr typ.isNil")
+      result.typ = n.typ
+
+      for i in 0 ..< n.len:
+        if n[i].kind == nkStmtListExpr:
+          let (st, ex) = exprToStmtList(n[i])
+          result.add(st)
+          n[i] = ex
+      result.add(n)
+
+  of nkIfStmt, nkIfExpr:
+    var ns = false
+    for i in 0 ..< n.len:
+      n[i] = ctx.lowerStmtListExprs(n[i], ns)
+
+    if ns:
+      needsSplit = true
+      var tmp: PSym
+      var s: PNode
+      let isExpr = not isEmptyType(n.typ)
+      if isExpr:
+        tmp = ctx.newTempVar(n.typ)
+        result = newNode(nkStmtListExpr)
+        result.typ = n.typ
+      else:
+        result = newNode(nkStmtList)
+
+      var curS = result
+
+      for branch in n:
+        case branch.kind
+        of nkElseExpr, nkElse:
+          if isExpr:
+            var newBranch = newNodeI(nkElse, branch.info)
+            let branchBody = newNode(nkStmtList)
+            ctx.addExprAssgn(branchBody, branch[0], tmp)
+            newBranch.add(branchBody)
+            curS.add(newBranch)
+          else:
+            curS.add(branch)
+
+        of nkElifExpr, nkElifBranch:
+          var newBranch: PNode
+          if branch[0].kind == nkStmtListExpr:
+            let elseBody = newNode(nkStmtList)
+
+            let (st, res) = exprToStmtList(branch[0])
+            elseBody.add(st)
+
+            newBranch = newNodeI(nkElifBranch, branch.info)
+            newBranch.add(res)
+            newBranch.add(branch[1])
+
+            let newIf = newNodeI(nkIfStmt, branch.info)
+            newIf.add(newBranch)
+            elseBody.add(newIf)
+            if curS.kind == nkIfStmt:
+              let newElse = newNodeI(nkElse, branch.info)
+              newElse.add(elseBody)
+              curS.add(newElse)
+            else:
+              curS.add(elseBody)
+            curS = newIf
+          else:
+            newBranch = branch
+            if curS.kind == nkIfStmt:
+              curS.add(newBranch)
+            else:
+              let newIf = newNodeI(nkIfStmt, branch.info)
+              newIf.add(newBranch)
+              curS.add(newIf)
+              curS = newIf
+
+          if isExpr:
+            let branchBody = newNode(nkStmtList)
+            ctx.addExprAssgn(branchBody, branch[1], tmp)
+            newBranch[1] = branchBody
+
+        else:
+          internalError("lowerStmtListExpr(nkIf): " & $branch.kind)
+
+      if isExpr: result.add(ctx.newEnvVarAccess(tmp))
+
+  of nkTryStmt:
+    var ns = false
+    for i in 0 ..< n.len:
+      n[i] = ctx.lowerStmtListExprs(n[i], ns)
+
+    if ns:
+      needsSplit = true
+      let isExpr = not isEmptyType(n.typ)
+
+      if isExpr:
+        result = newNodeI(nkStmtListExpr, n.info)
+        result.typ = n.typ
+        let tmp = ctx.newTempVar(n.typ)
+
+        n[0] = ctx.convertExprBodyToAsgn(n[0], tmp)
+        for i in 1 ..< n.len:
+          let branch = n[i]
+          case branch.kind
+          of nkExceptBranch:
+            if branch[0].kind == nkType:
+              branch[1] = ctx.convertExprBodyToAsgn(branch[1], tmp)
+            else:
+              branch[0] = ctx.convertExprBodyToAsgn(branch[0], tmp)
+          of nkFinally:
+            discard
+          else:
+            internalError("lowerStmtListExpr(nkTryStmt): " & $branch.kind)
+        result.add(n)
+        result.add(ctx.newEnvVarAccess(tmp))
+
+  of nkCaseStmt:
+    var ns = false
+    for i in 0 ..< n.len:
+      n[i] = ctx.lowerStmtListExprs(n[i], ns)
+
+    if ns:
+      needsSplit = true
+
+      let isExpr = not isEmptyType(n.typ)
+
+      if isExpr:
+        let tmp = ctx.newTempVar(n.typ)
+        result = newNodeI(nkStmtListExpr, n.info)
+        result.typ = n.typ
+
+        if n[0].kind == nkStmtListExpr:
+          let (st, ex) = exprToStmtList(n[0])
+          result.add(st)
+          n[0] = ex
+
+        for i in 1 ..< n.len:
+          let branch = n[i]
+          case branch.kind
+          of nkOfBranch:
+            branch[1] = ctx.convertExprBodyToAsgn(branch[1], tmp)
+          of nkElse:
+            branch[0] = ctx.convertExprBodyToAsgn(branch[0], tmp)
+          else:
+            internalError("lowerStmtListExpr(nkCaseStmt): " & $branch.kind)
+        result.add(n)
+        result.add(ctx.newEnvVarAccess(tmp))
+
+  of nkCallKinds:
+    var ns = false
+    for i in 0 ..< n.len:
+      n[i] = ctx.lowerStmtListExprs(n[i], ns)
+
+    if ns:
+      needsSplit = true
+      let isExpr = not isEmptyType(n.typ)
+
+      if isExpr:
+        result = newNodeI(nkStmtListExpr, n.info)
+        result.typ = n.typ
+      else:
+        result = newNode(nkStmtList, n.info)
+
+      if n[0].kind == nkSym and n[0].sym.magic in {mAnd, mOr}: # `and`/`or` short cirquiting
+        var cond = n[1]
+        if cond.kind == nkStmtListExpr:
+          let (st, ex) = exprToStmtList(cond)
+          result.add(st)
+          cond = ex
+
+        let tmp = ctx.newTempVar(cond.typ)
+        result.add(ctx.newEnvVarAsgn(tmp, cond))
+
+        let ifNode = newNode(nkIfStmt)
+        let ifBranch = newNode(nkElifBranch)
+
+        var check = ctx.newEnvVarAccess(tmp)
+        if n[0].sym.magic == mOr:
+          check = newNotCall(check)
+        ifBranch.add(check)
+
+        cond = n[2]
+        let ifBody = newNode(nkStmtList)
+        if cond.kind == nkStmtListExpr:
+          let (st, ex) = exprToStmtList(cond)
+          ifBody.add(st)
+          cond = ex
+        ifBody.add(ctx.newEnvVarAsgn(tmp, cond))
+        ifBranch.add(ifBody)
+        ifNode.add(ifBranch)
+        result.add(ifNode)
+        result.add(ctx.newEnvVarAccess(tmp))
+      else:
+        for i in 0 ..< n.len:
+          if n[i].kind == nkStmtListExpr:
+            let (st, ex) = exprToStmtList(n[i])
+            result.add(st)
+            n[i] = ex
+
+          if n[i].kind in nkCallKinds: # XXX: This should better be some sort of side effect tracking
+            let tmp = ctx.newTempVar(n[i].typ)
+            result.add(ctx.newEnvVarAsgn(tmp, n[i]))
+            n[i] = ctx.newEnvVarAccess(tmp)
+
+        result.add(n)
+
+  of nkVarSection, nkLetSection:
+    result = newNodeI(nkStmtList, n.info)
+    for c in n:
+      let varSect = newNodeI(n.kind, n.info)
+      varSect.add(c)
+      var ns = false
+      c[^1] = ctx.lowerStmtListExprs(c[^1], ns)
+      if ns:
+        needsSplit = true
+        assert(c[^1].kind == nkStmtListExpr)
+        let (st, ex) = exprToStmtList(c[^1])
+        result.add(st)
+        c[^1] = ex
+      result.add(varSect)
+
+  of nkDiscardStmt, nkReturnStmt, nkRaiseStmt:
+    var ns = false
+    for i in 0 ..< n.len:
+      n[i] = ctx.lowerStmtListExprs(n[i], ns)
+
+    if ns:
+      needsSplit = true
+      result = newNodeI(nkStmtList, n.info)
+      let (st, ex) = exprToStmtList(n[0])
+      result.add(st)
+      n[0] = ex
+      result.add(n)
+
+  of nkCast:
+    var ns = false
+    for i in 0 ..< n.len:
+      n[i] = ctx.lowerStmtListExprs(n[i], ns)
+
+    if ns:
+      needsSplit = true
+      result = newNodeI(nkStmtListExpr, n.info)
+      result.typ = n.typ
+      let (st, ex) = exprToStmtList(n[1])
+      result.add(st)
+      n[1] = ex
+      result.add(n)
+
+  of nkAsgn, nkFastAsgn:
+    var ns = false
+    for i in 0 ..< n.len:
+      n[i] = ctx.lowerStmtListExprs(n[i], ns)
+
+    if ns:
+      needsSplit = true
+      result = newNodeI(nkStmtList, n.info)
+      if n[0].kind == nkStmtListExpr:
+        let (st, ex) = exprToStmtList(n[0])
+        result.add(st)
+        n[0] = ex
+
+      if n[1].kind == nkStmtListExpr:
+        let (st, ex) = exprToStmtList(n[1])
+        result.add(st)
+        n[1] = ex
+
+      result.add(n)
+
+  of nkWhileStmt:
+    var ns = false
+
+    var condNeedsSplit = false
+    n[0] = ctx.lowerStmtListExprs(n[0], condNeedsSplit)
+    var bodyNeedsSplit = false
+    n[1] = ctx.lowerStmtListExprs(n[1], bodyNeedsSplit)
+
+    if condNeedsSplit or bodyNeedsSplit:
+      needsSplit = true
+
+      if condNeedsSplit:
+        let newBody = newNode(nkStmtList)
+
+        let (st, ex) = exprToStmtList(n[0])
+        newBody.add(st)
+        let check = newNode(nkIfStmt)
+        let branch = newNode(nkElifBranch)
+        branch.add(newNotCall(ex))
+        let brk = newNode(nkBreakStmt)
+        brk.add(emptyNode)
+        branch.add(brk)
+        check.add(branch)
+        newBody.add(check)
+        newBody.add(n[1])
+
+        n[0] = newSymNode(getSysSym("true"))
+        n[1] = newBody
   else:
     for i in 0 ..< n.len:
-      n[i] = ctx.lowerStmtListExpr(n[i])
+      n[i] = ctx.lowerStmtListExprs(n[i], needsSplit)
 
 proc newEndFinallyNode(ctx: var Ctx): PNode =
   # Generate the following code:
@@ -448,7 +793,7 @@ proc newEndFinallyNode(ctx: var Ctx): PNode =
   branch.add(cmp)
 
   let retStmt = newNode(nkReturnStmt)
-  let asgn = newNode(nkAsgn)
+  let asgn = newNode(nkFastAsgn)
   addSon(asgn, newSymNode(getClosureIterResult(ctx.fn)))
   addSon(asgn, ctx.newTmpResultAccess())
   retStmt.add(asgn)
@@ -482,7 +827,7 @@ proc transformReturnsInTry(ctx: var Ctx, n: PNode): PNode =
       asgn.add(newIntTypeNode(nkIntLit, 1, getSysType(tyBool)))
       result.add(asgn)
 
-    if n[0].kind != nkEmpty: # TODO: And not void!
+    if n[0].kind != nkEmpty:
       let asgnTmpResult = newNodeI(nkAsgn, n.info)
       asgnTmpResult.add(ctx.newTmpResultAccess())
       asgnTmpResult.add(n[0])
@@ -508,17 +853,15 @@ proc transformClosureIteratorBody(ctx: var Ctx, n: PNode, gotoOut: PNode): PNode
         nkSym, nkIdent, procDefs, nkTemplateDef:
       discard
 
-    of nkStmtList:
+    of nkStmtList, nkStmtListExpr:
+      assert(isEmptyType(n.typ), "nkStmtListExpr not lowered")
+
       result = addGotoOut(result, gotoOut)
       for i in 0 ..< n.len:
         if n[i].hasYieldsInExpressions:
           # Lower nkStmtListExpr nodes inside `n[i]` first
-          assert(ctx.loweredStmtListExpr.isNil)
-          ctx.loweredStmtListExpr = newNodeI(nkStmtList, n.info)
-          n[i] = ctx.lowerStmtListExpr(n[i])
-          ctx.loweredStmtListExpr.add(n[i])
-          n[i] = ctx.loweredStmtListExpr
-          ctx.loweredStmtListExpr = nil
+          var ns = false
+          n[i] = ctx.lowerStmtListExprs(n[i], ns)
 
         if n[i].hasYields:
           # Create a new split
@@ -534,9 +877,6 @@ proc transformClosureIteratorBody(ctx: var Ctx, n: PNode, gotoOut: PNode): PNode
           discard ctx.transformClosureIteratorBody(s, gotoOut)
           break
 
-    of nkStmtListExpr:
-      assert(false, "nkStmtListExpr not lowered")
-
     of nkYieldStmt:
       result = newNodeI(nkStmtList, n.info)
       result.add(n)
diff --git a/tests/iter/tyieldintry.nim b/tests/iter/tyieldintry.nim
index 9cb199c5b8..31ec65a830 100644
--- a/tests/iter/tyieldintry.nim
+++ b/tests/iter/tyieldintry.nim
@@ -5,197 +5,368 @@ output: "ok"
 var closureIterResult = newSeq[int]()
 
 proc checkpoint(arg: int) =
-    closureIterResult.add(arg)
+  closureIterResult.add(arg)
 
 type
-    TestException = object of Exception
-    AnotherException = object of Exception
+  TestException = object of Exception
+  AnotherException = object of Exception
 
 proc testClosureIterAux(it: iterator(): int, exceptionExpected: bool, expectedResults: varargs[int]) =
-    closureIterResult.setLen(0)
+  closureIterResult.setLen(0)
 
-    var exceptionCaught = false
+  var exceptionCaught = false
 
-    try:
-        for i in it():
-            closureIterResult.add(i)
-    except TestException:
-        exceptionCaught = true
+  try:
+    for i in it():
+      closureIterResult.add(i)
+  except TestException:
+    exceptionCaught = true
 
-    if closureIterResult != @expectedResults or exceptionCaught != exceptionExpected:
-        if closureIterResult != @expectedResults:
-            echo "Expected: ", @expectedResults
-            echo "Actual: ", closureIterResult
-        if exceptionCaught != exceptionExpected:
-            echo "Expected exception: ", exceptionExpected
-            echo "Got exception: ", exceptionCaught
-        doAssert(false)
+  if closureIterResult != @expectedResults or exceptionCaught != exceptionExpected:
+    if closureIterResult != @expectedResults:
+      echo "Expected: ", @expectedResults
+      echo "Actual: ", closureIterResult
+    if exceptionCaught != exceptionExpected:
+      echo "Expected exception: ", exceptionExpected
+      echo "Got exception: ", exceptionCaught
+    doAssert(false)
 
 proc test(it: iterator(): int, expectedResults: varargs[int]) =
-    testClosureIterAux(it, false, expectedResults)
+  testClosureIterAux(it, false, expectedResults)
 
 proc testExc(it: iterator(): int, expectedResults: varargs[int]) =
-    testClosureIterAux(it, true, expectedResults)
+  testClosureIterAux(it, true, expectedResults)
 
 proc raiseException() =
-    raise newException(TestException, "Test exception!")
+  raise newException(TestException, "Test exception!")
 
 block:
-    iterator it(): int {.closure.} =
-        var i = 5
-        while i != 0:
-            yield i
-            if i == 3:
-                yield 123
-            dec i
+  iterator it(): int {.closure.} =
+    var i = 5
+    while i != 0:
+      yield i
+      if i == 3:
+        yield 123
+      dec i
 
-    test(it, 5, 4, 3, 123, 2, 1)
+  test(it, 5, 4, 3, 123, 2, 1)
 
 block:
-    iterator it(): int {.closure.} =
-        yield 0
-        try:
-            checkpoint(1)
-            raiseException()
-        except TestException:
-            checkpoint(2)
-            yield 3
-            checkpoint(4)
-        finally:
-            checkpoint(5)
+  iterator it(): int {.closure.} =
+    yield 0
+    try:
+      checkpoint(1)
+      raiseException()
+    except TestException:
+      checkpoint(2)
+      yield 3
+      checkpoint(4)
+    finally:
+      checkpoint(5)
 
-        checkpoint(6)
+    checkpoint(6)
 
-    test(it, 0, 1, 2, 3, 4, 5, 6)
+  test(it, 0, 1, 2, 3, 4, 5, 6)
 
 block:
-    iterator it(): int {.closure.} =
-        yield 0
-        try:
-            yield 1
-            checkpoint(2)
-        finally:
-            checkpoint(3)
-            yield 4
-            checkpoint(5)
-            yield 6
+  iterator it(): int {.closure.} =
+    yield 0
+    try:
+      yield 1
+      checkpoint(2)
+    finally:
+      checkpoint(3)
+      yield 4
+      checkpoint(5)
+      yield 6
 
-    test(it, 0, 1, 2, 3, 4, 5, 6)
+  test(it, 0, 1, 2, 3, 4, 5, 6)
 
 block:
-    iterator it(): int {.closure.} =
-        yield 0
-        try:
-            yield 1
-            raiseException()
-            yield 2
-        finally:
-            checkpoint(3)
-            yield 4
-            checkpoint(5)
-            yield 6
+  iterator it(): int {.closure.} =
+    yield 0
+    try:
+      yield 1
+      raiseException()
+      yield 2
+    finally:
+      checkpoint(3)
+      yield 4
+      checkpoint(5)
+      yield 6
 
-    testExc(it, 0, 1, 3, 4, 5, 6)
+  testExc(it, 0, 1, 3, 4, 5, 6)
 
 block:
-    iterator it(): int {.closure.} =
-        try:
-            try:
-                raiseException()
-            except AnotherException:
-                yield 123
-            finally:
-                checkpoint(3)
-        finally:
-            checkpoint(4)
-
-    testExc(it, 3, 4)
-
-block:
-    iterator it(): int {.closure.} =
-        try:
-            yield 1
-            raiseException()
-        except AnotherException:
-            checkpoint(123)
-        finally:
-            checkpoint(2)
+  iterator it(): int {.closure.} =
+    try:
+      try:
+        raiseException()
+      except AnotherException:
+        yield 123
+      finally:
         checkpoint(3)
+    finally:
+      checkpoint(4)
 
-    testExc(it, 1, 2)
+  testExc(it, 3, 4)
 
 block:
-    iterator it(): int {.closure.} =
-        try:
-            yield 0
-            try:
-                yield 1
-                try:
-                    yield 2
-                    raiseException()
-                except AnotherException:
-                    yield 123
-                finally:
-                    yield 3
-            except AnotherException:
-                yield 124
-            finally:
-                yield 4
-            checkpoint(1234)
-        except:
-            yield 5
-            checkpoint(6)
-        finally:
-            checkpoint(7)
-            yield 8
-        checkpoint(9)
+  iterator it(): int {.closure.} =
+    try:
+      yield 1
+      raiseException()
+    except AnotherException:
+      checkpoint(123)
+    finally:
+      checkpoint(2)
+    checkpoint(3)
 
-    test(it, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9)
+  testExc(it, 1, 2)
 
 block:
-    iterator it(): int {.closure.} =
+  iterator it(): int {.closure.} =
+    try:
+      yield 0
+      try:
+        yield 1
         try:
-            yield 0
-            return 2
-        finally:
-            checkpoint(1)
-        checkpoint(123)
-
-    test(it, 0, 1)
-
-block:
-    iterator it(): int {.closure.} =
-        try:
-            try:
-                yield 0
-                raiseException()
-            finally:
-                checkpoint(1)
-        except TestException:
-            yield 2
-            return
-        finally:
-            yield 3
-
-        checkpoint(123)
-
-    test(it, 0, 1, 2, 3)
-
-block:
-    iterator it(): int {.closure.} =
-        try:
-            try:
-                yield 0
-                raiseException()
-            finally:
-                return # Return in finally should stop exception propagation
+          yield 2
+          raiseException()
         except AnotherException:
-            yield 2
-            return
+          yield 123
         finally:
-            yield 3
-        checkpoint(123)
+          yield 3
+      except AnotherException:
+        yield 124
+      finally:
+        yield 4
+      checkpoint(1234)
+    except:
+      yield 5
+      checkpoint(6)
+    finally:
+      checkpoint(7)
+      yield 8
+    checkpoint(9)
+
+  test(it, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9)
+
+block:
+  iterator it(): int {.closure.} =
+    try:
+      yield 0
+      return 2
+    finally:
+      checkpoint(1)
+    checkpoint(123)
+
+  test(it, 0, 1)
+
+block:
+  iterator it(): int {.closure.} =
+    try:
+      try:
+        yield 0
+        raiseException()
+      finally:
+        checkpoint(1)
+    except TestException:
+      yield 2
+      return
+    finally:
+      yield 3
+
+    checkpoint(123)
+
+  test(it, 0, 1, 2, 3)
+
+block:
+  iterator it(): int {.closure.} =
+    try:
+      try:
+        yield 0
+        raiseException()
+      finally:
+        return # Return in finally should stop exception propagation
+    except AnotherException:
+      yield 2
+      return
+    finally:
+      yield 3
+    checkpoint(123)
+
+  test(it, 0, 3)
+
+block: # Yield in yield
+  iterator it(): int {.closure.} =
+    template foo(): int =
+      yield 1
+      2
+
+    for i in 0 .. 2:
+      checkpoint(0)
+      yield foo()
+
+  test(it, 0, 1, 2, 0, 1, 2, 0, 1, 2)
+
+block:
+  iterator it(): int {.closure.} =
+    let i = if true:
+        yield 0
+        1
+      else:
+        2
+    yield i
+
+  test(it, 0, 1)
+
+block:
+  iterator it(): int {.closure.} =
+    var foo = 123
+    let i = try:
+        yield 0
+        raiseException()
+        1
+      except TestException as e:
+        assert(e.msg == "Test exception!")
+        case foo
+        of 1:
+          yield 123
+          2
+        of 123:
+          yield 5
+          6
+        else:
+          7
+    yield i
+
+  test(it, 0, 5, 6)
+
+block:
+  iterator it(): int {.closure.} =
+    proc voidFoo(i1, i2, i3: int) =
+      checkpoint(i1)
+      checkpoint(i2)
+      checkpoint(i3)
+
+    proc foo(i1, i2, i3: int): int =
+      voidFoo(i1, i2, i3)
+      i3
+
+    proc bar(i1: int): int =
+      checkpoint(i1)
+
+    template tryexcept: int =
+      try:
+        yield 1
+        raiseException()
+        123
+      except TestException:
+        yield 2
+        checkpoint(3)
+        4
+
+    let e1 = true
+
+    template ifelse1: int =
+      if e1:
+        yield 10
+        11
+      else:
+        12
+
+    template ifelse2: int =
+      if ifelse1() == 12:
+        yield 20
+        21
+      else:
+        yield 22
+        23
+
+    let i = foo(bar(0), tryexcept, ifelse2)
+    discard foo(bar(0), tryexcept, ifelse2)
+    voidFoo(bar(0), tryexcept, ifelse2)
+    yield i
+
+  test(it,
+
+    # let i = foo(bar(0), tryexcept, ifelse2)
+    0, # bar(0)
+    1, 2, 3, # tryexcept
+    10, # ifelse1
+    22, # ifelse22
+    0, 4, 23, # foo
+
+    # discard foo(bar(0), tryexcept, ifelse2)
+    0, # bar(0)
+    1, 2, 3, # tryexcept
+    10, # ifelse1
+    22, # ifelse22
+    0, 4, 23, # foo
+
+    # voidFoo(bar(0), tryexcept, ifelse2)
+    0, # bar(0)
+    1, 2, 3, # tryexcept
+    10, # ifelse1
+    22, # ifelse22
+    0, 4, 23, # foo
+
+    23 # i
+  )
+
+block:
+  iterator it(): int {.closure.} =
+    checkpoint(0)
+    for i in 0 .. 1:
+      try:
+        yield 1
+        raiseException()
+      except TestException as e:
+        doAssert(e.msg == "Test exception!")
+        yield 2
+      except AnotherException:
+        yield 123
+      except:
+        yield 1234
+      finally:
+        yield 3
+        checkpoint(4)
+        yield 5
+
+  test(it, 0, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5)
+
+block:
+  iterator it(): int {.closure.} =
+    var i = 5
+    template foo(): bool =
+      yield i
+      true
+
+    while foo():
+      dec i
+      if i == 0:
+        break
+
+  test(it, 5, 4, 3, 2, 1)
+
+block: # Short cirquits
+  iterator it(): int {.closure.} =
+    template trueYield: bool =
+      yield 1
+      true
+
+    template falseYield: bool =
+      yield 0
+      false
+
+    if trueYield or falseYield:
+      discard falseYield and trueYield
+
+    if falseYield and trueYield:
+      checkpoint(123)
+
+  test(it, 1, 0, 0)
 
-    test(it, 0, 3)
 
 echo "ok"

From c854865d3ea24d40e0823791eda5ec6669375614 Mon Sep 17 00:00:00 2001
From: Yuriy Glukhov <yuriy.glukhov@gmail.com>
Date: Tue, 8 May 2018 01:47:19 +0300
Subject: [PATCH 09/14] Corrected nkExceptBranch transformation

---
 compiler/closureiters.nim | 37 ++++++++++++++++++++++++-------------
 1 file changed, 24 insertions(+), 13 deletions(-)

diff --git a/compiler/closureiters.nim b/compiler/closureiters.nim
index a30b4e10ee..30c03bbe1e 100644
--- a/compiler/closureiters.nim
+++ b/compiler/closureiters.nim
@@ -313,33 +313,44 @@ proc newNullifyCurExc(ctx: var Ctx): PNode =
   nilnode.typ = curExc.typ
   result.add(nilnode)
 
+proc newOr(a, b: PNode): PNode =
+  result = newNode(nkCall)
+  result.add(newSymNode(getSysMagic("or", mOr)))
+  result.add(a)
+  result.add(b)
+  result.typ = getSysType(tyBool)
+
 proc collectExceptState(ctx: var Ctx, n: PNode): PNode =
   var ifStmt = newNode(nkIfStmt)
   for c in n:
     if c.kind == nkExceptBranch:
       var ifBranch: PNode
-      var branchBody: PNode
 
-      if c[0].kind == nkType:
-        assert(c.len == 2)
+      if c.len > 1:
+        var cond: PNode
+        for i in 0 .. c.len - 2:
+          assert(c[i].kind == nkType)
+          let nextCond = newNodeI(nkCall, n.info)
+          nextCond.add(newSymNode(getSysMagic("of", mOf)))
+          nextCond.add(callCodegenProc("getCurrentException", emptyNode))
+          nextCond.add(c[i])
+          nextCond.typ = getSysType(tyBool)
+
+          if cond.isNil:
+            cond = nextCond
+          else:
+            cond = newOr(cond, nextCond)
+
         ifBranch = newNode(nkElifBranch)
-        let expression = newNodeI(nkCall, n.info)
-        expression.add(newSymNode(getSysMagic("of", mOf)))
-        expression.add(callCodegenProc("getCurrentException", emptyNode))
-        expression.add(c[0])
-        expression.typ = getSysType(tyBool)
-        ifBranch.add(expression)
-        branchBody = c[1]
+        ifBranch.add(cond)
       else:
-        assert(c.len == 1)
         if ifStmt.len == 0:
           ifStmt = newNode(nkStmtList)
           ifBranch = newNode(nkStmtList)
         else:
           ifBranch = newNode(nkElse)
-        branchBody = c[0]
 
-      ifBranch.add(branchBody)
+      ifBranch.add(c[^1])
       ifStmt.add(ifBranch)
 
   if ifStmt.len != 0:

From fb965719a65809cbbfbcd9d83aabcc8216c6a54c Mon Sep 17 00:00:00 2001
From: Yuriy Glukhov <yuriy.glukhov@gmail.com>
Date: Tue, 8 May 2018 03:52:24 +0300
Subject: [PATCH 10/14] Fixed codegen (added blockLeaveActions) to closure
 iters

---
 compiler/ccgstmts.nim | 73 ++++++++++++++++++++++---------------------
 1 file changed, 38 insertions(+), 35 deletions(-)

diff --git a/compiler/ccgstmts.nim b/compiler/ccgstmts.nim
index 96f5b53a77..510dbfc185 100644
--- a/compiler/ccgstmts.nim
+++ b/compiler/ccgstmts.nim
@@ -157,6 +157,39 @@ proc genState(p: BProc, n: PNode) =
   elif n0.kind == nkStrLit:
     linefmt(p, cpsStmts, "$1: ;$n", n0.strVal.rope)
 
+proc blockLeaveActions(p: BProc, howManyTrys, howManyExcepts: int) =
+  # Called by return and break stmts.
+  # Deals with issues faced when jumping out of try/except/finally stmts,
+
+  var stack = newSeq[tuple[n: PNode, inExcept: bool]](0)
+
+  for i in countup(1, howManyTrys):
+    let tryStmt = p.nestedTryStmts.pop
+    if not p.module.compileToCpp or optNoCppExceptions in gGlobalOptions:
+      # Pop safe points generated by try
+      if not tryStmt.inExcept:
+        linefmt(p, cpsStmts, "#popSafePoint();$n")
+
+    # Pop this try-stmt of the list of nested trys
+    # so we don't infinite recurse on it in the next step.
+    stack.add(tryStmt)
+
+    # Find finally-stmt for this try-stmt
+    # and generate a copy of its sons
+    var finallyStmt = lastSon(tryStmt.n)
+    if finallyStmt.kind == nkFinally:
+      genStmts(p, finallyStmt.sons[0])
+
+  # push old elements again:
+  for i in countdown(howManyTrys-1, 0):
+    p.nestedTryStmts.add(stack[i])
+
+  if not p.module.compileToCpp or optNoCppExceptions in gGlobalOptions:
+    # Pop exceptions that was handled by the
+    # except-blocks we are in
+    for i in countdown(howManyExcepts-1, 0):
+      linefmt(p, cpsStmts, "#popCurrentException();$n")
+
 proc genGotoState(p: BProc, n: PNode) =
   # we resist the temptation to translate it into duff's device as it later
   # will be translated into computed gotos anyway for GCC at least:
@@ -167,7 +200,11 @@ proc genGotoState(p: BProc, n: PNode) =
   initLocExpr(p, n.sons[0], a)
   lineF(p, cpsStmts, "switch ($1) {$n", [rdLoc(a)])
   p.beforeRetNeeded = true
-  lineF(p, cpsStmts, "case -1: goto BeforeRet_;$n", [])
+  lineF(p, cpsStmts, "case -1:$n", [])
+  blockLeaveActions(p,
+    howManyTrys    = p.nestedTryStmts.len,
+    howManyExcepts = p.inExceptBlockLen)
+  lineF(p, cpsStmts, " goto BeforeRet_;$n", [])
   var statesCounter = lastOrd(n.sons[0].typ)
   if n.len >= 2 and n[1].kind == nkIntLit:
     statesCounter = n[1].intVal
@@ -328,40 +365,6 @@ proc genIf(p: BProc, n: PNode, d: var TLoc) =
     else: internalError(n.info, "genIf()")
   if sonsLen(n) > 1: fixLabel(p, lend)
 
-
-proc blockLeaveActions(p: BProc, howManyTrys, howManyExcepts: int) =
-  # Called by return and break stmts.
-  # Deals with issues faced when jumping out of try/except/finally stmts,
-
-  var stack = newSeq[tuple[n: PNode, inExcept: bool]](0)
-
-  for i in countup(1, howManyTrys):
-    let tryStmt = p.nestedTryStmts.pop
-    if not p.module.compileToCpp or optNoCppExceptions in gGlobalOptions:
-      # Pop safe points generated by try
-      if not tryStmt.inExcept:
-        linefmt(p, cpsStmts, "#popSafePoint();$n")
-
-    # Pop this try-stmt of the list of nested trys
-    # so we don't infinite recurse on it in the next step.
-    stack.add(tryStmt)
-
-    # Find finally-stmt for this try-stmt
-    # and generate a copy of its sons
-    var finallyStmt = lastSon(tryStmt.n)
-    if finallyStmt.kind == nkFinally:
-      genStmts(p, finallyStmt.sons[0])
-
-  # push old elements again:
-  for i in countdown(howManyTrys-1, 0):
-    p.nestedTryStmts.add(stack[i])
-
-  if not p.module.compileToCpp or optNoCppExceptions in gGlobalOptions:
-    # Pop exceptions that was handled by the
-    # except-blocks we are in
-    for i in countdown(howManyExcepts-1, 0):
-      linefmt(p, cpsStmts, "#popCurrentException();$n")
-
 proc genReturnStmt(p: BProc, t: PNode) =
   if nfPreventCg in t.flags: return
   p.beforeRetNeeded = true

From 0b5883c21ea0f99a6742fa37d7101c31bc34bcba Mon Sep 17 00:00:00 2001
From: Yuriy Glukhov <yuriy.glukhov@gmail.com>
Date: Tue, 8 May 2018 09:39:58 +0300
Subject: [PATCH 11/14] Small fix and cosmetics

---
 compiler/closureiters.nim  | 7 +------
 compiler/lambdalifting.nim | 3 +--
 2 files changed, 2 insertions(+), 8 deletions(-)

diff --git a/compiler/closureiters.nim b/compiler/closureiters.nim
index 30c03bbe1e..7172130baf 100644
--- a/compiler/closureiters.nim
+++ b/compiler/closureiters.nim
@@ -210,11 +210,6 @@ proc newCurExcAccess(ctx: var Ctx): PNode =
     ctx.curExcSym = ctx.newEnvVar(":curExc", callCodegenProc("getCurrentException", emptyNode).typ)
   ctx.newEnvVarAccess(ctx.curExcSym)
 
-proc setStateInAssgn(stateAssgn: PNode, stateNo: int) =
-  assert stateAssgn.kind == nkAsgn
-  assert stateAssgn[1].kind == nkIntLit
-  stateAssgn[1].intVal = stateNo
-
 proc newState(ctx: var Ctx, n, gotoOut: PNode): int =
   # Creates a new state, adds it to the context fills out `gotoOut` so that it
   # will goto this state.
@@ -710,7 +705,7 @@ proc lowerStmtListExprs(ctx: var Ctx, n: PNode, needsSplit: var bool): PNode =
       n[0] = ex
       result.add(n)
 
-  of nkCast:
+  of nkCast, nkHiddenStdConv, nkHiddenSubConv, nkConv:
     var ns = false
     for i in 0 ..< n.len:
       n[i] = ctx.lowerStmtListExprs(n[i], ns)
diff --git a/compiler/lambdalifting.nim b/compiler/lambdalifting.nim
index a118edf00f..3e4d09709b 100644
--- a/compiler/lambdalifting.nim
+++ b/compiler/lambdalifting.nim
@@ -881,9 +881,8 @@ proc liftForLoop*(body: PNode; owner: PSym): PNode =
       cl = createClosure()
       while true:
         let i = foo(cl)
-        if cl.state < 0:
+        if (nkBreakState(cl.state)):
           break
-        # nkBreakState(cl.state)
         ...
     """
   if liftingHarmful(owner): return body

From d99c82bc3b207952cdcce85fdd1f9033bfb7dbef Mon Sep 17 00:00:00 2001
From: Yuriy Glukhov <yuriy.glukhov@gmail.com>
Date: Tue, 8 May 2018 12:32:55 +0300
Subject: [PATCH 12/14] Cosmetics

---
 compiler/closureiters.nim  | 237 ++++++++++++-------------------------
 compiler/lambdalifting.nim |   4 +-
 2 files changed, 77 insertions(+), 164 deletions(-)

diff --git a/compiler/closureiters.nim b/compiler/closureiters.nim
index 7172130baf..f17cfbe25b 100644
--- a/compiler/closureiters.nim
+++ b/compiler/closureiters.nim
@@ -131,9 +131,8 @@
 #   break :stateLoop
 
 import
-  intsets, strutils, options, ast, astalgo, trees, treetab, msgs, os, options,
-  idents, renderer, types, magicsys, rodread, lowerings, tables, sequtils,
-  lambdalifting
+  intsets, strutils, options, ast, astalgo, trees, treetab, msgs, idents,
+  renderer, types, magicsys, rodread, lowerings, lambdalifting
 
 type
   Ctx = object
@@ -164,9 +163,7 @@ proc newStateAccess(ctx: var Ctx): PNode =
 proc newStateAssgn(ctx: var Ctx, toValue: PNode): PNode =
   # Creates state assignment:
   #   :state = toValue
-  result = newNode(nkAsgn)
-  result.add(ctx.newStateAccess())
-  result.add(toValue)
+  newTree(nkAsgn, ctx.newStateAccess(), toValue)
 
 proc newStateAssgn(ctx: var Ctx, stateNo: int = -2): PNode =
   # Creates state assignment:
@@ -300,22 +297,16 @@ proc transformBreaksInBlock(ctx: var Ctx, n: PNode, label, after: PNode): PNode
 
 proc newNullifyCurExc(ctx: var Ctx): PNode =
   # :curEcx = nil
-  result = newNode(nkAsgn)
   let curExc = ctx.newCurExcAccess()
-  result.add(curExc)
-
   let nilnode = newNode(nkNilLit)
   nilnode.typ = curExc.typ
-  result.add(nilnode)
+  result = newTree(nkAsgn, curExc, nilnode)
 
-proc newOr(a, b: PNode): PNode =
-  result = newNode(nkCall)
-  result.add(newSymNode(getSysMagic("or", mOr)))
-  result.add(a)
-  result.add(b)
+proc newOr(a, b: PNode): PNode {.inline.} =
+  result = newTree(nkCall, newSymNode(getSysMagic("or", mOr)), a, b)
   result.typ = getSysType(tyBool)
 
-proc collectExceptState(ctx: var Ctx, n: PNode): PNode =
+proc collectExceptState(ctx: var Ctx, n: PNode): PNode {.inline.} =
   var ifStmt = newNode(nkIfStmt)
   for c in n:
     if c.kind == nkExceptBranch:
@@ -349,9 +340,7 @@ proc collectExceptState(ctx: var Ctx, n: PNode): PNode =
       ifStmt.add(ifBranch)
 
   if ifStmt.len != 0:
-    result = newNode(nkStmtList)
-    result.add(ctx.newNullifyCurExc())
-    result.add(ifStmt)
+    result = newTree(nkStmtList, ctx.newNullifyCurExc(), ifStmt)
   else:
     result = emptyNode
 
@@ -362,21 +351,17 @@ proc addElseToExcept(ctx: var Ctx, n: PNode) =
     let branchBody = newNode(nkStmtList)
 
     block: # :unrollFinally = true
-      let asgn = newNode(nkAsgn)
-      asgn.add(ctx.newUnrollFinallyAccess())
-      asgn.add(newIntTypeNode(nkIntLit, 1, getSysType(tyBool)))
-      branchBody.add(asgn)
+      branchBody.add(newTree(nkAsgn,
+        ctx.newUnrollFinallyAccess(),
+        newIntTypeNode(nkIntLit, 1, getSysType(tyBool))))
 
     block: # :curExc = getCurrentException()
-      let asgn = newNode(nkAsgn)
-      asgn.add(ctx.newCurExcAccess)
-      asgn.add(callCodegenProc("getCurrentException", emptyNode))
-      branchBody.add(asgn)
+      branchBody.add(newTree(nkAsgn,
+        ctx.newCurExcAccess(),
+        callCodegenProc("getCurrentException", emptyNode)))
 
     block: # goto nearestFinally
-      let goto = newNode(nkGotoState)
-      goto.add(newIntLit(ctx.nearestFinally))
-      branchBody.add(goto)
+      branchBody.add(newTree(nkGotoState, newIntLit(ctx.nearestFinally)))
 
     elseBranch.add(branchBody)
     n[1].add(elseBranch)
@@ -421,9 +406,7 @@ proc exprToStmtList(n: PNode): tuple[s, res: PNode] =
   result.res = lastSon
 
 proc newEnvVarAsgn(ctx: Ctx, s: PSym, v: PNode): PNode =
-  result = newNode(nkFastAsgn)
-  result.add(ctx.newEnvVarAccess(s))
-  result.add(v)
+  newTree(nkFastAsgn, ctx.newEnvVarAccess(s), v)
 
 proc addExprAssgn(ctx: Ctx, output, input: PNode, sym: PSym) =
   if input.kind == nkStmtListExpr:
@@ -438,9 +421,7 @@ proc convertExprBodyToAsgn(ctx: Ctx, exprBody: PNode, res: PSym): PNode =
   ctx.addExprAssgn(result, exprBody, res)
 
 proc newNotCall(e: PNode): PNode =
-  result = newNode(nkCall)
-  result.add(newSymNode(getSysMagic("not", mNot)))
-  result.add(e)
+  result = newTree(nkCall, newSymNode(getSysMagic("not", mNot)), e)
   result.typ = getSysType(tyBool)
 
 proc lowerStmtListExprs(ctx: var Ctx, n: PNode, needsSplit: var bool): PNode =
@@ -751,19 +732,11 @@ proc lowerStmtListExprs(ctx: var Ctx, n: PNode, needsSplit: var bool): PNode =
       needsSplit = true
 
       if condNeedsSplit:
-        let newBody = newNode(nkStmtList)
-
         let (st, ex) = exprToStmtList(n[0])
-        newBody.add(st)
-        let check = newNode(nkIfStmt)
-        let branch = newNode(nkElifBranch)
-        branch.add(newNotCall(ex))
-        let brk = newNode(nkBreakStmt)
-        brk.add(emptyNode)
-        branch.add(brk)
-        check.add(branch)
-        newBody.add(check)
-        newBody.add(n[1])
+        let brk = newTree(nkBreakStmt, emptyNode)
+        let branch = newTree(nkElifBranch, newNotCall(ex), brk)
+        let check = newTree(nkIfStmt, branch)
+        let newBody = newTree(nkStmtList, st, check, n[1])
 
         n[0] = newSymNode(getSysSym("true"))
         n[1] = newBody
@@ -778,44 +751,26 @@ proc newEndFinallyNode(ctx: var Ctx): PNode =
   #         return :tmpResult
   #       else:
   #         raise
-
-  result = newNode(nkIfStmt)
-
-  let elifBranch = newNode(nkElifBranch)
-  elifBranch.add(ctx.newUnrollFinallyAccess())
-  result.add(elifBranch)
-
-  let ifBody = newNode(nkIfStmt)
-  let branch = newNode(nkElifBranch)
-
-  let cmp = newNode(nkCall)
-  cmp.add(getSysMagic("==", mEqRef).newSymNode)
   let curExc = ctx.newCurExcAccess()
   let nilnode = newNode(nkNilLit)
   nilnode.typ = curExc.typ
-  cmp.add(curExc)
-  cmp.add(nilnode)
+  let cmp = newTree(nkCall, getSysMagic("==", mEqRef).newSymNode, curExc, nilnode)
   cmp.typ = getSysType(tyBool)
-  branch.add(cmp)
 
-  let retStmt = newNode(nkReturnStmt)
-  let asgn = newNode(nkFastAsgn)
-  addSon(asgn, newSymNode(getClosureIterResult(ctx.fn)))
-  addSon(asgn, ctx.newTmpResultAccess())
-  retStmt.add(asgn)
-  branch.add(retStmt)
+  let asgn = newTree(nkFastAsgn,
+    newSymNode(getClosureIterResult(ctx.fn)),
+    ctx.newTmpResultAccess())
 
-  let elseBranch = newNode(nkElse)
-  let raiseStmt = newNode(nkRaiseStmt)
+  let retStmt = newTree(nkReturnStmt, asgn)
+  let branch = newTree(nkElifBranch, cmp, retStmt)
 
   # The C++ backend requires `getCurrentException` here.
-  raiseStmt.add(callCodegenProc("getCurrentException", emptyNode))
-  elseBranch.add(raiseStmt)
+  let raiseStmt = newTree(nkRaiseStmt, callCodegenProc("getCurrentException", emptyNode))
+  let elseBranch = newTree(nkElse, raiseStmt)
 
-  ifBody.add(branch)
-  ifBody.add(elseBranch)
-
-  elifBranch.add(ifBody)
+  let ifBody = newTree(nkIfStmt, branch, elseBranch)
+  let elifBranch = newTree(nkElifBranch, ctx.newUnrollFinallyAccess(), ifBody)
+  result = newTree(nkIfStmt, elifBranch)
 
 proc transformReturnsInTry(ctx: var Ctx, n: PNode): PNode =
   result = n
@@ -950,8 +905,7 @@ proc transformClosureIteratorBody(ctx: var Ctx, n: PNode, gotoOut: PNode): PNode
       result = newNode(nkGotoState)
       var tryBody = toStmtList(n[0])
       var exceptBody = ctx.collectExceptState(n)
-      var finallyBody = newNode(nkStmtList)
-      finallyBody.add(getFinallyNode(n))
+      var finallyBody = newTree(nkStmtList, getFinallyNode(n))
       finallyBody = ctx.transformReturnsInTry(finallyBody)
       finallyBody.add(ctx.newEndFinallyNode())
 
@@ -1125,15 +1079,13 @@ proc newArrayType(n: int, t: PType, owner: PSym): PType =
   result = newType(tyArray, owner)
 
   let rng = newType(tyRange, owner)
-  rng.n = newNode(nkRange)
-  rng.n.add(newIntLit(0))
-  rng.n.add(newIntLit(n))
+  rng.n = newTree(nkRange, newIntLit(0), newIntLit(n))
   rng.rawAddSon(t)
 
   result.rawAddSon(rng)
   result.rawAddSon(t)
 
-proc createExceptionTable(ctx: var Ctx): PNode =
+proc createExceptionTable(ctx: var Ctx): PNode {.inline.} =
   result = newNode(nkBracket)
   result.typ = newArrayType(ctx.exceptionTable.len, getSysType(tyInt16), ctx.fn)
 
@@ -1157,9 +1109,9 @@ proc newCatchBody(ctx: var Ctx): PNode {.inline.} =
   block:
 
     # exceptionTable[:state]
-    let getNextState = newNode(nkBracketExpr)
-    getNextState.add(ctx.createExceptionTable)
-    getNextState.add(ctx.newStateAccess())
+    let getNextState = newTree(nkBracketExpr,
+      ctx.createExceptionTable(),
+      ctx.newStateAccess())
     getNextState.typ = getSysType(tyInt)
 
     # :state = exceptionTable[:state]
@@ -1167,96 +1119,68 @@ proc newCatchBody(ctx: var Ctx): PNode {.inline.} =
 
   # if :state == 0: raise
   block:
-    let ifStmt = newNode(nkIfStmt)
-    let ifBranch = newNode(nkElifBranch)
-    let cond = newNode(nkCall)
-    cond.add(getSysMagic("==", mEqI).newSymNode)
-    cond.add(ctx.newStateAccess())
-    cond.add(newIntTypeNode(nkIntLit, 0, getSysType(tyInt)))
+    let cond = newTree(nkCall,
+      getSysMagic("==", mEqI).newSymNode(),
+      ctx.newStateAccess(),
+      newIntTypeNode(nkIntLit, 0, getSysType(tyInt)))
     cond.typ = getSysType(tyBool)
-    ifBranch.add(cond)
 
-    let raiseStmt = newNode(nkRaiseStmt)
-    raiseStmt.add(emptyNode)
-
-    ifBranch.add(raiseStmt)
-    ifStmt.add(ifBranch)
+    let raiseStmt = newTree(nkRaiseStmt, emptyNode)
+    let ifBranch = newTree(nkElifBranch, cond, raiseStmt)
+    let ifStmt = newTree(nkIfStmt, ifBranch)
     result.add(ifStmt)
 
   # :unrollFinally = :state > 0
   block:
-    let asgn = newNode(nkAsgn)
-    asgn.add(ctx.newUnrollFinallyAccess())
-
-    let cond = newNode(nkCall)
-    cond.add(getSysMagic("<", mLtI).newSymNode)
-    cond.add(newIntTypeNode(nkIntLit, 0, getSysType(tyInt)))
-    cond.add(ctx.newStateAccess())
+    let cond = newTree(nkCall,
+      getSysMagic("<", mLtI).newSymNode,
+      newIntTypeNode(nkIntLit, 0, getSysType(tyInt)),
+      ctx.newStateAccess())
     cond.typ = getSysType(tyBool)
-    asgn.add(cond)
+
+    let asgn = newTree(nkAsgn, ctx.newUnrollFinallyAccess(), cond)
     result.add(asgn)
 
   # if :state < 0: :state = -:state
   block:
-    let ifStmt = newNode(nkIfStmt)
-    let ifBranch = newNode(nkElifBranch)
-    let cond = newNode(nkCall)
-    cond.add(getSysMagic("<", mLtI).newSymNode)
-    cond.add(ctx.newStateAccess())
-    cond.add(newIntTypeNode(nkIntLit, 0, getSysType(tyInt)))
+    let cond = newTree(nkCall,
+      getSysMagic("<", mLtI).newSymNode,
+      ctx.newStateAccess(),
+      newIntTypeNode(nkIntLit, 0, getSysType(tyInt)))
     cond.typ = getSysType(tyBool)
-    ifBranch.add(cond)
 
-    let negateState = newNode(nkCall)
-    negateState.add(getSysMagic("-", mUnaryMinusI).newSymNode)
-    negateState.add(ctx.newStateAccess())
+    let negateState = newTree(nkCall,
+      getSysMagic("-", mUnaryMinusI).newSymNode,
+      ctx.newStateAccess())
     negateState.typ = getSysType(tyInt)
 
-    ifBranch.add(ctx.newStateAssgn(negateState))
-    ifStmt.add(ifBranch)
+    let ifBranch = newTree(nkElifBranch, cond, ctx.newStateAssgn(negateState))
+    let ifStmt = newTree(nkIfStmt, ifBranch)
     result.add(ifStmt)
 
   # :curExc = getCurrentException()
   block:
-    let getCurExc = callCodegenProc("getCurrentException", emptyNode)
-    let asgn = newNode(nkAsgn)
-    asgn.add(ctx.newCurExcAccess())
-    asgn.add(getCurExc)
-    result.add(asgn)
+    result.add(newTree(nkAsgn,
+      ctx.newCurExcAccess(),
+      callCodegenProc("getCurrentException", emptyNode)))
 
-proc wrapIntoTryExcept(ctx: var Ctx, n: PNode): PNode =
-  result = newNode(nkTryStmt)
+proc wrapIntoTryExcept(ctx: var Ctx, n: PNode): PNode {.inline.} =
+  let setupExc = newTree(nkCall,
+    newSymNode(getCompilerProc("closureIterSetupExc")),
+    ctx.newCurExcAccess())
 
-  let tryBody = newNode(nkStmtList)
+  let tryBody = newTree(nkStmtList, setupExc, n)
+  let exceptBranch = newTree(nkExceptBranch, ctx.newCatchBody())
 
-  let setupExc = newNode(nkCall)
-  setupExc.add(newSymNode(getCompilerProc("closureIterSetupExc")))
-
-  tryBody.add(setupExc)
-
-  tryBody.add(n)
-  result.add(tryBody)
-
-  let catchNode = newNode(nkExceptBranch)
-  result.add(catchNode)
-
-  let catchBody = newNode(nkStmtList)
-  catchBody.add(ctx.newCatchBody())
-  catchNode.add(catchBody)
-
-  setupExc.add(ctx.newCurExcAccess())
+  result = newTree(nkTryStmt, tryBody, exceptBranch)
 
 proc wrapIntoStateLoop(ctx: var Ctx, n: PNode): PNode =
   # while true:
   #   block :stateLoop:
   #     gotoState :state
   #     body # Might get wrapped in try-except
-
-  result = newNode(nkWhileStmt)
-  result.add(newSymNode(getSysSym("true")))
-
   let loopBody = newNodeI(nkStmtList, n.info)
-  result.add(loopBody)
+  result = newTree(nkWhileStmt, newSymNode(getSysSym("true")), loopBody)
 
   if not ctx.stateVarSym.isNil:
     let varSect = newNodeI(nkVarSection, n.info)
@@ -1269,26 +1193,19 @@ proc wrapIntoStateLoop(ctx: var Ctx, n: PNode): PNode =
   let blockStmt = newNodeI(nkBlockStmt, n.info)
   blockStmt.add(newSymNode(ctx.stateLoopLabel))
 
-  var blockBody = newNodeI(nkStmtList, n.info)
-
   let gs = newNodeI(nkGotoState, n.info)
   gs.add(ctx.newStateAccess())
   gs.add(newIntLit(ctx.states.len - 1))
 
-  blockBody.add(gs)
-  blockBody.add(n)
-
+  var blockBody = newTree(nkStmtList, gs, n)
   if ctx.hasExceptions:
     blockBody = ctx.wrapIntoTryExcept(blockBody)
 
   blockStmt.add(blockBody)
-
   loopBody.add(blockStmt)
 
 proc deleteEmptyStates(ctx: var Ctx) =
-  let goOut = newNode(nkGotoState)
-  goOut.add(newIntLit(-1))
-
+  let goOut = newTree(nkGotoState, newIntLit(-1))
   ctx.exitStateIdx = ctx.newState(goOut, nil)
 
   # Apply new state indexes and mark unused states with -1
@@ -1332,14 +1249,11 @@ proc transformClosureIterator*(fn: PSym, n: PNode): PNode =
     ctx.stateVarSym = newSym(skVar, getIdent(":state"), fn, fn.info)
     ctx.stateVarSym.typ = createClosureIterStateType(fn)
 
-  ctx.states = @[]
   ctx.stateLoopLabel = newSym(skLabel, getIdent(":stateLoop"), fn, fn.info)
-  ctx.exceptionTable = @[]
   let n = n.toStmtList
 
   discard ctx.newState(n, nil)
-  let gotoOut = newNode(nkGotoState)
-  gotoOut.add(newIntLit(-1))
+  let gotoOut = newTree(nkGotoState, newIntLit(-1))
 
   # Splitting transformation
   discard ctx.transformClosureIteratorBody(n, gotoOut)
@@ -1349,8 +1263,7 @@ proc transformClosureIterator*(fn: PSym, n: PNode): PNode =
 
   # Make new body by concating the list of states
   result = newNode(nkStmtList)
-  for i, s in ctx.states:
-    # result.add(s)
+  for s in ctx.states:
     let body = s[1]
     s.sons.del(1)
     result.add(s)
diff --git a/compiler/lambdalifting.nim b/compiler/lambdalifting.nim
index 3e4d09709b..43ff50190d 100644
--- a/compiler/lambdalifting.nim
+++ b/compiler/lambdalifting.nim
@@ -10,8 +10,8 @@
 # This file implements lambda lifting for the transformator.
 
 import
-  intsets, strutils, options, ast, astalgo, trees, treetab, msgs, os, options,
-  idents, renderer, types, magicsys, rodread, lowerings, tables, sequtils
+  intsets, strutils, options, ast, astalgo, trees, treetab, msgs,
+  idents, renderer, types, magicsys, rodread, lowerings, tables
 
 discard """
   The basic approach is that captured vars need to be put on the heap and

From c1dde282d68fb4e21025c1d14077283ddc2330fa Mon Sep 17 00:00:00 2001
From: Yuriy Glukhov <yuriy.glukhov@gmail.com>
Date: Tue, 8 May 2018 13:49:37 +0300
Subject: [PATCH 13/14] Fixed line info

---
 compiler/closureiters.nim | 1 +
 1 file changed, 1 insertion(+)

diff --git a/compiler/closureiters.nim b/compiler/closureiters.nim
index f17cfbe25b..2328e2c55a 100644
--- a/compiler/closureiters.nim
+++ b/compiler/closureiters.nim
@@ -1181,6 +1181,7 @@ proc wrapIntoStateLoop(ctx: var Ctx, n: PNode): PNode =
   #     body # Might get wrapped in try-except
   let loopBody = newNodeI(nkStmtList, n.info)
   result = newTree(nkWhileStmt, newSymNode(getSysSym("true")), loopBody)
+  result.info = n.info
 
   if not ctx.stateVarSym.isNil:
     let varSect = newNodeI(nkVarSection, n.info)

From 5d166fcc0ae1ab203f965f0650fd4834542a46f6 Mon Sep 17 00:00:00 2001
From: Yuriy Glukhov <yuriy.glukhov@gmail.com>
Date: Thu, 10 May 2018 00:27:50 +0300
Subject: [PATCH 14/14] Review comments addressed. More thorough line info
 tracking.

---
 compiler/closureiters.nim        | 130 +++++++++++++++----------------
 tests/async/tasync_traceback.nim |   4 +-
 2 files changed, 67 insertions(+), 67 deletions(-)

diff --git a/compiler/closureiters.nim b/compiler/closureiters.nim
index 2328e2c55a..193a812e1a 100644
--- a/compiler/closureiters.nim
+++ b/compiler/closureiters.nim
@@ -179,7 +179,7 @@ proc newEnvVar(ctx: var Ctx, name: string, typ: PType): PSym =
     # We haven't gone through labmda lifting yet, so just create a local var,
     # it will be lifted later
     if ctx.tempVars.isNil:
-      ctx.tempVars = newNode(nkVarSection)
+      ctx.tempVars = newNodeI(nkVarSection, ctx.fn.info)
       addVar(ctx.tempVars, newSymNode(result))
   else:
     let envParam = getEnvParam(ctx.fn)
@@ -295,9 +295,10 @@ proc transformBreaksInBlock(ctx: var Ctx, n: PNode, label, after: PNode): PNode
     for i in 0 ..< n.len:
       n[i] = ctx.transformBreaksInBlock(n[i], label, after)
 
-proc newNullifyCurExc(ctx: var Ctx): PNode =
+proc newNullifyCurExc(ctx: var Ctx, info: TLineInfo): PNode =
   # :curEcx = nil
   let curExc = ctx.newCurExcAccess()
+  curExc.info = info
   let nilnode = newNode(nkNilLit)
   nilnode.typ = curExc.typ
   result = newTree(nkAsgn, curExc, nilnode)
@@ -305,9 +306,10 @@ proc newNullifyCurExc(ctx: var Ctx): PNode =
 proc newOr(a, b: PNode): PNode {.inline.} =
   result = newTree(nkCall, newSymNode(getSysMagic("or", mOr)), a, b)
   result.typ = getSysType(tyBool)
+  result.info = a.info
 
 proc collectExceptState(ctx: var Ctx, n: PNode): PNode {.inline.} =
-  var ifStmt = newNode(nkIfStmt)
+  var ifStmt = newNodeI(nkIfStmt, n.info)
   for c in n:
     if c.kind == nkExceptBranch:
       var ifBranch: PNode
@@ -316,39 +318,39 @@ proc collectExceptState(ctx: var Ctx, n: PNode): PNode {.inline.} =
         var cond: PNode
         for i in 0 .. c.len - 2:
           assert(c[i].kind == nkType)
-          let nextCond = newNodeI(nkCall, n.info)
-          nextCond.add(newSymNode(getSysMagic("of", mOf)))
-          nextCond.add(callCodegenProc("getCurrentException", emptyNode))
-          nextCond.add(c[i])
+          let nextCond = newTree(nkCall,
+            newSymNode(getSysMagic("of", mOf)),
+            callCodegenProc("getCurrentException", emptyNode),
+            c[i])
           nextCond.typ = getSysType(tyBool)
+          nextCond.info = c.info
 
           if cond.isNil:
             cond = nextCond
           else:
             cond = newOr(cond, nextCond)
 
-        ifBranch = newNode(nkElifBranch)
+        ifBranch = newNodeI(nkElifBranch, c.info)
         ifBranch.add(cond)
       else:
         if ifStmt.len == 0:
-          ifStmt = newNode(nkStmtList)
-          ifBranch = newNode(nkStmtList)
+          ifStmt = newNodeI(nkStmtList, c.info)
+          ifBranch = newNodeI(nkStmtList, c.info)
         else:
-          ifBranch = newNode(nkElse)
+          ifBranch = newNodeI(nkElse, c.info)
 
       ifBranch.add(c[^1])
       ifStmt.add(ifBranch)
 
   if ifStmt.len != 0:
-    result = newTree(nkStmtList, ctx.newNullifyCurExc(), ifStmt)
+    result = newTree(nkStmtList, ctx.newNullifyCurExc(n.info), ifStmt)
   else:
     result = emptyNode
 
 proc addElseToExcept(ctx: var Ctx, n: PNode) =
   if n.kind == nkStmtList and n[1].kind == nkIfStmt and n[1][^1].kind != nkElse:
     # Not all cases are covered
-    let elseBranch = newNode(nkElse)
-    let branchBody = newNode(nkStmtList)
+    let branchBody = newNodeI(nkStmtList, n.info)
 
     block: # :unrollFinally = true
       branchBody.add(newTree(nkAsgn,
@@ -363,7 +365,7 @@ proc addElseToExcept(ctx: var Ctx, n: PNode) =
     block: # goto nearestFinally
       branchBody.add(newTree(nkGotoState, newIntLit(ctx.nearestFinally)))
 
-    elseBranch.add(branchBody)
+    let elseBranch = newTree(nkElse, branchBody)
     n[1].add(elseBranch)
 
 proc getFinallyNode(n: PNode): PNode =
@@ -406,7 +408,8 @@ proc exprToStmtList(n: PNode): tuple[s, res: PNode] =
   result.res = lastSon
 
 proc newEnvVarAsgn(ctx: Ctx, s: PSym, v: PNode): PNode =
-  newTree(nkFastAsgn, ctx.newEnvVarAccess(s), v)
+  result = newTree(nkFastAsgn, ctx.newEnvVarAccess(s), v)
+  result.info = v.info
 
 proc addExprAssgn(ctx: Ctx, output, input: PNode, sym: PSym) =
   if input.kind == nkStmtListExpr:
@@ -417,11 +420,11 @@ proc addExprAssgn(ctx: Ctx, output, input: PNode, sym: PSym) =
     output.add(ctx.newEnvVarAsgn(sym, input))
 
 proc convertExprBodyToAsgn(ctx: Ctx, exprBody: PNode, res: PSym): PNode =
-  result = newNode(nkStmtList)
+  result = newNodeI(nkStmtList, exprBody.info)
   ctx.addExprAssgn(result, exprBody, res)
 
 proc newNotCall(e: PNode): PNode =
-  result = newTree(nkCall, newSymNode(getSysMagic("not", mNot)), e)
+  result = newTree(nkCall, newSymNode(getSysMagic("not", mNot), e.info), e)
   result.typ = getSysType(tyBool)
 
 proc lowerStmtListExprs(ctx: var Ctx, n: PNode, needsSplit: var bool): PNode =
@@ -446,7 +449,7 @@ proc lowerStmtListExprs(ctx: var Ctx, n: PNode, needsSplit: var bool): PNode =
 
     needsSplit = true
 
-  of nkPar, nkObjConstr, nkTupleConstr, nkBracket, nkArgList:
+  of nkPar, nkObjConstr, nkTupleConstr, nkBracket:
     var ns = false
     for i in 0 ..< n.len:
       n[i] = ctx.lowerStmtListExprs(n[i], ns)
@@ -477,10 +480,10 @@ proc lowerStmtListExprs(ctx: var Ctx, n: PNode, needsSplit: var bool): PNode =
       let isExpr = not isEmptyType(n.typ)
       if isExpr:
         tmp = ctx.newTempVar(n.typ)
-        result = newNode(nkStmtListExpr)
+        result = newNodeI(nkStmtListExpr, n.info)
         result.typ = n.typ
       else:
-        result = newNode(nkStmtList)
+        result = newNodeI(nkStmtList, n.info)
 
       var curS = result
 
@@ -488,10 +491,9 @@ proc lowerStmtListExprs(ctx: var Ctx, n: PNode, needsSplit: var bool): PNode =
         case branch.kind
         of nkElseExpr, nkElse:
           if isExpr:
-            var newBranch = newNodeI(nkElse, branch.info)
-            let branchBody = newNode(nkStmtList)
+            let branchBody = newNodeI(nkStmtList, branch.info)
             ctx.addExprAssgn(branchBody, branch[0], tmp)
-            newBranch.add(branchBody)
+            let newBranch = newTree(nkElse, branchBody)
             curS.add(newBranch)
           else:
             curS.add(branch)
@@ -499,17 +501,12 @@ proc lowerStmtListExprs(ctx: var Ctx, n: PNode, needsSplit: var bool): PNode =
         of nkElifExpr, nkElifBranch:
           var newBranch: PNode
           if branch[0].kind == nkStmtListExpr:
-            let elseBody = newNode(nkStmtList)
-
             let (st, res) = exprToStmtList(branch[0])
-            elseBody.add(st)
+            let elseBody = newTree(nkStmtList, st)
 
-            newBranch = newNodeI(nkElifBranch, branch.info)
-            newBranch.add(res)
-            newBranch.add(branch[1])
+            newBranch = newTree(nkElifBranch, res, branch[1])
 
-            let newIf = newNodeI(nkIfStmt, branch.info)
-            newIf.add(newBranch)
+            let newIf = newTree(nkIfStmt, newBranch)
             elseBody.add(newIf)
             if curS.kind == nkIfStmt:
               let newElse = newNodeI(nkElse, branch.info)
@@ -523,13 +520,12 @@ proc lowerStmtListExprs(ctx: var Ctx, n: PNode, needsSplit: var bool): PNode =
             if curS.kind == nkIfStmt:
               curS.add(newBranch)
             else:
-              let newIf = newNodeI(nkIfStmt, branch.info)
-              newIf.add(newBranch)
+              let newIf = newTree(nkIfStmt, newBranch)
               curS.add(newIf)
               curS = newIf
 
           if isExpr:
-            let branchBody = newNode(nkStmtList)
+            let branchBody = newNodeI(nkStmtList, branch[1].info)
             ctx.addExprAssgn(branchBody, branch[1], tmp)
             newBranch[1] = branchBody
 
@@ -613,7 +609,7 @@ proc lowerStmtListExprs(ctx: var Ctx, n: PNode, needsSplit: var bool): PNode =
         result = newNodeI(nkStmtListExpr, n.info)
         result.typ = n.typ
       else:
-        result = newNode(nkStmtList, n.info)
+        result = newNodeI(nkStmtList, n.info)
 
       if n[0].kind == nkSym and n[0].sym.magic in {mAnd, mOr}: # `and`/`or` short cirquiting
         var cond = n[1]
@@ -625,23 +621,20 @@ proc lowerStmtListExprs(ctx: var Ctx, n: PNode, needsSplit: var bool): PNode =
         let tmp = ctx.newTempVar(cond.typ)
         result.add(ctx.newEnvVarAsgn(tmp, cond))
 
-        let ifNode = newNode(nkIfStmt)
-        let ifBranch = newNode(nkElifBranch)
-
         var check = ctx.newEnvVarAccess(tmp)
         if n[0].sym.magic == mOr:
           check = newNotCall(check)
-        ifBranch.add(check)
 
         cond = n[2]
-        let ifBody = newNode(nkStmtList)
+        let ifBody = newNodeI(nkStmtList, cond.info)
         if cond.kind == nkStmtListExpr:
           let (st, ex) = exprToStmtList(cond)
           ifBody.add(st)
           cond = ex
         ifBody.add(ctx.newEnvVarAsgn(tmp, cond))
-        ifBranch.add(ifBody)
-        ifNode.add(ifBranch)
+
+        let ifBranch = newTree(nkElifBranch, check, ifBody)
+        let ifNode = newTree(nkIfStmt, ifBranch)
         result.add(ifNode)
         result.add(ctx.newEnvVarAccess(tmp))
       else:
@@ -744,7 +737,7 @@ proc lowerStmtListExprs(ctx: var Ctx, n: PNode, needsSplit: var bool): PNode =
     for i in 0 ..< n.len:
       n[i] = ctx.lowerStmtListExprs(n[i], needsSplit)
 
-proc newEndFinallyNode(ctx: var Ctx): PNode =
+proc newEndFinallyNode(ctx: var Ctx, info: TLineInfo): PNode =
   # Generate the following code:
   #   if :unrollFinally:
   #       if :curExc.isNil:
@@ -754,11 +747,11 @@ proc newEndFinallyNode(ctx: var Ctx): PNode =
   let curExc = ctx.newCurExcAccess()
   let nilnode = newNode(nkNilLit)
   nilnode.typ = curExc.typ
-  let cmp = newTree(nkCall, getSysMagic("==", mEqRef).newSymNode, curExc, nilnode)
+  let cmp = newTree(nkCall, newSymNode(getSysMagic("==", mEqRef), info), curExc, nilnode)
   cmp.typ = getSysType(tyBool)
 
   let asgn = newTree(nkFastAsgn,
-    newSymNode(getClosureIterResult(ctx.fn)),
+    newSymNode(getClosureIterResult(ctx.fn), info),
     ctx.newTmpResultAccess())
 
   let retStmt = newTree(nkReturnStmt, asgn)
@@ -766,10 +759,12 @@ proc newEndFinallyNode(ctx: var Ctx): PNode =
 
   # The C++ backend requires `getCurrentException` here.
   let raiseStmt = newTree(nkRaiseStmt, callCodegenProc("getCurrentException", emptyNode))
+  raiseStmt.info = info
   let elseBranch = newTree(nkElse, raiseStmt)
 
   let ifBody = newTree(nkIfStmt, branch, elseBranch)
   let elifBranch = newTree(nkElifBranch, ctx.newUnrollFinallyAccess(), ifBody)
+  elifBranch.info = info
   result = newTree(nkIfStmt, elifBranch)
 
 proc transformReturnsInTry(ctx: var Ctx, n: PNode): PNode =
@@ -794,7 +789,7 @@ proc transformReturnsInTry(ctx: var Ctx, n: PNode): PNode =
       asgnTmpResult.add(n[0])
       result.add(asgnTmpResult)
 
-    result.add(ctx.newNullifyCurExc())
+    result.add(ctx.newNullifyCurExc(n.info))
 
     let goto = newNodeI(nkGotoState, n.info)
     goto.add(newIntLit(ctx.nearestFinally))
@@ -826,16 +821,17 @@ proc transformClosureIteratorBody(ctx: var Ctx, n: PNode, gotoOut: PNode): PNode
 
         if n[i].hasYields:
           # Create a new split
-          let go = newNode(nkGotoState)
+          let go = newNodeI(nkGotoState, n[i].info)
           n[i] = ctx.transformClosureIteratorBody(n[i], go)
 
-          let s = newNode(nkStmtList)
+          let s = newNodeI(nkStmtList, n[i + 1].info)
           for j in i + 1 ..< n.len:
             s.add(n[j])
 
           n.sons.setLen(i + 1)
           discard ctx.newState(s, go)
-          discard ctx.transformClosureIteratorBody(s, gotoOut)
+          if ctx.transformClosureIteratorBody(s, gotoOut) != s:
+            internalError("transformClosureIteratorBody != s")
           break
 
     of nkYieldStmt:
@@ -857,8 +853,7 @@ proc transformClosureIteratorBody(ctx: var Ctx, n: PNode, gotoOut: PNode): PNode
       if n[^1].kind != nkElse:
         # We don't have an else branch, but every possible branch has to end with
         # gotoOut, so add else here.
-        let elseBranch = newNode(nkElse)
-        elseBranch.add(gotoOut)
+        let elseBranch = newTree(nkElse, gotoOut)
         n.add(elseBranch)
 
     of nkWhileStmt:
@@ -888,8 +883,7 @@ proc transformClosureIteratorBody(ctx: var Ctx, n: PNode, gotoOut: PNode): PNode
       elifBranch.add(body)
       ifNode.add(elifBranch)
 
-      let elseBranch = newNode(nkElse)
-      elseBranch.add(gotoOut)
+      let elseBranch = newTree(nkElse, gotoOut)
       ifNode.add(elseBranch)
       s.add(ifNode)
 
@@ -902,12 +896,12 @@ proc transformClosureIteratorBody(ctx: var Ctx, n: PNode, gotoOut: PNode): PNode
       # See explanation above about how this works
       ctx.hasExceptions = true
 
-      result = newNode(nkGotoState)
+      result = newNodeI(nkGotoState, n.info)
       var tryBody = toStmtList(n[0])
       var exceptBody = ctx.collectExceptState(n)
       var finallyBody = newTree(nkStmtList, getFinallyNode(n))
       finallyBody = ctx.transformReturnsInTry(finallyBody)
-      finallyBody.add(ctx.newEndFinallyNode())
+      finallyBody.add(ctx.newEndFinallyNode(finallyBody.info))
 
       # The following index calculation is based on the knowledge how state
       # indexes are assigned
@@ -920,7 +914,7 @@ proc transformClosureIteratorBody(ctx: var Ctx, n: PNode, gotoOut: PNode): PNode
         exceptIdx = tryIdx + 1
         finallyIdx = tryIdx + 1
 
-      let outToFinally = newNode(nkGotoState)
+      let outToFinally = newNodeI(nkGotoState, finallyBody.info)
 
       block: # Create initial states.
         let oldExcHandlingState = ctx.curExcHandlingState
@@ -945,17 +939,22 @@ proc transformClosureIteratorBody(ctx: var Ctx, n: PNode, gotoOut: PNode): PNode
 
         ctx.curExcHandlingState = exceptIdx
 
-        discard ctx.transformReturnsInTry(tryBody)
-        discard ctx.transformClosureIteratorBody(tryBody, outToFinally)
+        if ctx.transformReturnsInTry(tryBody) != tryBody:
+          internalError("transformReturnsInTry != tryBody")
+        if ctx.transformClosureIteratorBody(tryBody, outToFinally) != tryBody:
+          internalError("transformClosureIteratorBody != tryBody")
 
         ctx.curExcHandlingState = finallyIdx
         ctx.addElseToExcept(exceptBody)
-        discard ctx.transformReturnsInTry(exceptBody)
-        discard ctx.transformClosureIteratorBody(exceptBody, outToFinally)
+        if ctx.transformReturnsInTry(exceptBody) != exceptBody:
+          internalError("transformReturnsInTry != exceptBody")
+        if ctx.transformClosureIteratorBody(exceptBody, outToFinally) != exceptBody:
+          internalError("transformClosureIteratorBody != exceptBody")
 
         ctx.curExcHandlingState = oldExcHandlingState
         ctx.nearestFinally = oldNearestFinally
-        discard ctx.transformClosureIteratorBody(finallyBody, gotoOut)
+        if ctx.transformClosureIteratorBody(finallyBody, gotoOut) != finallyBody:
+          internalError("transformClosureIteratorBody != finallyBody")
 
     of nkGotoState, nkForStmt:
       internalError("closure iter " & $n.kind)
@@ -1086,7 +1085,7 @@ proc newArrayType(n: int, t: PType, owner: PSym): PType =
   result.rawAddSon(t)
 
 proc createExceptionTable(ctx: var Ctx): PNode {.inline.} =
-  result = newNode(nkBracket)
+  result = newNodeI(nkBracket, ctx.fn.info)
   result.typ = newArrayType(ctx.exceptionTable.len, getSysType(tyInt16), ctx.fn)
 
   for i in ctx.exceptionTable:
@@ -1103,7 +1102,7 @@ proc newCatchBody(ctx: var Ctx): PNode {.inline.} =
   #   :state = -:state
   # :curExc = getCurrentException()
 
-  result = newNode(nkStmtList)
+  result = newNodeI(nkStmtList, ctx.fn.info)
 
   # :state = exceptionTable[:state]
   block:
@@ -1263,8 +1262,9 @@ proc transformClosureIterator*(fn: PSym, n: PNode): PNode =
   ctx.deleteEmptyStates()
 
   # Make new body by concating the list of states
-  result = newNode(nkStmtList)
+  result = newNodeI(nkStmtList, n.info)
   for s in ctx.states:
+    assert(s.len == 2)
     let body = s[1]
     s.sons.del(1)
     result.add(s)
diff --git a/tests/async/tasync_traceback.nim b/tests/async/tasync_traceback.nim
index e4c8a67b34..618a1dc769 100644
--- a/tests/async/tasync_traceback.nim
+++ b/tests/async/tasync_traceback.nim
@@ -82,7 +82,7 @@ Async traceback:
     asyncmacro\.nim\(\d+?\)\s+?a
     asyncmacro\.nim\(\d+?\)\s+?a_continue
       ## Resumes an async procedure
-    asyncmacro\.nim\(\d+?\)\s+?aIter
+    tasync_traceback\.nim\(\d+?\)\s+?aIter
     asyncfutures\.nim\(\d+?\)\s+?read
   \]#
 Exception message: b failure
@@ -110,7 +110,7 @@ Async traceback:
       ## Executes pending callbacks
     asyncmacro\.nim\(\d+?\)\s+?foo_continue
       ## Resumes an async procedure
-    asyncmacro\.nim\(\d+?\)\s+?fooIter
+    tasync_traceback\.nim\(\d+?\)\s+?fooIter
     asyncfutures\.nim\(\d+?\)\s+?read
   \]#
 Exception message: bar failure