feat: propagate maxHeartbeats to kernel (#4113)

Co-authored-by: Sebastian Ullrich <sebasti@nullri.ch>

src/Lean.lean

@@ -37,3 +37,4 @@ import Lean.Log
 import Lean.Linter
 import Lean.SubExpr
 import Lean.LabelAttribute
+import Lean.AddDecl

src/Lean/AddDecl.lean

@@ -0,0 +1,31 @@
+/-
+Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Leonardo de Moura
+-/
+prelude
+import Lean.CoreM
+
+namespace Lean
+
+def Environment.addDecl (env : Environment) (opts : Options) (decl : Declaration) : Except KernelException Environment :=
+  addDeclCore env (Core.getMaxHeartbeats opts).toUSize decl
+
+def Environment.addAndCompile (env : Environment) (opts : Options) (decl : Declaration) : Except KernelException Environment := do
+  let env ← addDecl env opts decl
+  compileDecl env opts decl
+
+def addDecl (decl : Declaration) : CoreM Unit := do
+  profileitM Exception "type checking" (← getOptions) 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
+      | .ok    env => setEnv env
+      | .error ex  => throwKernelException ex
+
+def addAndCompile (decl : Declaration) : CoreM Unit := do
+  addDecl decl
+  compileDecl decl
+
+end Lean

src/Lean/Compiler/InitAttr.lean

@@ -4,6 +4,7 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 prelude
+import Lean.AddDecl
 import Lean.Elab.InfoTree.Main
 
 namespace Lean

src/Lean/Compiler/Old.lean

@@ -67,9 +67,4 @@ opaque compileDecls (env : Environment) (opt : @& Options) (decls : @& List Name
 def compileDecl (env : Environment) (opt : @& Options) (decl : @& Declaration) : Except KernelException Environment :=
   compileDecls env opt (Compiler.getDeclNamesForCodeGen decl)
 
-
-def addAndCompile (env : Environment) (opt : Options) (decl : Declaration) : Except KernelException Environment := do
-  let env ← addDecl env decl
-  compileDecl env opt decl
-
 end Environment

src/Lean/CoreM.lean

@@ -338,15 +338,6 @@ def mkArrow (d b : Expr) : CoreM Expr :=
 /-- Iterated `mkArrow`, creates the expression `a₁ → a₂ → … → aₙ → b`. Also see `arrowDomainsN`. -/
 def mkArrowN (ds : Array Expr) (e : Expr) : CoreM Expr := ds.foldrM mkArrow e
 
-def addDecl (decl : Declaration) : CoreM Unit := do
-  profileitM Exception "type checking" (← getOptions) do
-    withTraceNode `Kernel (fun _ => return m!"typechecking declaration") do
-      if !(← MonadLog.hasErrors) && decl.hasSorry then
-        logWarning "declaration uses 'sorry'"
-      match (← getEnv).addDecl decl with
-      | Except.ok    env => setEnv env
-      | Except.error ex  => throwKernelException ex
-
 private def supportedRecursors :=
   #[``Empty.rec, ``False.rec, ``Eq.ndrec, ``Eq.rec, ``Eq.recOn, ``Eq.casesOn, ``False.casesOn, ``Empty.casesOn, ``And.rec, ``And.casesOn]
 
@@ -400,10 +391,6 @@ def compileDecls (decls : List Name) : CoreM Unit := do
   | Except.error ex =>
     throwKernelException ex
 
-def addAndCompile (decl : Declaration) : CoreM Unit := do
-  addDecl decl;
-  compileDecl decl
-
 def getDiag (opts : Options) : Bool :=
   diagnostics.get opts
 

src/Lean/Elab/BuiltinCommand.lean

@@ -123,7 +123,7 @@ private partial def elabChoiceAux (cmds : Array Syntax) (i : Nat) : CommandElabM
   n[1].forArgsM addUnivLevel
 
 @[builtin_command_elab «init_quot»] def elabInitQuot : CommandElab := fun _ => do
-  match (← getEnv).addDecl Declaration.quotDecl with
+  match (← getEnv).addDecl (← getOptions) Declaration.quotDecl with
   | Except.ok env   => setEnv env
   | Except.error ex => throwError (ex.toMessageData (← getOptions))
 

src/Lean/Environment.lean

@@ -246,7 +246,7 @@ namespace Environment
 
 /-- Type check given declaration and add it to the environment -/
 @[extern "lean_add_decl"]
-opaque addDecl (env : Environment) (decl : @& Declaration) : Except KernelException Environment
+opaque addDeclCore (env : Environment) (maxHeartbeats : USize) (decl : @& Declaration) : Except KernelException Environment
 
 end Environment
 

src/Lean/Meta/Closure.lean

@@ -6,6 +6,7 @@ Authors: Leonardo de Moura
 prelude
 import Lean.MetavarContext
 import Lean.Environment
+import Lean.AddDecl
 import Lean.Util.FoldConsts
 import Lean.Meta.Basic
 import Lean.Meta.Check

src/Lean/Meta/Constructions.lean

@@ -5,6 +5,7 @@ Authors: Leonardo de Moura
 -/
 prelude
 import Lean.AuxRecursor
+import Lean.AddDecl
 import Lean.Meta.AppBuilder
 
 namespace Lean

src/Lean/Meta/Eqns.lean

@@ -5,6 +5,7 @@ Authors: Leonardo de Moura
 -/
 prelude
 import Lean.ReservedNameAction
+import Lean.AddDecl
 import Lean.Meta.Basic
 import Lean.Meta.AppBuilder
 import Lean.Meta.Match.MatcherInfo

src/Lean/Meta/Eval.lean

@@ -4,6 +4,7 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Sebastian Ullrich, Leonardo de Moura
 -/
 prelude
+import Lean.AddDecl
 import Lean.Meta.Check
 
 namespace Lean.Meta

src/Lean/Meta/SizeOf.lean

@@ -5,6 +5,7 @@ Authors: Leonardo de Moura
 -/
 prelude
 import Init.Data.List.BasicAux
+import Lean.AddDecl
 import Lean.Meta.AppBuilder
 import Lean.Meta.Instances
 

src/Lean/Meta/Tactic/AuxLemma.lean

@@ -4,6 +4,7 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 prelude
+import Lean.AddDecl
 import Lean.Meta.Basic
 
 namespace Lean.Meta

src/Lean/Replay.lean

@@ -5,6 +5,7 @@ Authors: Scott Morrison
 -/
 prelude
 import Lean.CoreM
+import Lean.AddDecl
 import Lean.Util.FoldConsts
 
 /-!
@@ -56,7 +57,7 @@ def throwKernelException (ex : KernelException) : M Unit := do
 
 /-- Add a declaration, possibly throwing a `KernelException`. -/
 def addDecl (d : Declaration) : M Unit := do
-  match (← get).env.addDecl d with
+  match (← get).env.addDecl {} d with
   | .ok env => modify fun s => { s with env := env }
   | .error ex => throwKernelException ex
 

src/kernel/environment.cpp

@@ -292,7 +292,8 @@ environment environment::add(declaration const & d, bool check) const {
     lean_unreachable();
 }
 
-extern "C" LEAN_EXPORT object * lean_add_decl(object * env, object * decl) {
+extern "C" LEAN_EXPORT object * lean_add_decl(object * env, size_t max_heartbeat, object * decl) {
+    scope_max_heartbeat s(max_heartbeat);
     return catch_kernel_exceptions<environment>([&]() {
             return environment(env).add(declaration(decl, true));
         });

src/runtime/interrupt.cpp

@@ -26,7 +26,7 @@ size_t get_max_heartbeat() { return g_max_heartbeat; }
 void set_max_heartbeat_thousands(unsigned max) { g_max_heartbeat = static_cast<size_t>(max) * 1000; }
 
 scope_heartbeat::scope_heartbeat(size_t max):flet<size_t>(g_heartbeat, max) {}
-scope_max_heartbeat::scope_max_heartbeat(size_t max):flet<size_t>(g_max_heartbeat, max) {}
+LEAN_EXPORT scope_max_heartbeat::scope_max_heartbeat(size_t max):flet<size_t>(g_max_heartbeat, max) {}
 
 // separate definition to allow breakpoint in debugger
 void throw_heartbeat_exception() {

src/runtime/interrupt.h

@@ -35,9 +35,9 @@ LEAN_EXPORT void set_max_heartbeat_thousands(unsigned max);
 LEAN_EXPORT size_t get_max_heartbeat();
 
 /* Update the thread local max heartbeat */
-class scope_max_heartbeat : flet<size_t> {
+class LEAN_EXPORT scope_max_heartbeat : flet<size_t> {
 public:
-    scope_max_heartbeat(size_t max);
+    LEAN_EXPORT scope_max_heartbeat(size_t max);
 };
 
 LEAN_EXPORT void check_heartbeat();

tests/lean/partialIssue.lean

@@ -1,5 +1,5 @@
 import Lean.CoreM
-
+import Lean.AddDecl
 #eval Lean.addDecl <| .mutualDefnDecl [{
   name := `False_intro
   levelParams := []

tests/lean/run/kernel_maxheartbeats.lean

@@ -0,0 +1,22 @@
+import Lean
+
+def ack : Nat → Nat → Nat
+  | 0,   y   => y+1
+  | x+1, 0   => ack x 1
+  | x+1, y+1 => ack x (ack (x+1) y)
+
+set_option maxHeartbeats 500
+open Lean Meta
+
+/--
+error: (kernel) deterministic timeout
+-/
+#guard_msgs in
+run_meta do
+  let type ← mkEq (← mkAppM ``ack #[mkNatLit 4, mkNatLit 4]) (mkNatLit 100000)
+  let value ← mkEqRefl (mkNatLit 100000)
+  addDecl <| .thmDecl {
+    name        := `ack_4_4
+    levelParams := []
+    type, value
+  }

tests/lean/run/structInstFast.lean

@@ -92,7 +92,12 @@ def foo : String := "foo"
 
 deep_wide_struct_inst_with String foo 50 20
 
-/- Structure instances using the `with` pattern should be fast. Without #2478, this takes over 700
-  heartbeats. -/
-set_option maxHeartbeats 200 in
+/-
+Structure instances using the `with` pattern should be fast.
+Without #2478, this takes over 700 heartbeats in the elaborator.
+
+Remark: we are now propagating heartbeats to the kernel, and
+had to increase it value to 1000
+ -/
+set_option maxHeartbeats 1000 in
 example : a0 = a'0 := rfl