fix: adapt kernel interruption to new cancellation system (#4584)

Kernel checks were not canceled on edit after #3014

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

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
 

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.

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<environment>([&]() {
             return environment(env).add(declaration(decl, true));
         });

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<lean_object *>(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<bool>(lean_io_cancel_token_is_set(g_cancel_tk, lean_io_mk_world())) &&
+            !std::uncaught_exception()) {
+            throw interrupted();
+        }
     }
 }
 

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<lean_object *> {
+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();
 

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