fix tests

Co-authored-by: Rob23oba <152706811+Rob23oba@users.noreply.github.com>

src/Lean/Elab/Tactic/Try.lean

@@ -225,6 +225,7 @@ structure Ctx where
   root : TSyntax `tactic
   terminal : Bool
   config : Try.Config
+  originalMaxHeartbeats : Nat
 
 abbrev TryTacticM := ReaderT Ctx TacticM
 abbrev TryTactic := TSyntax `tactic → TryTacticM (TSyntax `tactic)
@@ -324,28 +325,28 @@ private def expandUserTactic (tac : TSyntax `tactic) (goal : MVarId) : MetaM (Ar
     let initialLog ← Core.getMessageLog
     let initialMsgCount := initialLog.toList.length
 
-    let result ← try
-      -- Run the tactic to capture its "Try this" messages
-      discard <| Tactic.run goal do
-        evalTactic tac
+    let result ← tryCatchRuntimeEx
+      (do
+        -- Run the tactic to capture its "Try this" messages
+        discard <| Tactic.run goal do
+          evalTactic tac
 
-      -- Extract tactic suggestions from new messages
-      -- This parses the format produced by TryThis.addSuggestions: "Try this:\n  [apply] tactic"
-      let newMsgs := (← Core.getMessageLog).toList.drop initialMsgCount
-      let mut suggestions : Array (TSyntax `tactic) := #[]
-      for msg in newMsgs do
-        if msg.severity == MessageSeverity.information then
-          let msgText ← msg.data.toString
-          for line in msgText.splitOn "\n" do
-            if line.startsWith "  [apply] " then
-              let tacticText := line.drop "  [apply] ".length
-              let env ← getEnv
-              if let .ok stx := Parser.runParserCategory env `tactic tacticText then
-                suggestions := suggestions.push ⟨stx⟩
+        -- Extract tactic suggestions from new messages
+        -- This parses the format produced by TryThis.addSuggestions: "Try this:\n  [apply] tactic"
+        let newMsgs := (← Core.getMessageLog).toList.drop initialMsgCount
+        let mut suggestions : Array (TSyntax `tactic) := #[]
+        for msg in newMsgs do
+          if msg.severity == MessageSeverity.information then
+            let msgText ← msg.data.toString
+            for line in msgText.splitOn "\n" do
+              if line.startsWith "  [apply] " then
+                let tacticText := line.drop "  [apply] ".length
+                let env ← getEnv
+                if let .ok stx := Parser.runParserCategory env `tactic tacticText then
+                  suggestions := suggestions.push ⟨stx⟩
 
-      pure (some suggestions)
-    catch _ =>
-      pure none
+        pure (some suggestions))
+      (fun _ => pure none)
 
     initialState.restore
     Core.setMessageLog initialLog
@@ -367,6 +368,31 @@ def observing (x : TryTacticM α) : TryTacticM (TacticResult α) := do
       s.restore (restoreInfo := true)
       return .error ex sNew
 
+/--
+Executes a tactic with heartbeat management:
+- Restores the original maxHeartbeats limit (recorded at try? start)
+- Gives the tactic a fresh heartbeat budget via withCurrHeartbeats
+- Catches heartbeat exceptions and converts them to regular errors
+-/
+def withOriginalHeartbeats (x : TryTacticM α) : TryTacticM α := do
+  let originalMaxHeartbeats := (← read).originalMaxHeartbeats
+  tryCatchRuntimeEx
+    (withTheReader Core.Context (fun c => { c with maxHeartbeats := originalMaxHeartbeats }) do
+      withCurrHeartbeats x)
+    (fun ex => do
+      if Exception.isMaxHeartbeat ex then
+        throwError "tactic exceeded heartbeat limit"
+      else
+        throw ex)
+
+/--
+Executes code with unlimited heartbeats (maxHeartbeats set to 0).
+Used by try? infrastructure itself so it doesn't timeout while testing tactics.
+-/
+def withUnlimitedHeartbeats {m : Type → Type} [Monad m] [MonadWithReaderOf Core.Context m] [MonadControlT CoreM m] (x : m α) : m α :=
+  withTheReader Core.Context (fun c => { c with maxHeartbeats := 0 }) do
+    withCurrHeartbeats x
+
 private def mergeParams (ps1 ps2 : Array Syntax) : Array Syntax := Id.run do
   let mut r := ps1
   for p in ps2 do
@@ -498,89 +524,92 @@ where
       modify (·.push tac)
 
 private def evalSuggestGrindTrace : TryTactic := fun tac => do
-  let `(tactic| grind? $configStx:optConfig $[only%$only]?  $[ [$params:grindParam,*] ]?) := tac | throwUnsupportedSyntax
-  let tacs ← evalGrindTraceCore tac (trace := (← read).config.only) (verbose := false) (useSorry := false)
-  let tac ← grindTraceToGrind tac
-  trace[try.debug] "`grind` succeeded"
-  replaceMainGoal []
-  for tac1 in tacs do
-    trace[try.debug] ">> {tac1}"
-  if (← read).config.only then
-    -- If config has +suggestions, only return the 'only' version, not the original
-    if configHasSuggestions configStx then
-      mkTrySuggestions tacs
+  withOriginalHeartbeats do
+    let `(tactic| grind? $configStx:optConfig $[only%$only]?  $[ [$params:grindParam,*] ]?) := tac | throwUnsupportedSyntax
+    let tacs ← evalGrindTraceCore tac (trace := (← read).config.only) (verbose := false) (useSorry := false)
+    let tac ← grindTraceToGrind tac
+    trace[try.debug] "`grind` succeeded"
+    replaceMainGoal []
+    for tac1 in tacs do
+      trace[try.debug] ">> {tac1}"
+    if (← read).config.only then
+      -- If config has +suggestions, only return the 'only' version, not the original
+      if configHasSuggestions configStx then
+        mkTrySuggestions tacs
+      else
+        mkTrySuggestions (#[tac] ++ tacs)
     else
-      mkTrySuggestions (#[tac] ++ tacs)
-  else
-    return tac
+      return tac
 
 private def evalSuggestSimpTrace : TryTactic := fun tac => do (← getMainGoal).withContext do
   match tac with
   | `(tactic| simp? $configStx:optConfig $[only%$only]? $[[$args,*]]? $(loc)?) =>
-    let tac ← simpTraceToSimp tac
-    let { ctx, simprocs, .. } ← mkSimpContext tac (eraseLocal := false)
-    let stats ← simpLocation ctx (simprocs := simprocs) none <| (loc.map expandLocation).getD (.targets #[] true)
-    trace[try.debug] "`simp` succeeded"
-    if (← read).config.only then
-      let tac' ← mkSimpCallStx tac stats.usedTheorems
-      -- If config has +suggestions, only return the 'only' version, not the original
-      if configHasSuggestions configStx then
-        mkTrySuggestions #[tac']
+    withOriginalHeartbeats do
+      let tac ← simpTraceToSimp tac
+      let { ctx, simprocs, .. } ← mkSimpContext tac (eraseLocal := false)
+      let stats ← simpLocation ctx (simprocs := simprocs) none <| (loc.map expandLocation).getD (.targets #[] true)
+      trace[try.debug] "`simp` succeeded"
+      if (← read).config.only then
+        let tac' ← mkSimpCallStx tac stats.usedTheorems
+        -- If config has +suggestions, only return the 'only' version, not the original
+        if configHasSuggestions configStx then
+          mkTrySuggestions #[tac']
+        else
+          mkTrySuggestions #[tac, tac']
       else
-        mkTrySuggestions #[tac, tac']
-    else
-      return tac
+        return tac
   | _ => throwUnsupportedSyntax
 
 private def evalSuggestSimpAllTrace : TryTactic := fun tac => do
   match tac with
   | `(tactic| simp_all? $[!%$_bang]? $configStx:optConfig $(_discharger)? $[only%$_only]? $[[$_args,*]]?) =>
     (← getMainGoal).withContext do
-      let hasSuggestions := configHasSuggestions configStx
+      withOriginalHeartbeats do
+        let hasSuggestions := configHasSuggestions configStx
 
-      -- Get library suggestions if +suggestions is present
-      let config ← elabSimpConfig configStx (kind := .simpAll)
-      let mut argsArray : TSyntaxArray [`Lean.Parser.Tactic.simpErase, `Lean.Parser.Tactic.simpLemma] := #[]
-      if config.suggestions then
-        let suggestions ← Lean.LibrarySuggestions.select (← getMainGoal)
-        for sugg in suggestions do
-          let ident := mkIdent sugg.name
-          let candidates ← resolveGlobalConst ident
-          for candidate in candidates do
-            let arg ← `(Parser.Tactic.simpLemma| $(mkCIdentFrom ident candidate (canonical := true)):term)
-            argsArray := argsArray.push arg
+        -- Get library suggestions if +suggestions is present
+        let config ← elabSimpConfig configStx (kind := .simpAll)
+        let mut argsArray : TSyntaxArray [`Lean.Parser.Tactic.simpErase, `Lean.Parser.Tactic.simpLemma] := #[]
+        if config.suggestions then
+          let suggestions ← Lean.LibrarySuggestions.select (← getMainGoal)
+          for sugg in suggestions do
+            let ident := mkIdent sugg.name
+            let candidates ← resolveGlobalConst ident
+            for candidate in candidates do
+              let arg ← `(Parser.Tactic.simpLemma| $(mkCIdentFrom ident candidate (canonical := true)):term)
+              argsArray := argsArray.push arg
 
-      -- Build tactic with resolved suggestions for execution
-      -- If +suggestions was present, we need to create a tactic without +suggestions,
-      -- either with the resolved suggestions or without arguments if none were found
-      -- Note: We use simp_all (without ?) for execution, not simp_all?
-      let tacForExec ← if config.suggestions then
-        if argsArray.isEmpty then
-          `(tactic| simp_all)
+        -- Build tactic with resolved suggestions for execution
+        -- If +suggestions was present, we need to create a tactic without +suggestions,
+        -- either with the resolved suggestions or without arguments if none were found
+        -- Note: We use simp_all (without ?) for execution, not simp_all?
+        let tacForExec ← if config.suggestions then
+          if argsArray.isEmpty then
+            `(tactic| simp_all)
+          else
+            `(tactic| simp_all [$argsArray,*])
         else
-          `(tactic| simp_all [$argsArray,*])
-      else
-        pure tac
+          pure tac
 
-      let { ctx, simprocs, .. } ← mkSimpContext tacForExec (eraseLocal := true) (kind := .simpAll) (ignoreStarArg := true)
-      let (result?, stats) ← simpAll (← getMainGoal) ctx (simprocs := simprocs)
-      match result? with
-      | none => replaceMainGoal []
-      | some mvarId => replaceMainGoal [mvarId]
-      trace[try.debug] "`simp_all` succeeded"
-      if (← read).config.only then
-        -- Remove +suggestions from config for the output (similar to SimpTrace.lean)
-        let filteredCfg ← filterSuggestionsFromSimpConfig configStx
-        -- Convert simp_all? to simp_all for mkSimpCallStx (similar to simpTraceToSimp)
-        let tacWithoutTrace ← `(tactic| simp_all $filteredCfg:optConfig $[only%$_only]? $[[$_args,*]]?)
-        let tac' ← mkSimpCallStx tacWithoutTrace stats.usedTheorems
-        -- If config has +suggestions, only return the 'only' version, not the original
-        if hasSuggestions then
-          mkTrySuggestions #[tac']
+        let { ctx, simprocs, .. } ← mkSimpContext tacForExec (eraseLocal := true) (kind := .simpAll) (ignoreStarArg := true)
+        let (result?, stats) ← simpAll (← getMainGoal) ctx (simprocs := simprocs)
+        match result? with
+        | none => replaceMainGoal []
+        | some mvarId => replaceMainGoal [mvarId]
+        trace[try.debug] "`simp_all` succeeded"
+        if (← read).config.only then
+          -- Remove +suggestions from config for the output (similar to SimpTrace.lean)
+          let filteredCfg ← filterSuggestionsFromSimpConfig configStx
+          -- Convert simp_all? to simp_all for mkSimpCallStx (similar to simpTraceToSimp)
+          let tacWithoutTrace ← `(tactic| simp_all $filteredCfg:optConfig $[only%$_only]? $[[$_args,*]]?)
+          let tac' ← mkSimpCallStx tacWithoutTrace stats.usedTheorems
+          -- If config has +suggestions, only return the 'only' version, not the original
+          if hasSuggestions then
+            mkTrySuggestions #[tac']
+          else
+            mkTrySuggestions #[tac, tac']
         else
-          mkTrySuggestions #[tac, tac']
-      else
-        return tac
+          return tac
   | _ => throwUnsupportedSyntax
 
 
@@ -672,7 +701,8 @@ where
 private partial def evalSuggestDefault (tac : TSyntax `tactic) : TryTacticM (TSyntax `tactic) := do
   let kind := tac.raw.getKind
   match (← getEvalFns kind) with
-  | [] => evalSuggestAtomic tac -- lift regular tactic
+  | [] =>
+    withOriginalHeartbeats (evalSuggestAtomic tac)
   | evalFns => eval (← Tactic.saveState) evalFns #[]
 where
   throwExs (failures : Array EvalTacticFailure) : TryTacticM (TSyntax `tactic) := do
@@ -726,7 +756,7 @@ private partial def evalSuggestImpl : TryTactic := fun tac => do
       else if k == ``Parser.Tactic.simpAllTrace then
         evalSuggestSimpAllTrace tac
       else if k == ``Parser.Tactic.exact? then
-        evalSuggestExact
+        withOriginalHeartbeats evalSuggestExact
       else
         evalSuggestDefault tac
       if (← read).terminal then
@@ -756,10 +786,10 @@ private def addSuggestions (tk : Syntax) (s : Array Tactic.TryThis.Suggestion) :
   else
     Tactic.TryThis.addSuggestions tk (s.map fun stx => stx) (origSpan? := (← getRef))
 
-def evalAndSuggest (tk : Syntax) (tac : TSyntax `tactic) (config : Try.Config := {}) : TacticM Unit := do
+def evalAndSuggest (tk : Syntax) (tac : TSyntax `tactic) (originalMaxHeartbeats : Nat) (config : Try.Config := {}) : TacticM Unit := do
   let initialLog ← Core.getMessageLog
   let tac' ← try
-    evalSuggest tac |>.run { terminal := true, root := tac, config }
+    evalSuggest tac |>.run { terminal := true, root := tac, config, originalMaxHeartbeats }
   catch _ =>
     throwEvalAndSuggestFailed config
   -- Restore message log to suppress "Try this" messages from intermediate tactic executions
@@ -920,10 +950,10 @@ private def wrapSuggestionWithBy (sugg : Tactic.TryThis.Suggestion) : TacticM Ta
   | _ => return sugg
 
 /-- Version of `evalAndSuggest` that wraps tactic suggestions with `by` for term mode. -/
-private def evalAndSuggestWithBy (tk : Syntax) (tac : TSyntax `tactic) (config : Try.Config) : TacticM Unit := do
+private def evalAndSuggestWithBy (tk : Syntax) (tac : TSyntax `tactic) (originalMaxHeartbeats : Nat) (config : Try.Config) : TacticM Unit := do
   let initialLog ← Core.getMessageLog
   let tac' ← try
-    evalSuggest tac |>.run { terminal := true, root := tac, config }
+    evalSuggest tac |>.run { terminal := true, root := tac, config, originalMaxHeartbeats }
   catch _ =>
     throwEvalAndSuggestFailed config
   -- Restore message log to suppress "Try this" messages from intermediate tactic executions
@@ -943,13 +973,15 @@ private def evalAndSuggestWithBy (tk : Syntax) (tac : TSyntax `tactic) (config :
   match stx with
   | `(tactic| try?%$tk $config:optConfig) => Tactic.focus do withMainContext do
     let config ← elabTryConfig config
-    let goal ← getMainGoal
-    let info ← Try.collect goal config
-    let tacStx ← mkTryEvalSuggestStx goal info
-    if config.wrapWithBy then
-      evalAndSuggestWithBy tk tacStx config
-    else
-      evalAndSuggest tk tacStx config
+    let originalMaxHeartbeats ← getMaxHeartbeats
+    withUnlimitedHeartbeats do
+      let goal ← getMainGoal
+      let info ← Try.collect goal config
+      let stx ← mkTryEvalSuggestStx goal info
+      if config.wrapWithBy then
+        evalAndSuggestWithBy tk stx originalMaxHeartbeats config
+      else
+        evalAndSuggest tk stx originalMaxHeartbeats config
   | _ => throwUnsupportedSyntax
 
 end Lean.Elab.Tactic.Try

tests/lean/run/grind_eval_suggest.lean

@@ -5,7 +5,7 @@ public import Std.Tactic.BVDecide
 
 open Lean Elab Tactic Try
 elab tk:"eval_suggest" tac:tactic : tactic => do
-  evalAndSuggest tk tac
+  evalAndSuggest tk tac (originalMaxHeartbeats := 10^8)
 
 set_option hygiene false in
 macro "try_simple?" : tactic => `(tactic| eval_suggest (intros; attempt_all | rfl | (first | simp?; done | simp? +arith; done | simp_all) | grind?))

tests/lean/run/grind_try_extend.lean

@@ -24,7 +24,7 @@ example (h : False) : False := by
 elab stx:"my_try?" : tactic => do
   -- Things to try
   let toTry ← `(tactic| attempt_all | assumption | apply True | rfl)
-  evalAndSuggest stx toTry
+  evalAndSuggest stx toTry (originalMaxHeartbeats := 10^8)
 
 /--
 info: Try these:

tests/lean/run/try_heartbeats.lean

@@ -0,0 +1,88 @@
+/-
+Tests for heartbeat management in `try?`
+
+This file tests that:
+1. Each tactic execution gets a fresh heartbeat budget with the original limit
+2. Tactics that exceed heartbeat limits are filtered out from try? suggestions
+3. `attempt_all` continues past heartbeat failures to find working tactics
+-/
+
+module
+public import Lean
+public meta import Lean.Elab.Tactic.Try
+
+open Lean Meta Elab Tactic Try
+
+-- Define a custom proposition that built-in tactics won't solve
+inductive CustomGoal : Prop where
+  | mk : CustomGoal
+
+-- Create a tactic that consumes heartbeats through meta operations
+elab "expensive_meta_tactic" : tactic => do
+  for _ in [:10000] do
+    let mvar ← mkFreshExprMVar (some (mkConst ``Nat))
+    let _ ← instantiateMVars mvar
+  evalTactic (← `(tactic| exact CustomGoal.mk))
+
+-- Verify that expensive_meta_tactic actually consumes heartbeats:
+-- It times out at 100 heartbeats
+set_option maxHeartbeats 100 in
+/--
+error: (deterministic) timeout at `«tactic execution»`, maximum number of heartbeats (100) has been reached
+
+Note: Use `set_option maxHeartbeats <num>` to set the limit.
+
+Hint: Additional diagnostic information may be available using the `set_option diagnostics true` command.
+-/
+#guard_msgs in
+example : CustomGoal := by
+  expensive_meta_tactic
+
+-- It succeeds with 500 heartbeats (establishes threshold)
+set_option maxHeartbeats 500 in
+#guard_msgs in
+example : CustomGoal := by
+  expensive_meta_tactic
+
+-- Now register it as a try? suggestion along with a cheap alternative
+@[local try_suggestion 900]
+meta def expensiveSolver (_ : MVarId) (_ : Try.Info) : MetaM (Array (TSyntax `tactic)) := do
+  return #[← `(tactic| expensive_meta_tactic)]
+
+@[local try_suggestion 700]
+meta def cheapSolver (_ : MVarId) (_ : Try.Info) : MetaM (Array (TSyntax `tactic)) := do
+  return #[← `(tactic| exact CustomGoal.mk)]
+
+-- With 2 heartbeats, `try?` still works, but can't find anything.
+set_option maxHeartbeats 2 in
+/--
+error: Tactic `try?` failed: consider using `grind` manually, or `try? +missing` for partial proofs containing `sorry`
+
+⊢ CustomGoal
+-/
+#guard_msgs in
+example : CustomGoal := by
+  try?
+
+-- With 100 heartbeats
+-- expensive_meta_tactic should be filtered out from try? suggestions
+-- Only `CustomGoal.mk`` should appear
+set_option maxHeartbeats 100 in
+/--
+info: Try this:
+  [apply] exact CustomGoal.mk✝
+-/
+#guard_msgs in
+example : CustomGoal := by
+  try?
+
+-- With 500 heartbeats, both should appear.
+/--
+info: Try these:
+  [apply] expensive_meta_tactic
+  [apply] exact CustomGoal.mk✝
+-/
+#guard_msgs in
+set_option maxHeartbeats 500 in
+example : CustomGoal := by
+  try?