tasks that support return values (#21859)

tasks.nim: Code cleanups and support expressions that produce a value

2 files changed, 53 insertions, 25 deletions

diff --git a/lib/std/tasks.nim b/lib/std/tasks.nim
index 055ddf144..dadb2bc97 100644
--- a/lib/std/tasks.nim
+++ b/lib/std/tasks.nim
@@ -11,7 +11,6 @@
 ## A `Task` should be only owned by a single Thread, it cannot be shared by threads.
 
 import std/[macros, isolation, typetraits]
-import system/ansi_c
 
 when defined(nimPreviewSlimSystem):
   import std/assertions
@@ -62,7 +61,7 @@ when compileOption("threads"):
 
 type
   Task* = object ## `Task` contains the callback and its arguments.
-    callback: proc (args: pointer) {.nimcall, gcsafe.}
+    callback: proc (args, res: pointer) {.nimcall, gcsafe.}
     args: pointer
     destroy: proc (args: pointer) {.nimcall, gcsafe.}
 
@@ -74,12 +73,12 @@ proc `=destroy`*(t: var Task) {.inline, gcsafe.} =
   if t.args != nil:
     if t.destroy != nil:
       t.destroy(t.args)
-    c_free(t.args)
+    deallocShared(t.args)
 
-proc invoke*(task: Task) {.inline, gcsafe.} =
+proc invoke*(task: Task; res: pointer = nil) {.inline, gcsafe.} =
   ## Invokes the `task`.
   assert task.callback != nil
-  task.callback(task.args)
+  task.callback(task.args, res)
 
 template checkIsolate(scratchAssignList: seq[NimNode], procParam, scratchDotExpr: NimNode) =
   # block:
@@ -110,8 +109,8 @@ macro toTask*(e: typed{nkCall | nkInfix | nkPrefix | nkPostfix | nkCommand | nkC
     let b = toTask hello(13)
     assert b is Task
 
-  if getTypeInst(e).typeKind != ntyVoid:
-    error("'toTask' cannot accept a call with a return value", e)
+  let retType = getTypeInst(e)
+  let returnsVoid = retType.typeKind == ntyVoid
 
   when compileOption("threads"):
     if not isGcSafe(e[0]):
@@ -188,27 +187,18 @@ macro toTask*(e: typed{nkCall | nkInfix | nkPrefix | nkPostfix | nkCommand | nkC
 
 
     let scratchObjPtrType = quote do:
-      cast[ptr `scratchObjType`](c_calloc(csize_t 1, csize_t sizeof(`scratchObjType`)))
+      cast[ptr `scratchObjType`](allocShared0(sizeof(`scratchObjType`)))
 
-    let scratchLetSection = newLetStmt(
-      scratchIdent,
-      scratchObjPtrType
-    )
-
-    let scratchCheck = quote do:
-      if `scratchIdent`.isNil:
-        raise newException(OutOfMemDefect, "Could not allocate memory")
+    let scratchLetSection = newLetStmt(scratchIdent, scratchObjPtrType)
 
     var stmtList = newStmtList()
     stmtList.add(scratchObj)
     stmtList.add(scratchLetSection)
-    stmtList.add(scratchCheck)
     stmtList.add(nnkBlockStmt.newTree(newEmptyNode(), newStmtList(scratchAssignList)))
 
     var functionStmtList = newStmtList()
     let funcCall = newCall(e[0], callNode)
     functionStmtList.add tempAssignList
-    functionStmtList.add funcCall
 
     let funcName = genSym(nskProc, e[0].strVal)
     let destroyName = genSym(nskProc, "destroyScratch")
@@ -216,12 +206,24 @@ macro toTask*(e: typed{nkCall | nkInfix | nkPrefix | nkPostfix | nkCommand | nkC
     let tempNode = quote("@") do:
         `=destroy`(@objTemp2[])
 
+    var funcDecl: NimNode
+    if returnsVoid:
+      funcDecl = quote do:
+        proc `funcName`(args, res: pointer) {.gcsafe, nimcall.} =
+          let `objTemp` = cast[ptr `scratchObjType`](args)
+          `functionStmtList`
+          `funcCall`
+    else:
+      funcDecl = quote do:
+        proc `funcName`(args, res: pointer) {.gcsafe, nimcall.} =
+          let `objTemp` = cast[ptr `scratchObjType`](args)
+          `functionStmtList`
+          cast[ptr `retType`](res)[] = `funcCall`
+
     result = quote do:
       `stmtList`
 
-      proc `funcName`(args: pointer) {.gcsafe, nimcall.} =
-        let `objTemp` = cast[ptr `scratchObjType`](args)
-        `functionStmtList`
+      `funcDecl`
 
       proc `destroyName`(args: pointer) {.gcsafe, nimcall.} =
         let `objTemp2` = cast[ptr `scratchObjType`](args)
@@ -232,11 +234,19 @@ macro toTask*(e: typed{nkCall | nkInfix | nkPrefix | nkPostfix | nkCommand | nkC
     let funcCall = newCall(e[0])
     let funcName = genSym(nskProc, e[0].strVal)
 
-    result = quote do:
-      proc `funcName`(args: pointer) {.gcsafe, nimcall.} =
-        `funcCall`
+    if returnsVoid:
+      result = quote do:
+        proc `funcName`(args, res: pointer) {.gcsafe, nimcall.} =
+          `funcCall`
+
+        Task(callback: `funcName`, args: nil)
+    else:
+      result = quote do:
+        proc `funcName`(args, res: pointer) {.gcsafe, nimcall.} =
+          cast[ptr `retType`](res)[] = `funcCall`
+
+        Task(callback: `funcName`, args: nil)
 
-      Task(callback: `funcName`, args: nil)
 
   when defined(nimTasksDebug):
     echo result.repr
diff --git a/tests/stdlib/ttasks.nim b/tests/stdlib/ttasks.nim
index 4889d49d9..347c3347a 100644
--- a/tests/stdlib/ttasks.nim
+++ b/tests/stdlib/ttasks.nim
@@ -505,3 +505,21 @@ block:
       b.invoke()
 
     doAssert called == 36
+
+  block:
+    proc returnsSomething(a, b: int): int = a + b
+
+    proc noArgsButReturnsSomething(): string = "abcdef"
+
+    proc testReturnValues() =
+      let t = toTask returnsSomething(2233, 11)
+      var res: int
+      t.invoke(addr res)
+      doAssert res == 2233+11
+
+      let tb = toTask noArgsButReturnsSomething()
+      var resB: string
+      tb.invoke(addr resB)
+      doAssert resB == "abcdef"
+
+    testReturnValues()