fix: bug fixes to rw? (#5563)

As reported on [zulip](https://leanprover.zulipchat.com/#narrow/stream/287929-mathlib4/topic/panic.20and.20error.20with.20rw.3F/near/370495531)



Co-authored-by: Scott Morrison <scott.morrison@gmail.com>

Mathlib/Lean/Meta/DiscrTree.lean

@@ -26,12 +26,28 @@ def insertIfSpecific [BEq α] (d : DiscrTree α s)
 /--
 Find keys which match the expression, or some subexpression.
 
+Note that repeated subexpressions will be visited each time they appear,
+making this operation potentially very expensive.
+It would be good to solve this problem!
+
 Implementation: we reverse the results from `getMatch`,
 so that we return lemmas matching larger subexpressions first,
 and amongst those we return more specific lemmas first.
 -/
 partial def getSubexpressionMatches (d : DiscrTree α s) (e : Expr) : MetaM (Array α) := do
-  e.foldlM (fun a f => do pure <| a ++ (← d.getSubexpressionMatches f)) (← d.getMatch e).reverse
+  match e with
+  | .bvar _ => return #[]
+  | .forallE _ _ _ _ => forallTelescope e (fun args body => do
+      args.foldlM (fun acc arg => do
+          pure <| acc ++ (← d.getSubexpressionMatches (← inferType arg)))
+        (← d.getSubexpressionMatches body).reverse)
+  | .lam _ _ _ _ => lambdaTelescope e (fun args body => do
+      args.foldlM (fun acc arg => do
+          pure <| acc ++ (← d.getSubexpressionMatches (← inferType arg)))
+        (← d.getSubexpressionMatches body).reverse)
+  | _ =>
+    e.foldlM (fun a f => do pure <| a ++ (← d.getSubexpressionMatches f)) (← d.getMatch e).reverse
+
 variable {m : Type → Type} [Monad m]
 
 /-- Apply a monadic function to the array of values at each node in a `DiscrTree`. -/

Mathlib/Tactic/Rewrites.lean

@@ -217,7 +217,8 @@ elab_rules : tactic |
             let replaceResult ← goal.replaceLocalDecl f r.result.eNew r.result.eqProof
             replaceMainGoal (replaceResult.mvarId :: r.result.mvarIds)
         | _ => failure
-    do
+    -- 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)
       reportOutOfHeartbeats `rewrites tk

test/rewrites.lean

@@ -13,14 +13,29 @@ import Mathlib.Algebra.Group.Basic
 -- `build/lib/MathlibExtras/Rewrites.extra`
 -- so that the cache is rebuilt.
 
+/--
+info: Try this: rw [@List.map_append]
+-- "no goals"
+-/
+#guard_msgs in
 example (f : α → β) (L M : List α) : (L ++ M).map f = L.map f ++ M.map f := by
   rw?!
 
 open CategoryTheory
 
+/--
+info: Try this: rw [@Category.id_comp]
+-- "no goals"
+-/
+#guard_msgs in
 example [Category C] {X Y Z : C} (f : X ⟶ Y) (g : Y ⟶ Z) : f ≫ 𝟙 _ ≫ g = f ≫ g := by
   rw?!
 
+/--
+info: Try this: rw [@mul_left_eq_self]
+-- "no goals"
+-/
+#guard_msgs in
 example [Group G] (h : G) : 1 * h = h := by
   rw?!
 
@@ -57,3 +72,34 @@ lemma prime_of_prime (n : ℕ) : Prime n ↔ Nat.Prime n := by
 example [Group G] (h : G) (hyp : g * 1 = h) : g = h := by
   rw?! at hyp
   assumption
+
+example : ∀ (x y : ℕ), x ≤ y := by
+  intros x y
+  rw? -- Used to be an error here https://leanprover.zulipchat.com/#narrow/stream/287929-mathlib4/topic/panic.20and.20error.20with.20rw.3F/near/370495531
+  admit
+
+example : ∀ (x y : ℕ), x ≤ y := by
+  -- Used to be a panic here https://leanprover.zulipchat.com/#narrow/stream/287929-mathlib4/topic/panic.20and.20error.20with.20rw.3F/near/370495531
+  success_if_fail_with_msg "Could not find any lemmas which can rewrite the goal" rw?
+  admit
+
+axiom K : Type
+@[instance] axiom K.ring : Ring K
+
+def foo : K → K := sorry
+
+example : foo x = 1 ↔ ∃ k : ℤ, x = k := by
+  rw? -- Used to panic, see https://leanprover.zulipchat.com/#narrow/stream/287929-mathlib4/topic/panic.20and.20error.20with.20rw.3F/near/370598036
+  admit
+
+lemma six_eq_seven : 6 = 7 := sorry
+
+example : ∀ (x : ℕ), x ≤ 6 := by
+  rw?!
+  guard_target = ∀ (x : ℕ), x ≤ 7
+  admit
+
+example : ∀ (x : ℕ) (w : x ≤ 6), x ≤ 8 := by
+  rw?!
+  guard_target = ∀ (x : ℕ) (w : x ≤ 7), x ≤ 8
+  admit