feat: add options maxHeartbeats and synthInstance.maxHeartbeats

src/Lean/CoreM.lean

@@ -5,6 +5,7 @@ Authors: Leonardo de Moura
 -/
 import Lean.Util.RecDepth
 import Lean.Util.Trace
+import Lean.Data.Options
 import Lean.Environment
 import Lean.Exception
 import Lean.InternalExceptionId
@@ -15,6 +16,18 @@ import Lean.ResolveName
 namespace Lean
 namespace Core
 
+def maxHeartbeatsDefault := 50000
+
+builtin_initialize
+  registerOption `maxHeartbeats {
+    defValue := DataValue.ofNat maxHeartbeatsDefault,
+    group := "",
+    descr := "maximum amount of heartbeats per command. A heartbeat is number of (small) memory allocations (in thousands), 0 means no limit"
+  }
+
+def getMaxHeartbeats (opts : Options) : Nat :=
+  opts.get `maxHeartbeats maxHeartbeatsDefault * 1000
+
 structure State where
   env             : Environment
   nextMacroScope  : MacroScope    := firstFrontendMacroScope + 1
@@ -29,6 +42,8 @@ structure Context where
   ref            : Syntax := Syntax.missing
   currNamespace  : Name := Name.anonymous
   openDecls      : List OpenDecl := []
+  initHeartbeats : Nat := 0
+  maxHeartbeats  : Nat := getMaxHeartbeats options
 
 abbrev CoreM := ReaderT Context $ StateRefT State (EIO Exception)
 
@@ -78,8 +93,8 @@ private def mkFreshNameImp (n : Name) : CoreM Name := do
   let env ← getEnv
   pure $ addMacroScope env.mainModule n fresh
 
-def mkFreshUserName [MonadLiftT CoreM m] (n : Name) : m Name :=
-  liftM $ mkFreshNameImp n
+def mkFreshUserName (n : Name) : CoreM Name :=
+  mkFreshNameImp n
 
 @[inline] def CoreM.run (x : CoreM α) (ctx : Context) (s : State) : EIO Exception (α × State) :=
   (x ctx).run s
@@ -88,7 +103,7 @@ def mkFreshUserName [MonadLiftT CoreM m] (n : Name) : m Name :=
   Prod.fst <$> x.run ctx s
 
 @[inline] def CoreM.toIO (x : CoreM α) (ctx : Context) (s : State) : IO (α × State) := do
