fix: improve performance of mvcgen by optimizing try (mpure_intro; trivial) (#10872)

This PR improves the performance of `mvcgen` by an optimized
implementation for `try (mpure_intro; trivial)`. This tactic sequence is
used to eagerly discharge VCs and in the process instantiates schematic
variables.

src/Lean/Elab/Tactic/Do/ProofMode/Cases.lean

@@ -6,8 +6,11 @@ Authors: Lars König, Mario Carneiro, Sebastian Graf
 module
 
 prelude
-public import Lean.Elab.Tactic.Do.ProofMode.Pure
-public import Lean.Elab.Tactic.Do.ProofMode.Intro
+public import Lean.Elab.Tactic.Do.ProofMode.MGoal
+public import Std.Tactic.Do.Syntax
+import Lean.Elab.Tactic.Do.ProofMode.Pure
+import Lean.Elab.Tactic.Do.ProofMode.Intro
+import Lean.Elab.Tactic.Do.ProofMode.Focus
 
 public section
 

src/Lean/Elab/Tactic/Do/ProofMode/Have.lean

@@ -6,8 +6,11 @@ Authors: Sebastian Graf
 module
 
 prelude
-public import Lean.Elab.Tactic.Do.ProofMode.Cases
-public import Lean.Elab.Tactic.Do.ProofMode.Specialize
+public import Std.Tactic.Do.Syntax
+public import Lean.Elab.Tactic.Basic
+import Lean.Elab.Tactic.Do.ProofMode.Focus
+import Lean.Elab.Tactic.Do.ProofMode.Cases
+import Lean.Elab.Tactic.Do.ProofMode.Specialize
 
 public section
 

src/Lean/Elab/Tactic/Do/ProofMode/Pure.lean

@@ -6,9 +6,13 @@ Authors: Lars König, Mario Carneiro, Sebastian Graf
 module
 
 prelude
-public import Std.Tactic.Do.Syntax
-public import Lean.Elab.Tactic.Do.ProofMode.Focus
-public import Lean.Elab.Tactic.Meta
+public import Lean.Elab.Tactic.Basic
+public import Lean.Elab.Tactic.Do.ProofMode.MGoal
+import Std.Tactic.Do.Syntax
+import Lean.Elab.Tactic.Meta
+import Lean.Elab.Tactic.Do.ProofMode.Basic
+import Lean.Elab.Tactic.Do.ProofMode.Focus
+import Lean.Meta.Tactic.Rfl
 
 public section
 
@@ -21,14 +25,16 @@ open Lean Elab Tactic Meta
 -- It will provide a proof for Q ∧ H ⊢ₛ T
 -- if `k` produces a proof for Q ⊢ₛ T that may range over a pure proof h : φ.
 -- It calls `k` with the φ in H = ⌜φ⌝ and a proof `h : φ` thereof.
-def mPureCore (σs : Expr) (hyp : Expr) (name : TSyntax ``binderIdent)
-  (k : Expr /-φ:Prop-/ → Expr /-h:φ-/ → MetaM (α × MGoal × Expr)) : MetaM (α × MGoal × Expr) := do
+def mPureCore
+  [Monad m] [MonadControlT MetaM m] [MonadLiftT MetaM m]
+  (σs : Expr) (hyp : Expr) (name : TSyntax ``binderIdent)
+  (k : Expr /-φ:Prop-/ → Expr /-h:φ-/ → m (α × MGoal × Expr)) : m (α × MGoal × Expr) := do
   let u ← mkFreshLevelMVar
   let φ ← mkFreshExprMVar (mkSort .zero)
   let inst ← synthInstance (mkApp3 (mkConst ``IsPure [u]) σs hyp φ)
-  let (name, ref) ← getFreshHypName name
+  let (name, ref) ← liftMetaM <| getFreshHypName name
   withLocalDeclD name φ fun h => do
-    addLocalVarInfo ref (← getLCtx) h φ
+    addLocalVarInfo ref (← liftMetaM <| getLCtx) h φ
     let (a, goal, prf /- : goal.toExpr -/) ← k φ h
     let prf ← mkLambdaFVars #[h] prf
     let prf := mkApp7 (mkConst ``Pure.thm [u]) σs goal.hyps hyp goal.target φ inst prf
@@ -38,10 +44,8 @@ def mPureCore (σs : Expr) (hyp : Expr) (name : TSyntax ``binderIdent)
 @[builtin_tactic Lean.Parser.Tactic.mpure]
 def elabMPure : Tactic
   | `(tactic| mpure $hyp) => do
-    let mvar ← getMainGoal
+    let (mvar, goal) ← mStartMainGoal
     mvar.withContext do
-    let g ← instantiateMVars <| ← mvar.getType
-    let some goal := parseMGoal? g | throwError "not in proof mode"
     let res ← goal.focusHypWithInfo hyp
     let (m, _new_goal, prf) ← mPureCore goal.σs res.focusHyp (← `(binderIdent| $hyp:ident)) fun _ _ => do
       let goal := res.restGoal goal
