leanprover · tydeu · Apr 23, 2025 · Jun 1, 2025 · Jun 1, 2025
diff --git a/src/Init/System/IO.lean b/src/Init/System/IO.lean
@@ -195,20 +195,33 @@ the value of an opaque constant.
 namespace BaseIO
 
 /--
-Runs `act` in a separate `Task`, with priority `prio`.
+Runs `act` in a separate `Task` with priority `prio`.
 
-Running the resulting `BaseIO` action causes the task to be started eagerly. Pure accesses to the
-`Task` do not influence the impure `act`.
-
-Unlike pure tasks created by `Task.spawn`, tasks created by this function will run even if the last
-reference to the task is dropped. The `act` should explicitly check for cancellation via
-`IO.checkCanceled` if it should be terminated or otherwise react to the last reference being
-dropped.
+Running the resulting `BaseIO` action causes the task to be started eagerly. Unlike pure tasks
+created by `Task.spawn`, tasks created by this function will run even if the last reference to the
+task is dropped. The `act` should explicitly check for cancellation via `IO.checkCanceled` if it
+should be terminated or otherwise react to the last reference being dropped.
 -/
 @[extern "lean_io_as_task"]
 opaque asTask (act : BaseIO α) (prio := Task.Priority.default) : BaseIO (Task α) :=
   Task.pure <$> act
 
+/--
+Runs `act` in a separate `Task` with priority `prio`.
+
+Unlike `BaseIO.asTask`, the `Task` produced by `act` is folded into the `Task` produced
+by this function. Thus, this is conceptually equivalent to `(← act.asTask prio).bind id`,
+but more efficient.
+
+Running the resulting `BaseIO` action causes the task to be started eagerly. Unlike pure tasks
+created by `Task.spawn`, tasks created by this function will run even if the last reference to the
+task is dropped. The `act` should explicitly check for cancellation via `IO.checkCanceled` if it
+should be terminated or otherwise react to the last reference being dropped.
+-/
+@[extern "lean_io_join_task"]
+opaque joinTask (act : BaseIO (Task α)) (prio := Task.Priority.default) : BaseIO (Task α) := do
+  return (← act.asTask prio).bind id
+
 /--
 Creates a new task that waits for `t` to complete and then runs the `BaseIO` action `f` on its
 result. This new task has priority `prio`.

diff --git a/src/runtime/io.cpp b/src/runtime/io.cpp
@@ -1403,6 +1403,23 @@ extern "C" LEAN_EXPORT obj_res lean_io_as_task(obj_arg act, obj_arg prio, obj_ar
     return io_result_mk_ok(t);
 }
 
+/* {α : Type} (act : BaseIO (Task α)) (_ : IO.RealWorld) */
+static obj_res lean_io_join_task_fn(obj_arg act, obj_arg) {
+    obj_res r = apply_1(act, io_mk_world());
+    b_obj_res t = io_result_get_value(r);
+    inc_ref(t);
+    dec_ref(r);
+    return task_join_core(t);
+}
+
+/* joinTask {α : Type} (act : BaseIO (Task α)) (prio : Nat) : BaseIO (Task α) */
+extern "C" LEAN_EXPORT obj_res lean_io_join_task(obj_arg act, obj_arg prio, obj_arg) {
+    object * c = lean_alloc_closure((void*)lean_io_join_task_fn, 2, 1);
+    lean_closure_set(c, 0, act);
+    object * t = lean_task_spawn_core(c, lean_unbox(prio), /* keep_alive */ true);
+    return io_result_mk_ok(t);
+}
+
 /* {α β : Type} (f : α → BaseIO β) (a : α) : β */
 static obj_res lean_io_bind_task_fn(obj_arg f, obj_arg a) {
     object_ref r(apply_2(f, a, io_mk_world()));

diff --git a/src/runtime/object.cpp b/src/runtime/object.cpp
@@ -1125,28 +1125,32 @@ static obj_res task_bind_fn2(obj_arg t, obj_arg) {
     return v;
 }
 
-static obj_res task_bind_fn1(obj_arg x, obj_arg f, obj_arg) {
-    b_obj_res v = lean_to_task(x)->m_value;
-    lean_assert(v != nullptr);
-    lean_inc(v);
-    lean_dec_ref(x);
-    obj_res new_task = lean_apply_1(f, v);
-    lean_assert(lean_is_task(new_task));
-    v = lean_to_task(new_task)->m_value;
+obj_res task_join_core(obj_arg t) {
+    lean_assert(lean_is_task(t));
+    b_obj_res v = lean_to_task(t)->m_value;
     if (v) {
         lean_inc(v);
-        lean_dec_ref(new_task);
+        lean_dec_ref(t);
         return v;
     } else {
         lean_assert(g_current_task_object->m_imp);
         lean_assert(g_current_task_object->m_imp->m_closure == nullptr);
-        obj_res c = mk_closure_2_1(task_bind_fn2, new_task);
+        obj_res c = mk_closure_2_1(task_bind_fn2, t);
         mark_mt(c);
         g_current_task_object->m_imp->m_closure = c;
         return nullptr; /* notify queue that task did not finish yet. */
     }
 }
 
+static obj_res task_bind_fn1(obj_arg x, obj_arg f, obj_arg) {
+    b_obj_res v = lean_to_task(x)->m_value;
+    lean_assert(v != nullptr);
+    lean_inc(v);
+    lean_dec_ref(x);
+    obj_res new_task = lean_apply_1(f, v);
+    return task_join_core(new_task);
+}
+
 extern "C" LEAN_EXPORT obj_res lean_task_bind_core(obj_arg x, obj_arg f, unsigned prio,
       bool sync, bool keep_alive) {
     if (!g_task_manager || (sync && lean_to_task(x)->m_value)) {

diff --git a/src/runtime/object.h b/src/runtime/object.h
@@ -278,6 +278,7 @@ class LEAN_EXPORT scoped_task_manager {
     ~scoped_task_manager();
 };
 
+obj_res task_join_core(obj_arg t);
 inline obj_res task_spawn(obj_arg c, unsigned prio = 0, bool keep_alive = false) { return lean_task_spawn_core(c, prio, keep_alive); }
 inline obj_res task_pure(obj_arg a) { return lean_task_pure(a); }
 inline obj_res task_bind(obj_arg x, obj_arg f, unsigned prio = 0, bool sync = false, bool keep_alive = false) { return lean_task_bind_core(x, f, prio, sync, keep_alive); }

diff --git a/tests/lean/run/task_test_io.lean b/tests/lean/run/task_test_io.lean
@@ -1,3 +1,9 @@
+/-- info: 3 -/
+#guard_msgs in #eval id (α := BaseIO _) do
+  let t1 ← BaseIO.joinTask $ pure $ Task.pure 1
+  let t2 ← BaseIO.joinTask $ pure $ Task.spawn fun _ => 2
+  return t1.get + t2.get
+
 #exit  -- TODO
 
 #eval id (α := IO _) do