Finish closures before completing futures (fix #415)

chronos/asyncfutures2.nim

@@ -165,6 +165,12 @@ template complete*(future: Future[void]) =
   ## Completes a void ``future``.
   complete(future, getSrcLocation())
 
+proc completeWithInternalData(future: FutureBase) =
+  if not(future.cancelled()):
+    checkFinished(future, future.internalLocation[LocationKind.Finish])
+    doAssert(isNil(future.internalError))
+    future.finish(FutureState.Completed)
+
 proc fail(future: FutureBase, error: ref CatchableError, loc: ptr SrcLoc) =
   if not(future.cancelled()):
     checkFinished(future, loc)
@@ -317,6 +323,11 @@ proc futureContinue*(fut: FutureBase) {.raises: [], gcsafe.} =
       # `await` typically) or completes / fails / is cancelled
       next = fut.internalClosure(fut)
       if fut.internalClosure.finished(): # Reached the end of the transformed proc
+
+        # The async macro will have filled the value and the location directly
+        # in the future. We just need to switch to completed state after the closure
+        # is _actually_ finished (ie all the `finally`s have been ran)
+        fut.completeWithInternalData()
         break
 
       if next == nil:

chronos/asyncmacro2.nim

@@ -10,34 +10,13 @@
 import std/[macros]
 
 # `quote do` will ruin line numbers so we avoid it using these helpers
-proc completeWithResult(fut, baseType: NimNode): NimNode {.compileTime.} =
-  # when `baseType` is void:
-  #   complete(`fut`)
-  # else:
-  #   complete(`fut`, result)
-  if baseType.eqIdent("void"):
-    # Shortcut if we know baseType at macro expansion time
-    newCall(ident "complete", fut)
-  else:
-    # `baseType` might be generic and resolve to `void`
-    nnkWhenStmt.newTree(
-      nnkElifExpr.newTree(
-        nnkInfix.newTree(ident "is", baseType, ident "void"),
-        newCall(ident "complete", fut)
-      ),
-      nnkElseExpr.newTree(
-        newCall(ident "complete", fut, ident "result")
-      )
-    )
-
-proc completeWithNode(fut, baseType, node: NimNode): NimNode {.compileTime.} =
+proc completeWithNode(node, resultNode: NimNode): NimNode {.compileTime.} =
   #   when typeof(`node`) is void:
   #     `node` # statement / explicit return
-  #     -> completeWithResult(fut, baseType)
   #   else: # expression / implicit return
-  #     complete(`fut`, `node`)
-  if node.kind == nnkEmpty: # shortcut when known at macro expanstion time
-    completeWithResult(fut, baseType)
+  #     result = `node`
+  if node.kind == nnkEmpty: # shortcut when known at macro expansion time
+    node
   else:
     # Handle both expressions and statements - since the type is not know at
     # macro expansion time, we delegate this choice to a later compilation stage
@@ -48,21 +27,40 @@ proc completeWithNode(fut, baseType, node: NimNode): NimNode {.compileTime.} =
           ident "is", nnkTypeOfExpr.newTree(node), ident "void"),
         newStmtList(
           node,
-          completeWithResult(fut, baseType)
         )
       ),
       nnkElseExpr.newTree(
-        newCall(ident "complete", fut, node)
+        newAssignment(resultNode, node)
       )
     )
 
-proc processBody(node, fut, baseType: NimNode): NimNode {.compileTime.} =
+proc processBody(node, resultNode, baseType: NimNode): NimNode {.compileTime.} =
+  # Rewrites the returns to fill the future instead
   #echo(node.treeRepr)
   case node.kind
   of nnkReturnStmt:
-    let
-      res = newNimNode(nnkStmtList, node)
-    res.add completeWithNode(fut, baseType, processBody(node[0], fut, baseType))
+    let res = newNimNode(nnkStmtList, node)
+    if node[0].kind != nnkEmpty:
+      # transforms return to:
+      # when `baseType` isnot void:
+      #   `resultNode` = processBody(node)
+      # else:
+      #   {.error: "no async return type declared".}
+      # return
+      res.add nnkWhenStmt.newTree(
+        nnkElifExpr.newTree(
+          nnkInfix.newTree(
+            ident "isnot", baseType, ident "void"),
+          newAssignment(resultNode, processBody(node[0], resultNode, baseType))
+        ),
+        nnkElse.newTree(
+          nnkPragma.newTree(
+            nnkExprColonExpr.newTree(
+              newIdentNode("error"),
+              newLit("no async return type declared"))
+          )
+        )
+      )
     res.add newNimNode(nnkReturnStmt, node).add(newNilLit())
 
     res
@@ -74,7 +72,7 @@ proc processBody(node, fut, baseType: NimNode): NimNode {.compileTime.} =
       # We must not transform nested procedures of any form, otherwise
       # `fut` will be used for all nested procedures as their own
       # `retFuture`.
-      node[i] = processBody(node[i], fut, baseType)
+      node[i] = processBody(node[i], resultNode, baseType)
     node
 
 proc getName(node: NimNode): string {.compileTime.} =
@@ -154,14 +152,19 @@ proc asyncSingleProc(prc: NimNode): NimNode {.compileTime.} =
         else: returnType
       castFutureSym = nnkCast.newTree(internalFutureType, internalFutureSym)
 
-      procBody = prc.body.processBody(castFutureSym, baseType)
+      resultNode = nnkDotExpr.newTree(castFutureSym, ident"internalValue")
+      procBody = prc.body.processBody(resultNode, baseType)
 
     # don't do anything with forward bodies (empty)
     if procBody.kind != nnkEmpty:
       let
         # fix #13899, `defer` should not escape its original scope
         procBodyBlck = nnkBlockStmt.newTree(newEmptyNode(), procBody)
 
