From e237e12478cd101915a232bfcac258ce014a0f21 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leomoura@amazon.com>
Date: Sat, 11 May 2024 00:34:09 +0200
Subject: [PATCH] refactor: add `tryCatchRuntimeEx` combinator (#4129)

see #4079
---
 src/Lean/CoreM.lean                 | 42 ++++++++++++++---------------
 src/Lean/Meta/SynthInstance.lean    |  9 +++----
 src/Lean/Meta/Tactic/Simp/Main.lean | 36 ++++++++++++-------------
 3 files changed, 42 insertions(+), 45 deletions(-)

diff --git a/src/Lean/CoreM.lean b/src/Lean/CoreM.lean
index c92760026c..fc1923be06 100644
--- a/src/Lean/CoreM.lean
+++ b/src/Lean/CoreM.lean
@@ -79,12 +79,6 @@ structure Context where
   maxHeartbeats  : Nat := getMaxHeartbeats options
   currMacroScope : MacroScope := firstFrontendMacroScope
   /--
-  If `catchRuntimeEx = false`, then given `try x catch ex => h ex`,
-  an runtime exception occurring in `x` is not handled by `h`.
-  Recall that runtime exceptions are `maxRecDepth` or `maxHeartbeats`.
-  -/
-  catchRuntimeEx : Bool := false
-  /--
   If `diag := true`, different parts of the system collect diagnostics.
   Use the `set_option diag true` to set it to true.
   -/
@@ -415,30 +409,36 @@ in these monads, but on `CommandElabM`. See issues #2775 and #2744 as well as `M
   try
     x
   catch ex =>
-    if ex.isRuntime && !(← read).catchRuntimeEx then
-      throw ex
+    if ex.isRuntime then
+      throw ex -- We should use `tryCatchRuntimeEx` for catching runtime exceptions
     else
       h ex
 
+@[inline] protected def Core.tryCatchRuntimeEx (x : CoreM α) (h : Exception → CoreM α) : CoreM α := do
+  try
+    x
+  catch ex =>
+    h ex
+
 instance : MonadExceptOf Exception CoreM where
   throw    := throw
   tryCatch := Core.tryCatch
 
-@[inline] def Core.withCatchingRuntimeEx (flag : Bool) (x : CoreM α) : CoreM α :=
-  withReader (fun ctx => { ctx with catchRuntimeEx := flag }) x
+class MonadRuntimeException (m : Type → Type) where
+  tryCatchRuntimeEx (body : m α) (handler : Exception → m α) : m α
+
+export MonadRuntimeException (tryCatchRuntimeEx)
+
+instance : MonadRuntimeException CoreM where
+  tryCatchRuntimeEx := Core.tryCatchRuntimeEx
+
+@[inline] instance [MonadRuntimeException m] : MonadRuntimeException (ReaderT ρ m) where
+  tryCatchRuntimeEx := fun x c r => tryCatchRuntimeEx (x r) (fun e => (c e) r)
+
+@[inline] instance [MonadRuntimeException m] : MonadRuntimeException (StateRefT' ω σ m) where
+  tryCatchRuntimeEx := fun x c s => tryCatchRuntimeEx (x s) (fun e => c e s)
 
 @[inline] def mapCoreM [MonadControlT CoreM m] [Monad m] (f : forall {α}, CoreM α → CoreM α) {α} (x : m α) : m α :=
   controlAt CoreM fun runInBase => f <| runInBase x
 
-/--
-Execute `x` with `catchRuntimeEx = flag`. That is, given `try x catch ex => h ex`,
-if `x` throws a runtime exception, the handler `h` will be invoked if `flag = true`
-Recall that
--/
-@[inline] def withCatchingRuntimeEx [MonadControlT CoreM m] [Monad m] (x : m α) : m α :=
-  mapCoreM (Core.withCatchingRuntimeEx true) x
-
-@[inline] def withoutCatchingRuntimeEx [MonadControlT CoreM m] [Monad m] (x : m α) : m α :=
-  mapCoreM (Core.withCatchingRuntimeEx false) x
-
 end Lean
diff --git a/src/Lean/Meta/SynthInstance.lean b/src/Lean/Meta/SynthInstance.lean
index 6a82fe7c57..feaace0c8c 100644
--- a/src/Lean/Meta/SynthInstance.lean
+++ b/src/Lean/Meta/SynthInstance.lean
@@ -633,12 +633,9 @@ def main (type : Expr) (maxResultSize : Nat) : MetaM (Option AbstractMVarsResult
      let action : SynthM (Option AbstractMVarsResult) := do
        newSubgoal (← getMCtx) key mvar Waiter.root
        synth
-     -- TODO: it would be nice to have a nice notation for the following idiom
-     withCatchingRuntimeEx
-       try
-         withoutCatchingRuntimeEx do
-           action.run { maxResultSize := maxResultSize, maxHeartbeats := getMaxHeartbeats (← getOptions) } |>.run' {}
-       catch ex =>
+     tryCatchRuntimeEx
+       (action.run { maxResultSize := maxResultSize, maxHeartbeats := getMaxHeartbeats (← getOptions) } |>.run' {})
+       fun ex =>
          if ex.isRuntime then
            throwError "failed to synthesize{indentExpr type}\n{ex.toMessageData}"
          else
diff --git a/src/Lean/Meta/Tactic/Simp/Main.lean b/src/Lean/Meta/Tactic/Simp/Main.lean
index a553e22d86..2a8bf686e0 100644
--- a/src/Lean/Meta/Tactic/Simp/Main.lean
+++ b/src/Lean/Meta/Tactic/Simp/Main.lean
@@ -662,30 +662,30 @@ def main (e : Expr) (ctx : Context) (stats : Stats := {}) (methods : Methods :=
     trace[Meta.Tactic.simp.numSteps] "{s.numSteps}"
     return (r, { s with })
 where
-  simpMain (e : Expr) : SimpM Result := withCatchingRuntimeEx do
-    try
-      withoutCatchingRuntimeEx <| simp e
-    catch ex =>
-      reportDiag (← get).diag
-      if ex.isRuntime then
-        throwNestedTacticEx `simp ex
-      else
-        throw ex
+  simpMain (e : Expr) : SimpM Result :=
+    tryCatchRuntimeEx
+      (simp e)
+      fun ex => do
+        reportDiag (← get).diag
+        if ex.isRuntime then
+          throwNestedTacticEx `simp ex
+        else
+          throw ex
 
 def dsimpMain (e : Expr) (ctx : Context) (stats : Stats := {}) (methods : Methods := {}) : MetaM (Expr × Stats) := do
   withSimpContext ctx do
     let (r, s) ← dsimpMain e methods.toMethodsRef ctx |>.run { stats with }
     pure (r, { s with })
 where
-  dsimpMain (e : Expr) : SimpM Expr := withCatchingRuntimeEx do
-    try
-      withoutCatchingRuntimeEx <| dsimp e
-    catch ex =>
-      reportDiag (← get).diag
-      if ex.isRuntime then
-        throwNestedTacticEx `simp ex
-      else
-        throw ex
+  dsimpMain (e : Expr) : SimpM Expr :=
+    tryCatchRuntimeEx
+      (dsimp e)
+      fun ex => do
+        reportDiag (← get).diag
+        if ex.isRuntime then
+          throwNestedTacticEx `simp ex
+        else
+          throw ex
 
 end Simp
 open Simp (SimprocsArray Stats)