Async await try statement fixes.

lib/pure/asyncdispatch.nim

@@ -1064,6 +1064,17 @@ proc accept*(socket: TAsyncFD,
 
 # -- Await Macro
 
+proc skipUntilStmtList(node: PNimrodNode): PNimrodNode {.compileTime.} =
+  # Skips a nest of StmtList's.
+  result = node
+  if node[0].kind == nnkStmtList:
+    result = skipUntilStmtList(node[0])
+
+proc skipStmtList(node: PNimrodNode): PNimrodNode {.compileTime.} =
+  result = node
+  if node[0].kind == nnkStmtList:
+    result = node[0]
+
 template createCb(retFutureSym, iteratorNameSym,
                    name: expr): stmt {.immediate.} =
   var nameIterVar = iteratorNameSym
@@ -1211,26 +1222,53 @@ proc processBody(node, retFutureSym: PNimrodNode,
   of nnkTryStmt:
     # try: await x; except: ...
     result = newNimNode(nnkStmtList, node)
+    template wrapInTry(n, tryBody: PNimrodNode) =
+      var temp = n
+      n[0] = tryBody
+      tryBody = temp
+
+      # Transform ``except`` body.
+      # TODO: Could we perform some ``await`` transformation here to get it
+      # working in ``except``?
+      tryBody[1] = processBody(n[1], retFutureSym, subTypeIsVoid, nil)
+
     proc processForTry(n: PNimrodNode, i: var int,
                        res: PNimrodNode): bool {.compileTime.} =
+      ## Transforms the body of the tryStmt. Does not transform the
+      ## body in ``except``.
+      ## Returns true if the tryStmt node was transformed into an ifStmt.
       result = false
-      while i < n[0].len:
+      var skipped = n.skipStmtList()
-        var processed = processBody(n[0][i], retFutureSym, subTypeIsVoid, n)
+      while i < skipped.len:
-        if processed.kind != n[0][i].kind or processed.len != n[0][i].len:
+        var processed = processBody(skipped[i], retFutureSym,
+                                    subTypeIsVoid, n)
+
+        # Check if we transformed the node into an exception check.
+        # This suggests skipped[i] contains ``await``.
+        if processed.kind != skipped[i].kind or processed.len != skipped[i].len:
+          processed = processed.skipUntilStmtList()
           expectKind(processed, nnkStmtList)
           expectKind(processed[2][1], nnkElse)
           i.inc
-          discard processForTry(n, i, processed[2][1][0])
+
+          if not processForTry(n, i, processed[2][1][0]):
+            # We need to wrap the nnkElse nodes back into a tryStmt.
+            # As they are executed if an exception does not happen
+            # inside the awaited future.
+            # The following code will wrap the nodes inside the
+            # original tryStmt.
+            wrapInTry(n, processed[2][1][0])
+
           res.add processed
           result = true
         else:
-          res.add n[0][i]
+          res.add skipped[i]
           i.inc
     var i = 0
     if not processForTry(node, i, result):
-      var temp = node
+      # If the tryStmt hasn't been transformed we can just put the body
-      temp[0] = result
+      # back into it.
-      result = temp
+      wrapInTry(node, result)
     return
   else: discard
 
@@ -1329,8 +1367,8 @@ macro async*(prc: stmt): stmt {.immediate.} =
   result[6] = outerProcBody
 
   #echo(treeRepr(result))
-  #if prc[0].getName == "catch":
+  if prc[0].getName == "test3":
-  #  echo(toStrLit(result))
+    echo(toStrLit(result))
 
 proc recvLine*(socket: TAsyncFD): Future[string] {.async.} =
   ## Reads a line of data from ``socket``. Returned future will complete once

tests/async/tasyncexceptions.nim

@@ -20,7 +20,6 @@ proc processClient(fd: int) {.async.} =
   if foo == 'g':
     raise newException(EBase, "foobar")
 
-
 proc serve() {.async.} =
 
   while true:

tests/async/tasynctry.nim

@@ -49,3 +49,44 @@ proc catch() {.async.} =
     assert false
 
 asyncCheck catch()
+
+proc test(): Future[bool] {.async.} =
+  result = false
+  try:
+    raise newException(OSError, "Foobar")
+  except:
+    result = true
+    return
+
+proc foo(): Future[bool] {.async.} = discard
+
+proc test2(): Future[bool] {.async.} =
+  result = false
+  try:
+    discard await foo()
+    raise newException(OSError, "Foobar")
+  except:
+    result = true
+    return
+
+proc test3(): Future[int] {.async.} =
+  result = 0
+  try:
+    try:
+      discard await foo()
+      raise newException(OSError, "Hello")
+    except:
+      result = 1
+      raise
+  except:
+    result = 2
+    return
+
+var x = test()
+assert x.read
+
+x = test2()
+assert x.read
+
+var y = test3()
+assert y.read == 2