-  match (← (x.run ctx s).toIO') with
+  match (← (x.run { ctx with initHeartbeats := (← IO.getNumHeartbeats) } s).toIO') with
   | Except.error (Exception.error _ msg)   => do let e ← msg.toString; throw $ IO.userError e
   | Except.error (Exception.internal id _) => throw $ IO.userError $ "internal exception #" ++ toString id.idx
   | Except.ok a => pure a
@@ -103,9 +118,25 @@ instance [MetaEval α] : MetaEval (CoreM α) where
 protected def withIncRecDepth [Monad m] [MonadControlT CoreM m] (x : m α) : m α :=
   controlAt CoreM fun runInBase => withIncRecDepth (runInBase x)
 
+def checkMaxHeartbeatsCore (moduleName : String) (optionName : Name) (max : Nat) : CoreM Unit := do
+  unless max == 0 do
+    let numHeartbeats ← IO.getNumHeartbeats (ε := Exception)
+    if numHeartbeats - (← read).initHeartbeats > max then
+      throwError! "(deterministic) timeout at '{moduleName}', maximum number of heartbeats ({max/1000}) has been reached (use 'set_option {optionName} <num>' to set the limit)"
+
+def checkMaxHeartbeats (moduleName : String) : CoreM Unit := do
+  checkMaxHeartbeatsCore moduleName `maxHeartbeats (← read).maxHeartbeats
+
+private def withCurrHeartbeatsImp (x : CoreM α) : CoreM α := do
+  let heartbeats ← IO.getNumHeartbeats (ε := Exception)
+  withReader (fun ctx => { ctx with initHeartbeats := heartbeats }) x
+
+def withCurrHeartbeats [Monad m] [MonadControlT CoreM m] (x : m α) : m α :=
+  controlAt CoreM fun runInBase => withCurrHeartbeatsImp (runInBase x)
+
 end Core
 
-export Core (CoreM mkFreshUserName)
+export Core (CoreM mkFreshUserName checkMaxHeartbeats withCurrHeartbeats)
 
 @[inline] def catchInternalId [Monad m] [MonadExcept Exception m] (id : InternalExceptionId) (x : m α) (h : Exception → m α) : m α := do
   try

src/Lean/Elab/Command.lean

@@ -99,21 +99,23 @@ instance : AddErrorMessageContext CommandElabM where
 def mkMessageAux (ctx : Context) (ref : Syntax) (msgData : MessageData) (severity : MessageSeverity) : Message :=
   mkMessageCore ctx.fileName ctx.fileMap msgData severity (ref.getPos?.getD ctx.cmdPos)
 
-private def mkCoreContext (ctx : Context) (s : State) : Core.Context :=
-  let scope       := s.scopes.head!
-  { options       := scope.opts
-    currRecDepth  := ctx.currRecDepth
-    maxRecDepth   := s.maxRecDepth
-    ref           := ctx.ref
-    currNamespace := scope.currNamespace
-    openDecls     := scope.openDecls }
+private def mkCoreContext (ctx : Context) (s : State) (heartbeats : Nat) : Core.Context :=
+  let scope        := s.scopes.head!
+  { options        := scope.opts
+    currRecDepth   := ctx.currRecDepth
+    maxRecDepth    := s.maxRecDepth
+    ref            := ctx.ref
+    currNamespace  := scope.currNamespace
+    openDecls      := scope.openDecls,
+    initHeartbeats := heartbeats }
 
 def liftCoreM {α} (x : CoreM α) : CommandElabM α := do
   let s ← get
   let ctx ← read
+  let heartbeats ← IO.getNumHeartbeats (ε := Exception)
   let Eα := Except Exception α
   let x : CoreM Eα := try let a ← x; pure $ Except.ok a catch ex => pure $ Except.error ex
-  let x : EIO Exception (Eα × Core.State) := (ReaderT.run x (mkCoreContext ctx s)).run { env := s.env, ngen := s.ngen }
+  let x : EIO Exception (Eα × Core.State) := (ReaderT.run x (mkCoreContext ctx s heartbeats)).run { env := s.env, ngen := s.ngen }
   let (ea, coreS) ← liftM x
   modify fun s => { s with env := coreS.env, ngen := coreS.ngen }
   match ea with
@@ -292,12 +294,14 @@ private def addTraceAsMessages (ctx : Context) (log : MessageLog) (traceState :
 def liftTermElabM {α} (declName? : Option Name) (x : TermElabM α) : CommandElabM α := do
   let ctx ← read
   let s   ← get
+  let heartbeats ← IO.getNumHeartbeats (ε := Exception)
+  -- dbgTrace! "heartbeats: {heartbeats}"
   let scope := s.scopes.head!
   -- We execute `x` with an empty message log. Thus, `x` cannot modify/view messages produced by previous commands.
   -- This is useful for implementing `runTermElabM` where we use `Term.resetMessageLog`
   let x : MetaM _      := (observing x).run (mkTermContext ctx s declName?) (mkTermState scope s)
   let x : CoreM _      := x.run mkMetaContext {}
-  let x : EIO _ _      := x.run (mkCoreContext ctx s) { env := s.env, ngen := s.ngen, nextMacroScope := s.nextMacroScope }
+  let x : EIO _ _      := x.run (mkCoreContext ctx s heartbeats) { env := s.env, ngen := s.ngen, nextMacroScope := s.nextMacroScope }
   let (((ea, termS), metaS), coreS) ← liftEIO x
   let infoTrees        := termS.infoState.trees.map fun tree =>
     let tree := tree.substitute termS.infoState.assignment

src/Lean/Elab/Term.lean

@@ -994,6 +994,7 @@ private partial def elabImplicitLambda (stx : Syntax) (catchExPostpone : Bool) :
 private partial def elabTermAux (expectedType? : Option Expr) (catchExPostpone : Bool) (implicitLambda : Bool) : Syntax → TermElabM Expr
   | stx => withFreshMacroScope $ withIncRecDepth do
     trace[Elab.step]! "expected type: {expectedType?}, term\n{stx}"
+    checkMaxHeartbeats "elaborator"
     withNestedTraces do
     let env ← getEnv
     let stxNew? ← catchInternalId unsupportedSyntaxExceptionId

src/Lean/Meta/ExprDefEq.lean

@@ -1306,6 +1306,7 @@ private def isDefEqProj : Expr → Expr → MetaM Bool
 
 partial def isExprDefEqAuxImpl (t : Expr) (s : Expr) : MetaM Bool := do
   trace[Meta.isDefEq.step]! "{t} =?= {s}"
+  checkMaxHeartbeats "isDefEq"
   withNestedTraces do
   whenUndefDo (isDefEqQuick t s) $
   whenUndefDo (isDefEqProofIrrel t s) do

src/Lean/Meta/SynthInstance.lean

@@ -17,6 +17,24 @@ namespace SynthInstance
 
 open Std (HashMap)
 
+def maxHeartbeatsDefault := 300
+
+builtin_initialize
+  registerOption `synthInstance.maxHeartbeats {
+    defValue := DataValue.ofNat maxHeartbeatsDefault,
+    group := "",
+    descr := "maximum amount of heartbeats per typeclass resolution problem. A heartbeat is number of (small) memory allocations (in thousands), 0 means no limit"
+  }
+
+def getMaxHeartbeats (opts : Options) : Nat :=
+  opts.get `synthInstance.maxHeartbeats maxHeartbeatsDefault * 1000
+
+def maxResultSizeDefault := 128
+builtin_initialize
+  registerOption `synthInstance.maxSize { defValue := maxResultSizeDefault, group := "", descr := "maximum number of instances used to construct a solution in the type class instance synthesis procedure" }
+private def getMaxSize (opts : Options) : Nat :=
+  opts.getNat `synthInstance.maxSize maxResultSizeDefault
+
 builtin_initialize inferTCGoalsRLAttr : TagAttribute ←
   registerTagAttribute `inferTCGoalsRL "instruct type class resolution procedure to solve goals from right to left for this instance"
 
@@ -143,6 +161,7 @@ structure TableEntry where
 
 structure Context where
   maxResultSize : Nat
+  maxHeartbeats : Nat
 
 /-
   Remark: the SynthInstance.State is not really an extension of `Meta.State`.
@@ -159,6 +178,9 @@ structure State where
 
 abbrev SynthM := ReaderT Context $ StateRefT State MetaM
 
+def checkMaxHeartbeats : SynthM Unit := do
+  Core.checkMaxHeartbeatsCore "typeclass" `synthInstance.maxHeartbeats (← read).maxHeartbeats
+
 @[inline] def mapMetaM (f : forall {α}, MetaM α → MetaM α) {α} : SynthM α → SynthM α :=
   monadMap @f
 
@@ -353,6 +375,7 @@ private def mkAnswer (cNode : ConsumerNode) : MetaM Answer :=
   withMCtx cNode.mctx do
     traceM `Meta.synthInstance.newAnswer do m!"size: {cNode.size}, {← inferType cNode.mvar}"
     let val ← instantiateMVars cNode.mvar
+    trace[Meta.synthInstance.newAnswer]! "val: {val}"
     let result ← abstractMVars val -- assignable metavariables become parameters
     let resultType ← inferType result.expr
     pure { result := result, resultType := resultType, size := cNode.size + 1 }
@@ -437,6 +460,7 @@ def resume : SynthM Unit := do
       consume { key := cNode.key, mvar := cNode.mvar, subgoals := rest, mctx := mctx, size := cNode.size + answer.size }
 
 def step : SynthM Bool := do
+  checkMaxHeartbeats
   let s ← get
   if !s.resumeStack.isEmpty then
     resume
@@ -460,7 +484,7 @@ partial def synth : SynthM (Option Expr) := do
     pure none
 
 def main (type : Expr) (maxResultSize : Nat) : MetaM (Option Expr) :=
-  traceCtx `Meta.synthInstance do
+  withCurrHeartbeats <| traceCtx `Meta.synthInstance do
      trace[Meta.synthInstance]! "main goal {type}"
      let mvar ← mkFreshExprMVar type
      let mctx ← getMCtx
@@ -468,7 +492,7 @@ def main (type : Expr) (maxResultSize : Nat) : MetaM (Option Expr) :=
      let action : SynthM (Option Expr) := do
        newSubgoal mctx key mvar Waiter.root
        synth
-     action.run { maxResultSize := maxResultSize } |>.run' {}
+     action.run { maxResultSize := maxResultSize, maxHeartbeats := getMaxHeartbeats (← getOptions) } |>.run' {}
 
 end SynthInstance
 
@@ -530,12 +554,6 @@ private def preprocessOutParam (type : Expr) : MetaM Expr :=
         mkForallFVars xs (mkAppN c args)
     | _ => pure type
 
-def maxResultSizeDefault := 128
-builtin_initialize
-  registerOption `maxInstSize { defValue := maxResultSizeDefault, group := "", descr := "maximum number of instances used to construct a solution in the type class instance synthesis procedure" }
-private def getMaxSize (opts : Options) : Nat :=
-  opts.getNat `maxInstSize maxResultSizeDefault
-
 /-
   Remark: when `maxResultSize? == none`, the configuration option `synthInstance.maxResultSize` is used.
   Remark: we use a different option for controlling the maximum result size for coercions.
@@ -543,7 +561,7 @@ private def getMaxSize (opts : Options) : Nat :=
 
 def synthInstance? (type : Expr) (maxResultSize? : Option Nat := none) : MetaM (Option Expr) := do profileitM Exception "typeclass inference" (← getOptions) ⟨0, 0⟩ do
   let opts ← getOptions
-  let maxResultSize := maxResultSize?.getD (getMaxSize opts)
+  let maxResultSize := maxResultSize?.getD (SynthInstance.getMaxSize opts)
   let inputConfig ← getConfig
   withConfig (fun config => { config with isDefEqStuckEx := true, transparency := TransparencyMode.reducible,
                                           foApprox := true, ctxApprox := true, constApprox := false }) do

src/Lean/Meta/WHNF.lean

@@ -513,6 +513,7 @@ private def cache (useCache : Bool) (e r : Expr) : MetaM Expr := do
 
 partial def whnfImp (e : Expr) : MetaM Expr :=
   whnfEasyCases e fun e => do
+    checkMaxHeartbeats "whnf"
     let useCache ← useWHNFCache e
     match (← cached? useCache e) with
     | some e' => pure e'

tests/lean/run/typeclass_coerce.lean

@@ -70,7 +70,8 @@ axiom Top (α : Type) (n : Nat) : Type
 
 #print "-----"
 
-set_option maxInstSize 256
+set_option synthInstance.maxSize 256
+set_option synthInstance.maxHeartbeats 5000
 
 #synth HasCoerce (Top Unit Nat.zero)
                  (Top Unit Nat.zero.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ)

tests/lean/run/typeclass_diamond.lean

@@ -1,4 +1,3 @@
-
 class Top₁   (n : Nat) : Type := (u : Unit := ())
 class Bot₁   (n : Nat) : Type := (u : Unit := ())
 class Left₁  (n : Nat) : Type := (u : Unit := ())
@@ -28,6 +27,8 @@ class Top (n : Nat) : Type := (u : Unit := ())
 instance Top₁ToTop (n : Nat) [Top₁ n] : Top n := {}
 instance Top₂ToTop (n : Nat) [Top₂ n] : Top n := {}
 
+set_option synthInstance.maxHeartbeats 500
+
 #synth Top Nat.zero.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ
 
 def tst : Top Nat.zero.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ.succ :=

tests/lean/tcloop.lean

@@ -0,0 +1,14 @@
+class A (α : Type u) where
+  a : Unit
+
+class B (α : Type u) where
+  a : Unit
+
+instance [s : B (Array α)] : A α  where
+  a := ()
+
+instance [s : A (Array α)] : B α where
+  a := ()
+
+def f : B Nat :=
+  inferInstance

tests/lean/tcloop.lean.expected.out

@@ -0,0 +1 @@
+tcloop.lean:14:2-14:15: error: (deterministic) timeout at 'typeclass', maximum number of heartbeats (300) has been reached (use 'set_option synthInstance.maxHeartbeats <num>' to set the limit)