try, salvage, catch

Stream: general

[Thomas Murrills]

Ever wanted to grab some data dependent (only) on the monad state after you've thrown an error in a try block? Well...probably not, actually. But if you ever find yourself in that specific predicament, here's some code that lets you write the following in most of Lean's major monads. :)

  try
    ...
  salvage ex =>
    ...
    -- salvage `b : β` by working in the monad after `ex` has been thrown,
    -- but before restoring the state with `catch`
    return b
  catch ex (b : β) =>
    <do something with b!>
    ...

(I'm not planning on PRing this—this is just for fun. Thought I'd share. :) )

code

import Lean

section

set_option autoImplicit true

open Lean Elab Tactic

section instances

/-!
Everything is structurally exactly the same as for `tryCatch`.
Note that we don't include `throw`—the `throw` in `MonadExcept` is enough.
-/

class MonadSalvageOf.{u, v, w} (ε : semiOutParam (Type u)) (m : Type v → Type w) :
    Type (max (max u (v + 1)) w) where
  trySalvageCatch {α : Type v} (x : m α) {β : Type v} (salvage : ε → m β)
      (handle : ε → β → m α) : m α

class MonadSalvage.{u, v, w} (ε : outParam (Type u)) (m : Type v → Type w) : Type (max (max u (v + 1)) w) where
  trySalvageCatch {α : Type v} (x : m α) {β : Type v} (salvage : ε → m β)
      (handle : ε → β → m α) : m α

export MonadSalvage (trySalvageCatch)

abbrev trySalvageCatchThe (ε : Type u) {m : Type v → Type w} [MonadSalvageOf ε m]
    {α : Type v} (x : m α) {β} (salvage : ε → m β) (handle : ε → β → m α) : m α :=
  MonadSalvageOf.trySalvageCatch x salvage handle

instance (ε : outParam (Type u)) (m : Type v → Type w) [MonadSalvageOf ε m] : MonadSalvage ε m where
  trySalvageCatch := trySalvageCatchThe ε

namespace EStateM

@[always_inline, inline]
def trySalvageCatch {δ} [Backtrackable δ σ] {α} (x : EStateM ε σ α) {β} (salvage : ε → EStateM ε σ β) (handle : ε → β → EStateM ε σ α) : EStateM ε σ α := fun s =>
  let d := Backtrackable.save s
  match x s with
  | Result.error e s => match salvage e s with
    | Result.ok b s' => handle e b (Backtrackable.restore s' d)
    | Result.error e' s' => Result.error e' s' -- Errors in `salvage` are preserved and uncaught.
  | ok               => ok

instance {δ} [Backtrackable δ σ] : MonadSalvageOf ε (EStateM ε σ) where
  trySalvageCatch := trySalvageCatch

end EStateM

instance : MonadSalvageOf ε (EIO ε) := inferInstanceAs (MonadSalvageOf ε (EStateM ε IO.RealWorld))

@[always_inline]
instance (ε) [MonadSalvageOf ε m] : MonadSalvageOf ε (ReaderT ρ m) where
  trySalvageCatch := fun x _ s c r =>
    trySalvageCatchThe ε (x r) (fun e => (s e) r) (fun e b => (c e b) r)

@[always_inline]
instance (ε) [MonadSalvageOf ε m] : MonadSalvageOf ε (StateRefT' ω σ m) where
  trySalvageCatch := fun x _ s c t =>
    trySalvageCatchThe ε (x t) (fun e => (s e) t) (fun e b => (c e b) t)

@[inline] protected def Core.trySalvageCatch (x : CoreM α) {β} (s : Exception → CoreM β) (h : Exception → β → CoreM α) : CoreM α := do
  trySalvageCatch x s fun ex b => do
    if ex.isRuntime && !(← read).catchRuntimeEx then
      throw ex
    else
      h ex b

instance : MonadSalvageOf Exception CoreM where
  trySalvageCatch := Core.trySalvageCatch

namespace Tactic

@[inline] protected def trySalvageCatch {α} (x : TacticM α) {β} (s : Exception → TacticM β)
    (h : Exception → β → TacticM α) (restoreInfo := false): TacticM α := do
  let t ← saveState
  trySalvageCatch x s fun ex b => do t.restore restoreInfo; h ex b

instance : MonadSalvage Exception TacticM where
  trySalvageCatch := fun x _ s h => Tactic.trySalvageCatch x s h (restoreInfo := false)

end Tactic

end instances

/-- Example of a 'useful' function built on `trySalvageCatch`. -/
def tryCatchSalvagingState {m : Type → Type} [Monad m] [MonadBacktrack σ m]
    [MonadSalvage ε m] (x : m α) (h : ε → σ → m α) : m α :=
  trySalvageCatch x (fun _ => saveState) h

end

section parser

open Lean Meta Parser Term

def doSalvage := leading_parser
  ppDedent ppLine >> atomic ("salvage " >> funBinder) >> optType >> darrow >> doSeq

def doCatchSalvage := leading_parser
  ppDedent ppLine >> atomic ("catch " >> funBinder >> funBinder) >> darrow >> doSeq

@[doElem_parser] def doTrySalvage := leading_parser
  "try " >> doSeq >> doSalvage >> many (doCatchSalvage <|> doCatchMatch) >> Parser.optional doFinally

macro_rules
| `(doTrySalvage| try $seq1 salvage $bi $[: $type]? => $seq2 catch $biex $bis => $seq3) =>
    `(doElem|trySalvageCatch
      (do $seq1)
      (fun $bi:funBinder $[: $type]? => do $seq2)
      (fun $biex $bis => do $seq3))
| `(doTrySalvage| try $seq1 salvage $bi $[: $type]? => $seq2 catch $alts:matchAlts) =>
    `(doElem|trySalvageCatch
      (do $seq1)
      (fun $bi:funBinder $[: $type]? => do $seq2)
      (fun $alts:matchAlts))

end parser

section

open Lean Meta

elab "test" : tactic => do
  try
    let _ ← mkFreshExprMVar (Expr.const ``Nat []) (userName := `foo)
    throwError "error!"
  salvage _ => getMCtx
  catch _ mctx =>
    if let .some foo := (← getMCtx).findUserName? `foo then
      throwError "somehow, {foo} survived"
    let .some foo := mctx.findUserName? `foo
      | throwError "could not find ?foo"
    let mvar := mctx.getDecl foo
    logInfo m!"found ?foo : {mvar.type}"

example : True := by
  test -- found ?foo : Nat
  trivial

end