@@ -52,8 +56,87 @@ def elabMPure : Tactic
     replaceMainGoal [m.mvarId!]
   | _ => throwUnsupportedSyntax
 
-def MGoal.triviallyPure (goal : MGoal) : OptionT MetaM Expr := do
+-- NB: We do not use MVarId.intro because that would mean we require all callers to supply an MVarId.
+-- This function only knows about the hypothesis H=⌜φ⌝ to destruct.
+-- It will provide a proof for Q ∧ H ⊢ₛ T
+-- if `k` produces a proof for Q ⊢ₛ T that may range over a pure proof h : φ.
+-- It calls `k` with the φ in H = ⌜φ⌝ and a proof `h : φ` thereof.
+def mPureIntroCore [Monad m] [MonadLiftT MetaM m]
+  (goal : MGoal)
+  (k : Expr /-φ:Prop-/ → m (α × Expr)) : m (α × Expr) := do
+  let φ ← mkFreshExprMVar (mkSort .zero)
+  let inst ← synthInstance (mkApp3 (mkConst ``IsPure [goal.u]) goal.σs goal.target φ)
+  let (a, hφ) ← k φ
+  let prf := mkApp6 (mkConst ``Pure.intro [goal.u]) goal.σs goal.hyps goal.target φ inst hφ
+  return (a, prf)
+
+@[builtin_tactic Lean.Parser.Tactic.mpureIntro]
+def elabMPureIntro : Tactic
+  | `(tactic| mpure_intro) => do
+    let (mvar, goal) ← mStartMainGoal
+    mvar.withContext do
+    let (mv, prf) ← mPureIntroCore goal fun φ => do
+      let m ← mkFreshExprSyntheticOpaqueMVar φ (← mvar.getTag)
+      return (m.mvarId!, m)
+    mvar.assign prf
+    replaceMainGoal [mv]
+  | _ => throwUnsupportedSyntax
+
+partial def _root_.Lean.MVarId.applyRflAndAndIntro (mvar : MVarId) : MetaM Unit := do
+  -- The target might look like `(⌜?n = nₛ ∧ ?m = b⌝ s).down`, which we reduce to
+  -- `?n = nₛ ∧ ?m = b` by `whnfD`.
+  -- (Recall that `⌜s = 4⌝ s` is `SPred.pure (σs:=[Nat]) (s = 4) s` and `SPred.pure` is
+  -- semi-reducible.)
+  let ty ← whnfD (← mvar.getType)
+  trace[Elab.Tactic.Do.spec] "whnf: {ty}"
+  if ty.isAppOf ``True then
+    mvar.assign (mkConst ``True.intro)
+  else if let some (lhs, rhs) := ty.app2? ``And then
+    let hlhs ← mkFreshExprMVar lhs
+    let hrhs ← mkFreshExprMVar rhs
+    applyRflAndAndIntro hlhs.mvarId!
+    applyRflAndAndIntro hrhs.mvarId!
+    mvar.assign (mkApp4 (mkConst ``And.intro) lhs rhs hlhs hrhs)
+  else
+    mvar.setType ty
+    mvar.applyRfl
+
+def MGoal.pureRflAndAndIntro (goal : MGoal) : OptionT MetaM Expr := do
+  trace[Elab.Tactic.Do.spec] "pureRflAndAndIntro: {goal.target}"
+  try
+    let (_, prf) ← mPureIntroCore goal fun φ => do
+      trace[Elab.Tactic.Do.spec] "discharge? {φ}"
+      let m ← mkFreshExprMVar φ
+      m.mvarId!.applyRflAndAndIntro
+      trace[Elab.Tactic.Do.spec] "discharged: {φ}"
+      return ((), m)
+    return prf
+  catch _ => failure
+
+def MGoal.pureTrivial (goal : MGoal) : OptionT MetaM Expr := do
+  try
+    let (_, prf) ← mPureIntroCore goal fun φ => do
+      let m ← mkFreshExprMVar φ
+      try
+        -- First try to use rfl and And.intro directly.
+        -- This is more efficient than to elaborate the `trivial` tactic.
+        m.mvarId!.applyRflAndAndIntro
+      catch _ =>
+        let ([], _) ← runTactic m.mvarId! (← `(tactic| trivial))
+          | failure
+      return ((), m)
+    return prf
+  catch _ => failure
+
+/-
+def MGoal.pureRfl (goal : MGoal) : OptionT MetaM Expr := do
   let mv ← mkFreshExprMVar goal.toExpr
