diff --git a/src/Init/System/IO.lean b/src/Init/System/IO.lean index 7071b16a65..1cacbe6009 100644 --- a/src/Init/System/IO.lean +++ b/src/Init/System/IO.lean @@ -814,6 +814,10 @@ def set (tk : CancelToken) : BaseIO Unit := def isSet (tk : CancelToken) : BaseIO Bool := tk.ref.get +-- separate definition as otherwise no unboxed version is generated +@[export lean_io_cancel_token_is_set] +private def isSetExport := @isSet + end CancelToken namespace FS diff --git a/src/Lean/AddDecl.lean b/src/Lean/AddDecl.lean index d046787fa4..128fb20a00 100644 --- a/src/Lean/AddDecl.lean +++ b/src/Lean/AddDecl.lean @@ -14,14 +14,16 @@ register_builtin_option debug.skipKernelTC : Bool := { descr := "skip kernel type checker. WARNING: setting this option to true may compromise soundness because your proofs will not be checked by the Lean kernel" } -def Environment.addDecl (env : Environment) (opts : Options) (decl : Declaration) : Except KernelException Environment := +def Environment.addDecl (env : Environment) (opts : Options) (decl : Declaration) + (cancelTk? : Option IO.CancelToken := none) : Except KernelException Environment := if debug.skipKernelTC.get opts then addDeclWithoutChecking env decl else - addDeclCore env (Core.getMaxHeartbeats opts).toUSize decl + addDeclCore env (Core.getMaxHeartbeats opts).toUSize decl cancelTk? -def Environment.addAndCompile (env : Environment) (opts : Options) (decl : Declaration) : Except KernelException Environment := do - let env ← addDecl env opts decl +def Environment.addAndCompile (env : Environment) (opts : Options) (decl : Declaration) + (cancelTk? : Option IO.CancelToken := none) : Except KernelException Environment := do + let env ← addDecl env opts decl cancelTk? compileDecl env opts decl def addDecl (decl : Declaration) : CoreM Unit := do @@ -29,7 +31,7 @@ def addDecl (decl : Declaration) : CoreM Unit := do withTraceNode `Kernel (fun _ => return m!"typechecking declaration") do if !(← MonadLog.hasErrors) && decl.hasSorry then logWarning "declaration uses 'sorry'" - match (← getEnv).addDecl (← getOptions) decl with + match (← getEnv).addDecl (← getOptions) decl (← read).cancelTk? with | .ok env => setEnv env | .error ex => throwKernelException ex diff --git a/src/Lean/Environment.lean b/src/Lean/Environment.lean index 5cb3b690e7..22be24fff5 100644 --- a/src/Lean/Environment.lean +++ b/src/Lean/Environment.lean @@ -248,7 +248,8 @@ namespace Environment Type check given declaration and add it to the environment -/ @[extern "lean_add_decl"] -opaque addDeclCore (env : Environment) (maxHeartbeats : USize) (decl : @& Declaration) : Except KernelException Environment +opaque addDeclCore (env : Environment) (maxHeartbeats : USize) (decl : @& Declaration) + (cancelTk? : @& Option IO.CancelToken) : Except KernelException Environment /-- Add declaration to kernel without type checking it. diff --git a/src/kernel/environment.cpp b/src/kernel/environment.cpp index 99f6a70405..7757a90a55 100644 --- a/src/kernel/environment.cpp +++ b/src/kernel/environment.cpp @@ -291,9 +291,14 @@ environment environment::add(declaration const & d, bool check) const { } lean_unreachable(); } - -extern "C" LEAN_EXPORT object * lean_add_decl(object * env, size_t max_heartbeat, object * decl) { +/* +addDeclCore (env : Environment) (maxHeartbeats : USize) (decl : @& Declaration) + (cancelTk? : @& Option IO.CancelToken) : Except KernelException Environment +*/ +extern "C" LEAN_EXPORT object * lean_add_decl(object * env, size_t max_heartbeat, object * decl, + object * opt_cancel_tk) { scope_max_heartbeat s(max_heartbeat); + scope_cancel_tk s2(is_scalar(opt_cancel_tk) ? nullptr : cnstr_get(opt_cancel_tk, 0)); return catch_kernel_exceptions([&]() { return environment(env).add(declaration(decl, true)); }); diff --git a/src/runtime/interrupt.cpp b/src/runtime/interrupt.cpp index 6a0751ad57..b66fd0c8d8 100644 --- a/src/runtime/interrupt.cpp +++ b/src/runtime/interrupt.cpp @@ -10,6 +10,7 @@ Author: Leonardo de Moura #include "runtime/exception.h" #include "runtime/memory.h" #include "lean/lean.h" +#include "util/io.h" namespace lean { LEAN_THREAD_VALUE(size_t, g_max_heartbeat, 0); @@ -39,9 +40,20 @@ void check_heartbeat() { throw_heartbeat_exception(); } +LEAN_THREAD_VALUE(lean_object *, g_cancel_tk, nullptr); + +scope_cancel_tk::scope_cancel_tk(lean_object * o):flet(g_cancel_tk, o) {} + +/* CancelToken.isSet : @& IO.CancelToken → BaseIO Bool */ +extern "C" lean_obj_res lean_io_cancel_token_is_set(b_lean_obj_arg cancel_tk, lean_obj_arg); + void check_interrupted() { - if (lean_io_check_canceled_core() && !std::uncaught_exception()) { - throw interrupted(); + if (g_cancel_tk) { + inc_ref(g_cancel_tk); + if (get_io_scalar_result(lean_io_cancel_token_is_set(g_cancel_tk, lean_io_mk_world())) && + !std::uncaught_exception()) { + throw interrupted(); + } } } diff --git a/src/runtime/interrupt.h b/src/runtime/interrupt.h index 0f6b10cc51..055363d057 100644 --- a/src/runtime/interrupt.h +++ b/src/runtime/interrupt.h @@ -42,8 +42,14 @@ public: LEAN_EXPORT void check_heartbeat(); +/* Update the thread local `IO.CancelToken` (`nullptr` if unset) */ +class scope_cancel_tk : flet { +public: + scope_cancel_tk(lean_object *); +}; + /** - \brief Throw an interrupted exception if the current task is marked cancelled. + \brief Throw an interrupted exception if the current thread's cancel token is set. */ LEAN_EXPORT void check_interrupted(); diff --git a/tests/lean/run/kernelInterrupt.lean b/tests/lean/run/kernelInterrupt.lean index 183fb14fe4..f81b1f0110 100644 --- a/tests/lean/run/kernelInterrupt.lean +++ b/tests/lean/run/kernelInterrupt.lean @@ -12,9 +12,16 @@ open Lean #eval show CoreM _ from do let env ← getEnv let envPromise ← IO.Promise.new + let tk ← IO.CancelToken.new let t := Task.spawn fun _ => let env := envPromise.result.get - Kernel.whnf env {} (mkApp2 (mkConst `Nat.add) (mkNatLit 1) (mkNatLit 2)) - IO.cancel t + let decl := .axiomDecl { + name := `test + levelParams := [] + type := mkConst `Nat + isUnsafe := false + } + env.addDeclCore 1000 decl tk + tk.set envPromise.resolve env assert! t.get matches .error .interrupted