feat: upstream rw? tactic (#3719)

This updates the rw? tactic from Mathlib to use lazy discriminator trees
and upstreams it.

---------

Co-authored-by: Scott Morrison <[email protected]>

src/Init/Tactics.lean

@@ -1318,6 +1318,22 @@ used when closing the goal.
 -/
 syntax (name := apply?) "apply?" (" using " (colGt term),+)? : tactic
 
+/--
+Syntax for excluding some names, e.g. `[-my_lemma, -my_theorem]`.
+-/
+syntax rewrites_forbidden := " [" (("-" ident),*,?) "]"
+
+/--
+`rw?` tries to find a lemma which can rewrite the goal.
+
+`rw?` should not be left in proofs; it is a search tool, like `apply?`.
+
+Suggestions are printed as `rw [h]` or `rw [← h]`.
+
+You can use `rw? [-my_lemma, -my_theorem]` to prevent `rw?` using the named lemmas.
+-/
+syntax (name := rewrites?) "rw?" (ppSpace location)? (rewrites_forbidden)? : tactic
+
 /--
 `show_term tac` runs `tac`, then prints the generated term in the form
 "exact X Y Z" or "refine X ?_ Z" if there are remaining subgoals.

src/Lean/Elab/Tactic.lean

@@ -39,3 +39,4 @@ import Lean.Elab.Tactic.SolveByElim
 import Lean.Elab.Tactic.LibrarySearch
 import Lean.Elab.Tactic.ShowTerm
 import Lean.Elab.Tactic.Rfl
+import Lean.Elab.Tactic.Rewrites

src/Lean/Elab/Tactic/Rewrites.lean

@@ -0,0 +1,69 @@
+/-
+Copyright (c) 2023 Scott Morrison. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Scott Morrison
+-/
+prelude
+import Lean.Elab.Tactic.Location
+import Lean.Meta.Tactic.Replace
+import Lean.Meta.Tactic.Rewrites
+
+/-!
+# The `rewrites` tactic.
+
+`rw?` tries to find a lemma which can rewrite the goal.
+
+`rw?` should not be left in proofs; it is a search tool, like `apply?`.
+
+Suggestions are printed as `rw [h]` or `rw [← h]`.
+
+-/
+namespace Lean.Elab.Rewrites
+
+open Lean Meta Rewrites
+open Lean.Parser.Tactic
+
+open Lean Elab Tactic
+
+@[builtin_tactic Lean.Parser.Tactic.rewrites?]
+def evalExact : Tactic := fun stx => do
+  let `(tactic| rw?%$tk $[$loc]? $[[ $[-$forbidden],* ]]?) := stx
+        | throwUnsupportedSyntax
+  let moduleRef ← createModuleTreeRef
+  let forbidden : NameSet :=
+    ((forbidden.getD #[]).map Syntax.getId).foldl (init := ∅) fun s n => s.insert n
+  reportOutOfHeartbeats `findRewrites tk
+  let goal ← getMainGoal
+  withLocation (expandOptLocation (Lean.mkOptionalNode loc))
+    fun f => do
+      let some a ← f.findDecl? | return
+      if a.isImplementationDetail then return
+      let target ← instantiateMVars (← f.getType)
+      let hyps ← localHypotheses (except := [f])
+      let results ← findRewrites hyps moduleRef goal target (stopAtRfl := false) forbidden
+      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 withMCtx r.mctx do
+        Tactic.TryThis.addRewriteSuggestion tk [(r.expr, r.symm)]
+          r.result.eNew (loc? := .some (.fvar f)) (origSpan? := ← getRef)
+      if let some r := results[0]? then
+        setMCtx r.mctx
+        let replaceResult ← goal.replaceLocalDecl f r.result.eNew r.result.eqProof
+        replaceMainGoal (replaceResult.mvarId :: r.result.mvarIds)
+    do
+      let target ← instantiateMVars (← goal.getType)
+      let hyps ← localHypotheses
+      let results ← findRewrites hyps moduleRef goal target (stopAtRfl := true) forbidden
+      reportOutOfHeartbeats `rewrites tk
+      if results.isEmpty then
+        throwError "Could not find any lemmas which can rewrite the goal"
+      results.forM (·.addSuggestion tk)
+      if let some r := results[0]? then
+        setMCtx r.mctx
+        replaceMainGoal
+          ((← goal.replaceTargetEq r.result.eNew r.result.eqProof) :: r.result.mvarIds)
+        evalTactic (← `(tactic| try rfl))
+    (fun _ => throwError "Failed to find a rewrite for some location")
+
+end Lean.Elab.Rewrites

src/Lean/Expr.lean

@@ -1881,6 +1881,22 @@ def letFunAppArgs? (e : Expr) : Option (Array Expr × Name × Expr × Expr × Ex
   | .lam n _ b _ => some (rest, n, t, v, b)
   | _ => some (rest, .anonymous, t, v, .app f (.bvar 0))
 
+/-- Maps `f` on each immediate child of the given expression. -/
+@[specialize]
+def traverseChildren [Applicative M] (f : Expr → M Expr) : Expr → M Expr
+  | e@(forallE _ d b _) => pure e.updateForallE! <*> f d <*> f b
+  | e@(lam _ d b _)     => pure e.updateLambdaE! <*> f d <*> f b
+  | e@(mdata _ b)       => e.updateMData! <$> f b
+  | e@(letE _ t v b _)  => pure e.updateLet! <*> f t <*> f v <*> f b
+  | e@(app l r)         => pure e.updateApp! <*> f l <*> f r
+  | e@(proj _ _ b)      => e.updateProj! <$> f b
+  | e                   => pure e
+
+/-- `e.foldlM f a` folds the monadic function `f` over the subterms of the expression `e`,
+with initial value `a`. -/
+def foldlM {α : Type} {m} [Monad m] (f : α → Expr → m α) (init : α) (e : Expr) : m α :=
+  Prod.snd <$> StateT.run (e.traverseChildren (fun e' => fun a => Prod.mk e' <$> f a e')) init
+
 end Expr
 
 /--