+        # workaround https://github.com/nim-lang/Nim/issues/22645
+        internalFutureTypeSym = genSym(nskType, "InternalFutureType")
+        castFutureSym2 = nnkCast.newTree(internalFutureTypeSym, internalFutureSym)
+
         resultDecl = nnkWhenStmt.newTree(
           # when `baseType` is void:
           nnkElifExpr.newTree(
@@ -173,34 +176,26 @@ proc asyncSingleProc(prc: NimNode): NimNode {.compileTime.} =
           ),
           # else:
           nnkElseExpr.newTree(
-            newStmtList(
-              quote do: {.push warning[resultshadowed]: off.},
-              # var result {.used.}: `baseType`
-              # In the proc body, result may or may not end up being used
-              # depending on how the body is written - with implicit returns /
-              # expressions in particular, it is likely but not guaranteed that
-              # it is not used. Ideally, we would avoid emitting it in this
-              # case to avoid the default initializaiton. {.used.} typically
-              # works better than {.push.} which has a tendency to leak out of
-              # scope.
-              # TODO figure out if there's a way to detect `result` usage in
-              #      the proc body _after_ template exapnsion, and therefore
-              #      avoid creating this variable - one option is to create an
-              #      addtional when branch witha fake `result` and check
-              #      `compiles(procBody)` - this is not without cost though
-              nnkVarSection.newTree(nnkIdentDefs.newTree(
-                nnkPragmaExpr.newTree(
-                  ident "result",
-                  nnkPragma.newTree(ident "used")),
-                baseType, newEmptyNode())
-                ),
-              quote do: {.pop.},
+            nnkStmtList.newTree(
+              # type `internalFutureTypeSym` = `internalFutureType` (workaround Nim#22645)
+              nnkTypeSection.newTree(nnkTypeDef.newTree(
+                internalFutureTypeSym, newEmptyNode(), internalFutureType
+              )),
+              # template result(): untyped {.used.} = `castFutureSym2`.internalValue
+              nnkTemplateDef.newTree(
+                ident"result",
+                newEmptyNode(),
+                newEmptyNode(),
+                nnkFormalParams.newTree(ident"untyped"),
+                nnkPragma.newTree(ident"used"),
+                newEmptyNode(),
+                nnkDotExpr.newTree(castFutureSym2, ident"internalValue")
+              )
             )
           )
         )
 
-        completeDecl = completeWithNode(castFutureSym, baseType, procBodyBlck)
-
+        completeDecl = completeWithNode(procBodyBlck, resultNode)
         closureBody = newStmtList(resultDecl, completeDecl)
 
         internalFutureParameter = nnkIdentDefs.newTree(
@@ -265,6 +260,17 @@ proc asyncSingleProc(prc: NimNode): NimNode {.compileTime.} =
           iteratorNameSym)
       )
 
+      # -> resultFuture.internalLocation[LocationKind.Finish] = getSrcLocation("", -1)
+      outerProcBody.add(
+        newAssignment(
+          nnkBracketExpr.newTree(
+            newDotExpr(retFutureSym, newIdentNode("internalLocation")),
+            newDotExpr(newIdentNode("LocationKind"), newIdentNode("Finish")),
+          ),
+          newCall("getSrcLocation", newLit("\"\""), newLit(-1))
+        )
+      )
+
       # -> futureContinue(resultFuture))
       outerProcBody.add(
           newCall(newIdentNode("futureContinue"), retFutureSym)

chronos/srcloc.nim

@@ -36,6 +36,6 @@ proc srcLocImpl(procedure: static string,
   )
   return addr(loc)
 
-template getSrcLocation*(procedure: static string = ""): ptr SrcLoc =
-  srcLocImpl(procedure,
-             instantiationInfo(-2).filename, instantiationInfo(-2).line)
+template getSrcLocation*(procedure: static string = "", level: static int = -2): ptr SrcLoc =
+  const instInfo = instantiationInfo(level)
+  srcLocImpl(procedure, instInfo.filename, instInfo.line)

tests/testfut.nim

@@ -1218,8 +1218,8 @@ suite "Future[T] behavior test suite":
   test "location test":
     # WARNING: This test is very sensitive to line numbers and module name.
 
-    proc macroFuture() {.async.} =       # LINE POSITION 1
-      let someVar {.used.} = 5           # LINE POSITION 2
+    proc macroFuture() {.async.} =       # LINE POSITION 1 & 2
+      let someVar {.used.} = 5
       let someOtherVar {.used.} = 4
       if true:
         let otherVar {.used.} = 3
@@ -1256,7 +1256,7 @@ suite "Future[T] behavior test suite":
 
     check:
       chk(loc10, "testfut.nim", 1221, "macroFuture")
-      chk(loc11, "testfut.nim", 1222, "")
+      chk(loc11, "testfut.nim", 1221, "")
       chk(loc20, "testfut.nim", 1234, "template")
       chk(loc21, "testfut.nim", 1237, "")
       chk(loc30, "testfut.nim", 1231, "procedure")

tests/testmacro.nim

@@ -136,6 +136,27 @@ suite "Macro transformations test suite":
     check:
       waitFor(gen(int)) == default(int)
 
+  test "Nested return":
+    proc nr: Future[int] {.async.} =
+      return
+        if 1 == 1:
+          return 42
+        else:
+          33
+
+    check waitFor(nr()) == 42
+
+  test "Issue #415 (run closure to completion on return)":
+    var x = 0
+    proc test415 {.async.} =
+      try:
+        return
+      finally:
+        await stepsAsync(1)
+        x = 5
+    waitFor(test415())
+    check: x == 5
+
   test "Implicit return":
     proc implicit(): Future[int] {.async.} =
       42