feat: IO.TaskState

src/Init/System/IO.lean

@@ -210,8 +210,33 @@ def sleep (ms : UInt32) : BaseIO Unit :=
 /-- Request cooperative cancellation of the task. The task must explicitly call `IO.checkCanceled` to react to the cancellation. -/
 @[extern "lean_io_cancel"] opaque cancel : @& Task α → BaseIO Unit
 
+/-- The current state of a `Task` in the Lean runtime's task manager. -/
+inductive TaskState
+  /--
+  The `Task` is waiting for dependencies to complete or
+  sitting in the task manager queue waiting for a thread to run on.
+  -/
+  | waiting
+  /-- The `Task` is actively running on a thread. -/
+  | running
+  /--
+  The `Task` has finished running and its result is available.
+  Calling `Task.get` or `IO.wait` on the task will not block.
+  -/
+  | finished
+  deriving Inhabited, Repr, DecidableEq, Ord
+
+instance : LT TaskState := ltOfOrd
+instance : LE TaskState := leOfOrd
+instance : Min TaskState := minOfLe
+instance : Max TaskState := maxOfLe
+
+/-- Returns current state of the `Task` in the Lean runtime's task manager. -/
+@[extern "lean_io_get_task_state"] opaque getTaskState : @& Task α → BaseIO TaskState
+
 /-- Check if the task has finished execution, at which point calling `Task.get` will return immediately. -/
-@[extern "lean_io_has_finished"] opaque hasFinished : @& Task α → BaseIO Bool
+@[inline] def hasFinished (task : Task α) : BaseIO Bool := do
+  return (← getTaskState task) matches .finished
 
 /-- Wait for the task to finish, then return its result. -/
 @[extern "lean_io_wait"] opaque wait (t : Task α) : BaseIO α :=

src/include/lean/lean.h

@@ -1110,8 +1110,8 @@ static inline lean_obj_res lean_task_get_own(lean_obj_arg t) {
 LEAN_EXPORT bool lean_io_check_canceled_core(void);
 /* primitive for implementing `IO.cancel : Task a -> IO Unit` */
 LEAN_EXPORT void lean_io_cancel_core(b_lean_obj_arg t);
-/* primitive for implementing `IO.hasFinished : Task a -> IO Unit` */
-LEAN_EXPORT bool lean_io_has_finished_core(b_lean_obj_arg t);
+/* primitive for implementing `IO.getTaskState : Task a -> IO TaskState` */
+LEAN_EXPORT uint8_t lean_io_get_task_state_core(b_lean_obj_arg t);
 /* primitive for implementing `IO.waitAny : List (Task a) -> IO (Task a)` */
 LEAN_EXPORT b_lean_obj_res lean_io_wait_any_core(b_lean_obj_arg task_list);
 

src/runtime/io.cpp

@@ -1066,8 +1066,8 @@ extern "C" LEAN_EXPORT obj_res lean_io_cancel(b_obj_arg t, obj_arg) {
     return io_result_mk_ok(box(0));
 }
 
-extern "C" LEAN_EXPORT obj_res lean_io_has_finished(b_obj_arg t, obj_arg) {
-    return io_result_mk_ok(box(lean_io_has_finished_core(t)));
+extern "C" LEAN_EXPORT obj_res lean_io_get_task_state(b_obj_arg t, obj_arg) {
+    return io_result_mk_ok(box(lean_io_get_task_state_core(t)));
 }
 
 extern "C" LEAN_EXPORT obj_res lean_io_wait(obj_arg t, obj_arg) {

src/runtime/object.cpp

@@ -1030,8 +1030,17 @@ extern "C" LEAN_EXPORT void lean_io_cancel_core(b_obj_arg t) {
     g_task_manager->cancel(lean_to_task(t));
 }
 
-extern "C" LEAN_EXPORT bool lean_io_has_finished_core(b_obj_arg t) {
-    return lean_to_task(t)->m_value != nullptr;
+extern "C" LEAN_EXPORT uint8_t lean_io_get_task_state_core(b_obj_arg t) {
+    lean_task_object * o = lean_to_task(t);
+    if (o->m_imp) {
+        if (o->m_imp->m_closure) {
+            return 0; // waiting (waiting/queued/promised)
+        } else {
+            return 1; // running
+        }
+    } else {
+        return 2; // finished
+    }
 }
 
 extern "C" LEAN_EXPORT b_obj_res lean_io_wait_any_core(b_obj_arg task_list) {

src/runtime/object.h

@@ -288,7 +288,7 @@ inline b_obj_res task_get(b_obj_arg t) { return lean_task_get(t); }
 
 inline bool io_check_canceled_core() { return lean_io_check_canceled_core(); }
 inline void io_cancel_core(b_obj_arg t) { return lean_io_cancel_core(t); }
-inline bool io_has_finished_core(b_obj_arg t) { return lean_io_has_finished_core(t); }
+inline bool io_get_task_state_core(b_obj_arg t) { return lean_io_get_task_state_core(t); }
 inline b_obj_res io_wait_any_core(b_obj_arg task_list) { return lean_io_wait_any_core(task_list); }
 
 // =======================================