-  let ([], _) ← try runTactic mv.mvarId! (← `(tactic| apply $(mkIdent ``Std.Do.SPred.Tactic.Pure.intro); trivial)) catch _ => failure
+  let ([], _) ← try runTactic mv.mvarId! (← `(tactic| apply $(mkIdent ``Std.Do.SPred.Tactic.Pure.intro); rfl)) catch _ => failure
     | failure
   return mv
+def MGoal.pureRfl (goal : MGoal) : OptionT MetaM Expr := do
+  let mv ← mkFreshExprMVar goal.toExpr
+  let ([], _) ← try runTactic mv.mvarId! (← `(tactic| apply $(mkIdent ``Std.Do.SPred.Tactic.Pure.intro); rfl)) catch _ => failure
+    | failure
+  return mv
+-/

src/Lean/Elab/Tactic/Do/ProofMode/Specialize.lean

@@ -6,9 +6,11 @@ Authors: Lars König, Mario Carneiro, Sebastian Graf
 module
 
 prelude
-public import Lean.Elab.Tactic.Do.ProofMode.Basic
-public import Lean.Elab.Tactic.Do.ProofMode.Pure
 public import Lean.Elab.Tactic.ElabTerm
+public import Lean.Elab.Tactic.Do.ProofMode.MGoal
+import Lean.Elab.Tactic.Do.ProofMode.Basic
+import Lean.Elab.Tactic.Do.ProofMode.Focus
+import Lean.Elab.Tactic.Do.ProofMode.Pure
 
 public section
 

src/Lean/Elab/Tactic/Do/Spec.lean

@@ -6,11 +6,15 @@ Authors: Lars König, Mario Carneiro, Sebastian Graf
 module
 
 prelude
+public import Lean.Elab.Tactic.Do.Attr
+public import Lean.Elab.Tactic.Do.ProofMode.MGoal
+
+-- All these should become private imports in the future:
+import Std.Tactic.Do.Syntax
 public import Lean.Elab.Tactic.Do.ProofMode.Intro
 public import Lean.Elab.Tactic.Do.ProofMode.Pure
 public import Lean.Elab.Tactic.Do.ProofMode.Frame
 public import Lean.Elab.Tactic.Do.ProofMode.Assumption
-public import Lean.Elab.Tactic.Do.Attr
 
 namespace Lean.Elab.Tactic.Do
 open Lean Elab Tactic Meta
@@ -82,7 +86,7 @@ mutual
 partial def dischargePostEntails (α : Expr) (ps : Expr) (Q : Expr) (Q' : Expr) (goalTag : Name) : n Expr := do
   -- Often, Q' is fully instantiated while Q contains metavariables. Try refl
   let u ← liftMetaM <| getLevel α >>= decLevel
-  if (← isDefEq Q Q') then
+  if (← withDefault <| isDefEqGuarded Q Q') then
     return mkApp3 (mkConst ``PostCond.entails.refl [u]) α ps Q'
   let Q ← whnfR Q
   let Q' ← whnfR Q'
@@ -104,11 +108,11 @@ partial def dischargePostEntails (α : Expr) (ps : Expr) (Q : Expr) (Q' : Expr)
 partial def dischargeFailEntails (u : Level) (ps : Expr) (Q : Expr) (Q' : Expr) (goalTag : Name) : n Expr := do
   if ps.isAppOf ``PostShape.pure then
     return mkConst ``True.intro
-  if ← isDefEq Q Q' then
+  if ← withDefault <| isDefEqGuarded Q Q' then
     return mkApp2 (mkConst ``ExceptConds.entails.refl [u]) ps Q
-  if ← isDefEq Q (mkApp (mkConst ``ExceptConds.false [u]) ps) then
+  if ← withDefault <| isDefEqGuarded Q (mkApp (mkConst ``ExceptConds.false [u]) ps) then
     return mkApp2 (mkConst ``ExceptConds.entails_false [u]) ps Q'
-  if ← isDefEq Q' (mkApp (mkConst ``ExceptConds.true [u]) ps) then
+  if ← withDefault <| isDefEqGuarded Q' (mkApp (mkConst ``ExceptConds.true [u]) ps) then
     return mkApp2 (mkConst ``ExceptConds.entails_true [u]) ps Q
   -- the remaining cases are recursive.
   if let some (_σ, ps) := ps.app2? ``PostShape.arg then
@@ -129,18 +133,24 @@ partial def dischargeFailEntails (u : Level) (ps : Expr) (Q : Expr) (Q' : Expr)
   mkFreshExprSyntheticOpaqueMVar (mkApp3 (mkConst ``ExceptConds.entails [u]) ps Q Q') goalTag
 end
 
-def dischargeMGoal (goal : MGoal) (goalTag : Name) : n Expr := do
+def dischargeMGoal (goal : MGoal) (goalTag : Name) (tryTrivial : Bool) : n Expr := do
   liftMetaM <| do trace[Elab.Tactic.Do.spec] "dischargeMGoal: {goal.target}"
   -- simply try one of the assumptions for now. Later on we might want to decompose conjunctions etc; full xsimpl
   -- The `withDefault` ensures that a hyp `⌜s = 4⌝` can be used to discharge `⌜s = 4⌝ s`.
   -- (Recall that `⌜s = 4⌝ s` is `SPred.pure (σs:=[Nat]) (s = 4) s` and `SPred.pure` is
   -- semi-reducible.)
-  -- We also try `mpure_intro; trivial` through `goal.triviallyPure` here because later on an
-  -- assignment like `⌜s = ?c⌝` becomes impossible to discharge because `?c` will get abstracted
-  -- over local bindings that depend on synthetic opaque MVars (such as loop invariants), and then
-  -- the type of the new `?c` will not be defeq to itself. A bug, but we need to work around it for
-  -- now.
-  let some prf ← liftMetaM (withDefault <| goal.assumption <|> goal.assumptionPure <|> goal.triviallyPure)
+  -- We also try `mpure_intro; trivial` through `goal.pureTrivial` here because later on an
+  -- assignment like `⌜s = ?c ∧ s₂ = ?d⌝` becomes impossible to discharge because `?c` will get
+  -- abstracted over local bindings that depend on synthetic opaque MVars (such as loop invariants),
+  -- and then the type of the new `?c` will not be defeq to itself. A bug, but we need to work
+  -- around it for now.
+  -- Even when we *do not* want to `tryTrivial`, we still use `goal.pureRflAndAndIntro` because
+  -- it might assign metavariables as above.
+  let some prf ← liftMetaM <| -- withDefault <|
+      if tryTrivial then
+        goal.pureTrivial <|> goal.assumption <|> goal.assumptionPure
+      else
+        goal.pureRflAndAndIntro
     | mkFreshExprSyntheticOpaqueMVar goal.toExpr goalTag
   liftMetaM <| do trace[Elab.Tactic.Do.spec] "proof: {prf}"
   return prf
@@ -148,7 +158,7 @@ def dischargeMGoal (goal : MGoal) (goalTag : Name) : n Expr := do
 /--
   Returns the proof and the list of new unassigned MVars.
 -/
-public def mSpec (goal : MGoal) (elabSpecAtWP : Expr → n SpecTheorem) (goalTag : Name) : n Expr := do
+public def mSpec (goal : MGoal) (elabSpecAtWP : Expr → n SpecTheorem) (goalTag : Name) (tryTrivial : Bool := true) : n Expr := do
   -- First instantiate `fun s => ...` in the target via repeated `mintro ∀s`.
   mIntroForallN goal goal.target.consumeMData.getNumHeadLambdas fun goal => do
   -- Elaborate the spec for the wp⟦e⟧ app in the target
@@ -189,7 +199,7 @@ public def mSpec (goal : MGoal) (elabSpecAtWP : Expr → n SpecTheorem) (goalTag
   let_expr f@Triple m ps instWP α prog P Q := specTy
     | liftMetaM <| throwError "target not a Triple application {specTy}"
   let wp' := mkApp5 (mkConst ``WP.wp f.constLevels!) m ps instWP α prog
-  unless (← withAssignableSyntheticOpaque <| isDefEq wp wp') do
+  unless (← withAssignableSyntheticOpaque <| isDefEqGuarded wp wp') do
     Term.throwTypeMismatchError none wp wp' spec
 
   -- Try synthesizing synthetic MVars. We don't have the convenience of `TermElabM`, hence
@@ -215,7 +225,7 @@ public def mSpec (goal : MGoal) (elabSpecAtWP : Expr → n SpecTheorem) (goalTag
   -- Often P or Q are schematic (i.e. an MVar app). Try to solve by rfl.
   -- We do `fullApproxDefEq` here so that `constApprox` is active; this is useful in
   -- `need_const_approx` of `doLogicTests.lean`.
-  let (HPRfl, QQ'Rfl) ← fullApproxDefEq <| do
+  let (HPRfl, QQ'Rfl) ← withDefault <| fullApproxDefEq <| do
     return (← isDefEqGuarded P goal.hyps, ← isDefEqGuarded Q Q')
 
   -- Discharge the validity proof for the spec if not rfl
@@ -224,7 +234,7 @@ public def mSpec (goal : MGoal) (elabSpecAtWP : Expr → n SpecTheorem) (goalTag
     -- let P := (← reduceProjBeta? P).getD P
     -- Try to avoid creating a longer name if the postcondition does not need to create a goal
     let tag := if !QQ'Rfl then goalTag ++ `pre else goalTag
-    let HPPrf ← dischargeMGoal { goal with target := P } tag
+    let HPPrf ← dischargeMGoal { goal with target := P } tag tryTrivial
     prePrf := mkApp6 (mkConst ``SPred.entails.trans [u]) goal.σs goal.hyps P goal.target HPPrf
 
   -- Discharge the entailment on postconditions if not rfl

src/Lean/Elab/Tactic/Do/VCGen.lean

@@ -145,7 +145,7 @@ where
         -- This is so that `mSpec` can frame hypotheses involving uninstantiated loop invariants.
         -- It is absolutely crucial that we do not lose these hypotheses in the inductive step.
         collectFreshMVars <| mIntroForallN goal (← TypeList.length goal.σs) fun goal =>
-          withDefault <| mSpec goal (fun _wp  => return specThm) name
+          mSpec goal (fun _wp  => return specThm) name (tryTrivial := false)
       catch ex =>
         trace[Elab.Tactic.Do.vcgen] "Failed to find spec for {wp}. Trying simp. Reason: {ex.toMessageData}"
         -- Last resort: Simp and try again

src/Std/Do/SPred/DerivedLaws.lean

@@ -236,7 +236,7 @@ theorem Pure.thm {P Q T : SPred σs} {φ : Prop} [IsPure Q φ] (h : φ → P ⊢
   · intro hp
     exact and_elim_l.trans (h hp)
 /-- A generalization of `pure_intro` exploiting `IsPure`. -/
-theorem Pure.intro {P Q : SPred σs} {φ : Prop} [IsPure Q φ] (hp : φ) : P ⊢ₛ Q := (pure_intro hp).trans IsPure.to_pure.mpr
+theorem Pure.intro {P Q : SPred σs} {φ : Prop} [IsPure Q φ] (hφ : φ) : P ⊢ₛ Q := (pure_intro hφ).trans IsPure.to_pure.mpr
 theorem Revert.revert {P Q H T : SPred σs} (hfoc : P ⊣⊢ₛ Q ∧ H) (h : Q ⊢ₛ H → T) : P ⊢ₛ T := hfoc.mp.trans (imp_elim h)
 theorem Specialize.imp_stateful {P P' Q R : SPred σs}
     (hrefocus : P ∧ (Q → R) ⊣⊢ₛ P' ∧ Q) : P ∧ (Q → R) ⊢ₛ P ∧ R := by

src/Std/Tactic/Do/Syntax.lean

@@ -105,8 +105,7 @@ syntax (name := mleft) "mleft" : tactic
 syntax (name := mpure) "mpure" colGt ident : tactic
 
 @[tactic_alt Lean.Parser.Tactic.mpureIntroMacro]
-macro (name := mpureIntro) "mpure_intro" : tactic =>
-  `(tactic| apply $(mkIdent ``Std.Do.SPred.Tactic.Pure.intro))
+syntax (name := mpureIntro) "mpure_intro" : tactic
 
 @[tactic_alt Lean.Parser.Tactic.mrenameIMacro]
 syntax (name := mrenameI) "mrename_i" (ppSpace colGt binderIdent)+ : tactic

tests/lean/run/doLogicTests.lean

@@ -429,6 +429,8 @@ theorem add_unfold [Monad m] [WPMonad m sh] :
 
 theorem mkFreshPair_triple : ⦃⌜True⌝⦄ mkFreshPair ⦃⇓ (a, b) => ⌜a ≠ b⌝⦄ := by
   mvcgen -elimLets +trivial [mkFreshPair]
+  -- this tests whether `mSpec` immediately discharges by `rfl` and `And.intro` and in the process
+  -- eagerly instantiates some schematic variables.
   simp_all
 
 theorem sum_loop_spec :