Fix processing of 'await' with a nnkCall.

Specifically, ``discard readMessages(disp, await disp.accept(server))``
works now, i.e. using 'await' as one of the params to a proc call.

2 files changed, 12 insertions, 9 deletions

diff --git a/lib/pure/asyncio2.nim b/lib/pure/asyncio2.nim
index d3e4bcc98..8541b2ba7 100644
--- a/lib/pure/asyncio2.nim
+++ b/lib/pure/asyncio2.nim
@@ -498,10 +498,14 @@ proc processBody(node, retFutureSym: PNimrodNode): PNimrodNode {.compileTime.} =
   of nnkCommand:
     if node[0].ident == !"await":
       case node[1].kind
-      of nnkIdent, nnkCall:
+      of nnkIdent:
         # await x
-        # await foo(p, x)
         result = newNimNode(nnkYieldStmt).add(node[1]) # -> yield x
+      of nnkCall:
+        # await foo(p, x)
+        var futureValue: PNimrodNode
+        createVar("future" & $node[1][0].toStrLit, node[1], futureValue)
+        result.add futureValue
       else:
         error("Invalid node kind in 'await', got: " & $node[1].kind)
     elif node[1].kind == nnkCommand and node[1][0].kind == nnkIdent and
@@ -510,6 +514,7 @@ proc processBody(node, retFutureSym: PNimrodNode): PNimrodNode {.compileTime.} =
       var newCommand = node
       createVar("future" & $node[0].ident, node[1][0], newCommand[1])
       result.add newCommand
+
   of nnkVarSection, nnkLetSection:
     case node[0][2].kind
     of nnkCommand:
@@ -534,11 +539,10 @@ proc processBody(node, retFutureSym: PNimrodNode): PNimrodNode {.compileTime.} =
     if node[0][0].ident == !"await":
       var dummy = newNimNode(nnkStmtList)
       createVar("futureDiscard_" & $toStrLit(node[0][1]), node[0][1], dummy)
-  else:
-    for i in 0 .. <node.len:
-      result[i] = processBody(node[i], retFutureSym)
-
-  assert(not result.isNil)
+  else: discard
+  
+  for i in 0 .. <result.len:
+    result[i] = processBody(result[i], retFutureSym)
   #echo(treeRepr(result))
 
 proc getName(node: PNimrodNode): string {.compileTime.} =
diff --git a/tests/async/tasyncawait.nim b/tests/async/tasyncawait.nim
index 84a4164d9..bcaffc287 100644
--- a/tests/async/tasyncawait.nim
+++ b/tests/async/tasyncawait.nim
@@ -52,8 +52,7 @@ proc createServer(disp: PDispatcher, port: TPort): PFuture[int] {.async.} =
   server.bindAddr(port)
   server.listen()
   while true:
-    var client = await disp.accept(server)
-    discard readMessages(disp, client)
+    discard readMessages(disp, await disp.accept(server))
 
 discard disp.createServer(TPort(10335))
 discard disp.launchSwarm(TPort(10335))