chore: fix bug in rw? (#6088)

This fixes a problem I encountered in another tactic, where the rewrites are not being tried with the original `MetavarContext`. 



Co-authored-by: Kevin Buzzard <k.buzzard@imperial.ac.uk>
Co-authored-by: Joachim Breitner <mail@joachim-breitner.de>
Co-authored-by: Chris Hughes <chrishughes24@gmail.com>
Co-authored-by: Eric Wieser <wieser.eric@gmail.com>
Co-authored-by: Anatole Dedecker <anatolededecker@gmail.com>
Co-authored-by: Thomas Browning <tb65536@uw.edu>
Co-authored-by: Oliver Nash <github@olivernash.org>
Co-authored-by: sgouezel <sebastien.gouezel@univ-rennes1.fr>
Co-authored-by: Alex J Best <alex.j.best@gmail.com>
Co-authored-by: Xavier Roblot <46200072+xroblot@users.noreply.github.com>
Co-authored-by: michaellee94 <michael_lee1@brown.edu>
Co-authored-by: Yuma Mizuno <mizuno.y.aj@gmail.com>
Co-authored-by: Markus Himmel <markus@himmel-villmar.de>
Co-authored-by: Scott Morrison <scott.morrison@gmail.com>
Co-authored-by: Peiran Wu <15968905+wupr@users.noreply.github.com>
Co-authored-by: Moritz Firsching <moritz.firsching@gmail.com>
Co-authored-by: Yury G. Kudryashov <urkud@urkud.name>
Co-authored-by: Ruben Van de Velde <65514131+Ruben-VandeVelde@users.noreply.github.com>
Co-authored-by: Anne Baanen <vierkantor@vierkantor.com>
Co-authored-by: Bolton Bailey <bolton.bailey@gmail.com>
Co-authored-by: Frédéric Dupuis <dupuisf@iro.umontreal.ca>
Co-authored-by: Moritz Doll <moritz.doll@googlemail.com>
Co-authored-by: kkytola <kalle.kytola@aalto.fi>
Co-authored-by: Joël Riou <joel.riou@universite-paris-saclay.fr>
Co-authored-by: damiano <adomani@gmail.com>
Co-authored-by: Kyle Miller <kmill31415@gmail.com>
Co-authored-by: Rémy Degenne <remydegenne@gmail.com>
Co-authored-by: Jz Pan <acme_pjz@hotmail.com>
Co-authored-by: Eric Rodriguez <ericrboidi@gmail.com>

Author: 29 people

Mathlib/Tactic/Rewrites.lean

@@ -118,15 +118,21 @@ structure RewriteResult where
   /-- Can the new goal in `result` be closed by `with_reducible rfl`? -/
   -- This is an `Option` so that it can be computed lazily.
   rfl? : Option Bool
+  /-- The metavariable context after the rewrite.
+  This needs to be stored as part of the result so we can backtrack the state. -/
+  mctx : MetavarContext
 
 /-- Update a `RewriteResult` by filling in the `rfl?` field if it is currently `none`,
 to reflect whether the remaining goal can be closed by `with_reducible rfl`. -/
 def RewriteResult.computeRfl (r : RewriteResult) : MetaM RewriteResult := do
   if let some _ := r.rfl? then
     return r
   try
-    (← mkFreshExprMVar r.result.eNew).mvarId!.rfl
-    pure { r with rfl? := some true }
+    withoutModifyingState <| withMCtx r.mctx do
+      -- We use `withReducible` here to follow the behaviour of `rw`.
+      withReducible (← mkFreshExprMVar r.result.eNew).mvarId!.rfl
+      -- We do not need to record the updated `MetavarContext` here.
+      pure { r with rfl? := some true }
   catch _ =>
     pure { r with rfl? := some false }
 
@@ -136,16 +142,18 @@ Find lemmas which can rewrite the goal.
 This core function returns a monadic list, to allow the caller to decide how long to search.
 See also `rewrites` for a more convenient interface.
 -/
+-- We need to supply the current `MetavarContext` (which will be reused for each lemma application)
+-- because `MLList.squash` executes lazily,
+-- so there is no opportunity for `← getMCtx` to record the context at the call site.
 def rewritesCore (lemmas : DiscrTree (Name × Bool × Nat) s × DiscrTree (Name × Bool × Nat) s)
-    (goal : MVarId) (target : Expr) : MLList MetaM RewriteResult := MLList.squash fun _ => do
-
+    (ctx : MetavarContext) (goal : MVarId) (target : Expr) :
+    MLList MetaM RewriteResult := MLList.squash fun _ => do
   -- Get all lemmas which could match some subexpression
   let candidates := (← lemmas.1.getSubexpressionMatches target)
     ++ (← lemmas.2.getSubexpressionMatches target)
 
   -- Sort them by our preferring weighting
   -- (length of discriminant key, doubled for the forward implication)
-
   let candidates := candidates.insertionSort fun r s => r.2.2 > s.2.2
 
   -- Now deduplicate. We can't use `Array.deduplicateSorted` as we haven't completely sorted,
@@ -167,34 +175,34 @@ def rewritesCore (lemmas : DiscrTree (Name × Bool × Nat) s × DiscrTree (Name
 
   -- Lift to a monadic list, so the caller can decide how much of the computation to run.
   let lazy := MLList.ofList deduped.toList
-  pure <| lazy.filterMapM fun ⟨lem, symm, weight⟩ => do
+  pure <| lazy.filterMapM fun ⟨lem, symm, weight⟩ => withMCtx ctx do
     trace[Tactic.rewrites] "considering {if symm then "←" else ""}{lem}"
     let some result ← try? do goal.rewrite target (← mkConstWithFreshMVarLevels lem) symm
       | return none
     return if result.mvarIds.isEmpty then
-      some ⟨lem, symm, weight, result, none⟩
+      some ⟨lem, symm, weight, result, none, ← getMCtx⟩
     else
       -- TODO Perhaps allow new goals? Try closing them with solveByElim?
       none
 
 /-- Find lemmas which can rewrite the goal. -/
 def rewrites (lemmas : DiscrTree (Name × Bool × Nat) s × DiscrTree (Name × Bool × Nat) s)
-    (goal : MVarId) (target : Expr) (stop_at_rfl : Bool := False) (max : Nat := 20)
+    (goal : MVarId) (target : Expr) (stopAtRfl : Bool := false) (max : Nat := 20)
     (leavePercentHeartbeats : Nat := 10) : MetaM (List RewriteResult) := do
-  let results ← rewritesCore lemmas goal target
+  let results ← rewritesCore lemmas (← getMCtx) goal target
     -- Don't report duplicate results.
     -- (TODO: we later pretty print results; save them here?)
     -- (TODO: a config flag to disable this,
     -- if distinct-but-pretty-print-the-same results are desirable?)
     |>.dedupBy (fun r => do pure <| (← ppExpr r.result.eNew).pretty)
     -- Stop if we find a rewrite after which `with_reducible rfl` would succeed.
-    |>.mapM RewriteResult.computeRfl -- TODO could simply not compute this if `stop_at_rfl` is False
-    |>.takeUpToFirst (fun r => stop_at_rfl && r.rfl? = some true)
+    |>.mapM RewriteResult.computeRfl -- TODO could simply not compute this if `stopAtRfl` is False
+    |>.takeUpToFirst (fun r => stopAtRfl && r.rfl? = some true)
     -- Don't use too many heartbeats.
     |>.whileAtLeastHeartbeatsPercent leavePercentHeartbeats
     -- Bound the number of results.
     |>.takeAsList max
-  return match results.filter (fun r => stop_at_rfl && r.rfl? = some true) with
+  return match results.filter (fun r => stopAtRfl && r.rfl? = some true) with
   | [] =>
     -- TODO consider sorting the results,
     -- e.g. if we use solveByElim to fill arguments,
@@ -226,39 +234,41 @@ elab_rules : tactic |
       let some a ← f.findDecl? | return
       if a.isImplementationDetail then return
       let target ← instantiateMVars (← f.getType)
-      let results ← rewrites lems goal target (stop_at_rfl := false)
+      let results ← rewrites lems goal target (stopAtRfl := false)
       reportOutOfHeartbeats `rewrites tk
       if results.isEmpty then
         throwError "Could not find any lemmas which can rewrite the hypothesis {
           ← f.getUserName}"
-      for r in results do
+      for r in results do withMCtx r.mctx do
         addRewriteSuggestion tk [(← mkConstWithFreshMVarLevels r.name, r.symm)]
           r.result.eNew (loc? := .some (.fvar f)) (origSpan? := ← getRef)
       if lucky.isSome then
         match results[0]? with
         | some r => do
+            setMCtx r.mctx
             let replaceResult ← goal.replaceLocalDecl f r.result.eNew r.result.eqProof
             replaceMainGoal (replaceResult.mvarId :: r.result.mvarIds)
         | _ => failure
     -- See https://github.com/leanprover/lean4/issues/2150
     do withMainContext do
       let target ← instantiateMVars (← goal.getType)
-      let results ← rewrites lems goal target (stop_at_rfl := true)
+      let results ← rewrites lems goal target (stopAtRfl := true)
       reportOutOfHeartbeats `rewrites tk
       if results.isEmpty then
         throwError "Could not find any lemmas which can rewrite the goal"
-      for r in results do
+      for r in results do withMCtx r.mctx do
         let newGoal := if r.rfl? = some true then Expr.lit (.strVal "no goals") else r.result.eNew
         addRewriteSuggestion tk [(← mkConstWithFreshMVarLevels r.name, r.symm)]
           newGoal (origSpan? := ← getRef)
       if lucky.isSome then
         match results[0]? with
         | some r => do
+            setMCtx r.mctx
             replaceMainGoal
               ((← goal.replaceTargetEq r.result.eNew r.result.eqProof) :: r.result.mvarIds)
             evalTactic (← `(tactic| try rfl))
         | _ => failure
-    (λ _ => throwError "Failed to find a rewrite for some location")
+    (fun _ => throwError "Failed to find a rewrite for some location")
 
 @[inherit_doc rewrites'] macro "rw?!" h:(ppSpace location)? : tactic =>
   `(tactic| rw? ! $[$h]?)

test/rewrites.lean

@@ -15,6 +15,8 @@ set_option autoImplicit true
 -- `build/lib/MathlibExtras/Rewrites.extra`
 -- so that the cache is rebuilt.
 
+set_option autoImplicit true
+
 /--
 info: Try this: rw [@List.map_append]
 -- "no goals"
@@ -148,3 +150,12 @@ info: Try this: rw [f_eq]
 lemma test : f n = f m := by
   rw?
   rw [f_eq, f_eq]
+
+
+def zero : Nat := 0
+
+-- This used to (incorrectly!) succeed because `rw?` would try `rfl`,
+-- rather than `withReducible` `rfl`.
+example : zero = 0 := by
+  rw?!
+  sorry