Async await try statement fixes.

3 files changed, 89 insertions, 11 deletions

diff --git a/lib/pure/asyncdispatch.nim b/lib/pure/asyncdispatch.nim
index bbd8ed895..34c4b5f70 100644
--- a/lib/pure/asyncdispatch.nim
+++ b/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 processed = processBody(n[0][i], retFutureSym, subTypeIsVoid, n)
-        if processed.kind != n[0][i].kind or processed.len != n[0][i].len:
+      var skipped = n.skipStmtList()
+      while i < skipped.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
-      temp[0] = result
-      result = temp
+      # If the tryStmt hasn't been transformed we can just put the body
+      # back into it.
+      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":
-  #  echo(toStrLit(result))
+  if prc[0].getName == "test3":
+    echo(toStrLit(result))
 
 proc recvLine*(socket: TAsyncFD): Future[string] {.async.} =
   ## Reads a line of data from ``socket``. Returned future will complete once
diff --git a/tests/async/tasyncexceptions.nim b/tests/async/tasyncexceptions.nim
index 30ef41756..c4379f7d8 100644
--- a/tests/async/tasyncexceptions.nim
+++ b/tests/async/tasyncexceptions.nim
@@ -19,7 +19,6 @@ proc processClient(fd: int) {.async.} =
   var foo = line[0]
   if foo == 'g':
     raise newException(EBase, "foobar")
-  
 
 proc serve() {.async.} =
 
diff --git a/tests/async/tasynctry.nim b/tests/async/tasynctry.nim
index 66ea40d49..99433b9d8 100644
--- a/tests/async/tasynctry.nim
+++ b/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