Skip to content

Commit f2c85c3

Browse files
kim-emParcly-TaxelScott Morrison
committed
feat: rewrites, a tactic to suggest rewrites. (#3119)
Unlike the very slow `rw_hint` long ago proposed (but never merged) for mathlib3, this uses discrimination trees (with keys given by LHS and RHS of lemmas), and looks up all subexpressions of the target expression. - [x] depends on: #2898 - [x] depends on: #3181 [![Open in Gitpod](https://gitpod.io/button/open-in-gitpod.svg)](https://gitpod.io/from-referrer/) Co-authored-by: Scott Morrison <scott.morrison@gmail.com> Co-authored-by: Jeremy Tan Jie Rui <reddeloostw@gmail.com> Co-authored-by: Scott Morrison <scott.morrison@anu.edu.au>
1 parent ebf0b75 commit f2c85c3

File tree

8 files changed

+226
-9
lines changed

8 files changed

+226
-9
lines changed

Mathlib.lean

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1889,6 +1889,7 @@ import Mathlib.Tactic.Rename
18891889
import Mathlib.Tactic.RenameBVar
18901890
import Mathlib.Tactic.Replace
18911891
import Mathlib.Tactic.RestateAxiom
1892+
import Mathlib.Tactic.Rewrites
18921893
import Mathlib.Tactic.Ring
18931894
import Mathlib.Tactic.Ring.Basic
18941895
import Mathlib.Tactic.Ring.RingNF

Mathlib/Lean/CoreM.lean

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,3 +22,9 @@ Return the percentage of the max heartbeats allowed
2222
that have been consumed so far in this computation.
2323
-/
2424
def heartbeatsPercent : CoreM Nat := do pure <| (← IO.getNumHeartbeats) * 100 / (← getMaxHeartbeats)
25+
26+
/-- Log a message if it looks like we ran out of time. -/
27+
def reportOutOfHeartbeats (tac : Name) (stx : Syntax) : CoreM Unit := do
28+
if (← heartbeatsPercent) ≥ 90 then
29+
logInfoAt stx (s!"`{tac}` stopped because it was running out of time.\n" ++
30+
"You may get better results using `set_option maxHeartbeats 0`.")

Mathlib/Lean/Meta/DiscrTree.lean

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ Released under Apache 2.0 license as described in the file LICENSE.
44
Authors: Scott Morrison
55
-/
66
import Lean.Meta.DiscrTree
7+
import Mathlib.Lean.Expr.Traverse
78

89
/-!
910
# Additions to `Lean.Meta.DiscrTree`
@@ -22,6 +23,11 @@ def insertIfSpecific [BEq α] (d : DiscrTree α s)
2223
else
2324
d.insertCore keys v
2425

26+
/--
27+
Find keys which match the expression, or some subexpression.
28+
-/
29+
partial def getSubexpressionMatches (d : DiscrTree α s) (e : Expr) : MetaM (Array α) := do
30+
e.foldlM (fun a f => do pure <| a ++ (← d.getSubexpressionMatches f)) (← d.getMatch e)
2531
variable {m : TypeType} [Monad m]
2632

2733
/-- Apply a monadic function to the array of values at each node in a `DiscrTree`. -/

Mathlib/Tactic.lean

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -95,6 +95,7 @@ import Mathlib.Tactic.Rename
9595
import Mathlib.Tactic.RenameBVar
9696
import Mathlib.Tactic.Replace
9797
import Mathlib.Tactic.RestateAxiom
98+
import Mathlib.Tactic.Rewrites
9899
import Mathlib.Tactic.Ring
99100
import Mathlib.Tactic.Ring.Basic
100101
import Mathlib.Tactic.Ring.RingNF

Mathlib/Tactic/LibrarySearch.lean

Lines changed: 3 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -242,12 +242,6 @@ def librarySearch (goal : MVarId) (required : List Expr)
242242
setMCtx ctx
243243
return none)
244244

245-
/-- Log a message if it looks like we ran out of time. -/
246-
def reportOutOfHeartbeats (stx : Syntax) : MetaM Unit := do
247-
if (← heartbeatsPercent) ≥ 90 then
248-
logInfoAt stx ("`library_search` stopped because it was running out of time.\n" ++
249-
"You may get better results using `set_option maxHeartbeats 0`.")
250-
251245
open Lean.Parser.Tactic
252246

253247
-- TODO: implement the additional options for `library_search` from Lean 3,
@@ -268,7 +262,7 @@ elab_rules : tactic | `(tactic| library_search%$tk $[using $[$required:term],*]?
268262
goal.withContext do
269263
let required := (← (required.getD #[]).mapM getFVarId).toList.map .fvar
270264
if let some suggestions ← librarySearch goal required then
271-
reportOutOfHeartbeats tk
265+
reportOutOfHeartbeats `library_search tk
272266
for suggestion in suggestions do
273267
withMCtx suggestion.1 do
274268
addRefineSuggestion tk (← instantiateMVars (mkMVar mvar)).headBeta
@@ -283,7 +277,7 @@ elab tk:"library_search%" : term <= expectedType => do
283277
let (_, introdGoal) ← goal.mvarId!.intros
284278
introdGoal.withContext do
285279
if let some suggestions ← librarySearch introdGoal [] then
286-
reportOutOfHeartbeats tk
280+
reportOutOfHeartbeats `library_search tk
287281
for suggestion in suggestions do
288282
withMCtx suggestion.1 do
289283
addTermSuggestion tk (← instantiateMVars goal).headBeta
@@ -313,7 +307,7 @@ elab_rules : tactic |
313307
let (type, _) ← elabTermWithHoles t none (← getMainTag) true
314308
let .mvar goal ← mkFreshExprMVar type | failure
315309
if let some _ ← librarySearch goal [] then
316-
reportOutOfHeartbeats tk
310+
reportOutOfHeartbeats `library_search tk
317311
throwError "observe did not find a solution"
318312
else
319313
let v := (← instantiateMVars (mkMVar goal)).headBeta

Mathlib/Tactic/Rewrites.lean

Lines changed: 153 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,153 @@
1+
/-
2+
Copyright (c) 2023 Scott Morrison. All rights reserved.
3+
Released under Apache 2.0 license as described in the file LICENSE.
4+
Authors: Scott Morrison
5+
-/
6+
import Mathlib.Data.ListM.Heartbeats
7+
import Mathlib.Lean.Meta.DiscrTree
8+
import Mathlib.Tactic.Cache
9+
import Mathlib.Tactic.SolveByElim
10+
import Mathlib.Tactic.TryThis
11+
import Mathlib.Control.Basic
12+
13+
/-!
14+
# The `rewrites` tactic.
15+
16+
`rewrites` tries to find a lemma which can rewrite the goal.
17+
18+
`rewrites` should not be left in proofs; it is a search tool, like `library_search`.
19+
20+
Suggestions are printed as `rw [h]` or `rw [←h]`.
21+
22+
## Future work
23+
It would be nice to have `rewrites at h`.
24+
25+
We could also try discharging side goals via `assumption` or `solve_by_elim`.
26+
27+
-/
28+
29+
namespace Mathlib.Tactic.Rewrites
30+
31+
open Lean Meta Std.Tactic.TryThis
32+
33+
initialize registerTraceClass `Tactic.rewrites
34+
35+
initialize rewriteLemmas : DeclCache (DiscrTree (Name × Bool × Nat) true) ←
36+
DeclCache.mk "rewrites: init cache" {} fun name constInfo lemmas => do
37+
if constInfo.isUnsafe then return lemmas
38+
if ← name.isBlackListed then return lemmas
39+
if name matches .str _ "injEq" then return lemmas
40+
if name matches .str _ "sizeOf_spec" then return lemmas
41+
match name with
42+
| .str _ n => if n.endsWith "_inj" ∨ n.endsWith "_inj'" then return lemmas
43+
| _ => pure ()
44+
let forwardWeight := 2
45+
let backwardWeight := 1
46+
withNewMCtxDepth do withReducible do
47+
let (_, _, type) ← forallMetaTelescopeReducing constInfo.type
48+
match type.getAppFnArgs with
49+
| (``Eq, #[_, lhs, rhs]) => do
50+
let lhsKey ← DiscrTree.mkPath lhs
51+
let rhsKey ← DiscrTree.mkPath rhs
52+
pure <| lemmas.insertIfSpecific lhsKey (name, false, forwardWeight * lhsKey.size)
53+
|>.insertIfSpecific rhsKey (name, true, backwardWeight * rhsKey.size)
54+
| (``Iff, #[lhs, rhs]) => do
55+
let lhsKey ← DiscrTree.mkPath lhs
56+
let rhsKey ← DiscrTree.mkPath rhs
57+
pure <| lemmas.insertIfSpecific lhsKey (name, false, forwardWeight * lhsKey.size)
58+
|>.insertIfSpecific rhsKey (name, true, backwardWeight * rhsKey.size)
59+
| _ => pure lemmas
60+
61+
/-- Data structure recording a potential rewrite to report from the `rewrites` tactic. -/
62+
structure RewriteResult where
63+
name : Name
64+
symm : Bool
65+
weight : Nat
66+
result : Meta.RewriteResult
67+
/-- Can the new goal in `result` be closed by `with_reducible rfl`? -/
68+
-- This is an `Option` so that it can be computed lazily.
69+
refl? : Option Bool
70+
71+
/-- Update a `RewriteResult` by filling in the `refl?` field if it is currently `none`,
72+
to reflect whether the remaining goal can be closed by `with_reducible rfl`. -/
73+
def RewriteResult.computeRefl (r : RewriteResult) : MetaM RewriteResult := do
74+
if let some _ := r.refl? then
75+
return r
76+
try
77+
(← mkFreshExprMVar r.result.eNew).mvarId!.refl
78+
pure { r with refl? := some true }
79+
catch _ =>
80+
pure { r with refl? := some false }
81+
82+
/--
83+
Find lemmas which can rewrite the goal.
84+
85+
This core function returns a monadic list, to allow the caller to decide how long to search.
86+
See also `rewrites` for a more convenient interface.
87+
-/
88+
def rewritesCore (lemmas : DiscrTree (Name × Bool × Nat) s) (goal : MVarId) :
89+
ListM MetaM RewriteResult := ListM.squash do
90+
let type ← instantiateMVars (← goal.getType)
91+
-- Get all lemmas which could match some subexpression
92+
let candidates ← lemmas.getSubexpressionMatches type
93+
-- Sort them by our preferring weighting
94+
-- (length of discriminant key, doubled for the forward implication)
95+
let candidates := candidates.insertionSort fun r s => r.2.2 > s.2.2
96+
-- Lift to a monadic list, so the caller can decide how much of the computation to run.
97+
let candidates := ListM.ofList candidates.toList
98+
pure <| candidates.filterMapM fun ⟨lem, symm, weight⟩ => do
99+
trace[Tactic.rewrites] "considering {if symm then "" else ""}{lem}"
100+
let some result ← try? <| goal.rewrite type (← mkConstWithFreshMVarLevels lem) symm
101+
| return none
102+
return if result.mvarIds.isEmpty then
103+
some ⟨lem, symm, weight, result, none⟩
104+
else
105+
-- TODO Perhaps allow new goals? Try closing them with solveByElim?
106+
none
107+
108+
/-- Find lemmas which can rewrite the goal. -/
109+
def rewrites (lemmas : DiscrTree (Name × Bool × Nat) s) (goal : MVarId)
110+
(max : Nat := 10) (leavePercentHeartbeats : Nat := 10) : MetaM (List RewriteResult) :=
111+
rewritesCore lemmas goal
112+
-- Don't use too many heartbeats.
113+
|>.whileAtLeastHeartbeatsPercent leavePercentHeartbeats
114+
-- Stop if we find a rewrite after which `with_reducible rfl` would succeed.
115+
|>.mapM RewriteResult.computeRefl
116+
|>.takeUpToFirst (fun r => r.refl? = some true)
117+
-- Bound the number of results.
118+
|>.takeAsList max
119+
120+
open Lean.Parser.Tactic
121+
122+
/--
123+
`rewrites` tries to find a lemma which can rewrite the goal.
124+
125+
`rewrites` should not be left in proofs; it is a search tool, like `library_search`.
126+
127+
Suggestions are printed as `rw [h]` or `rw [←h]`.
128+
`rewrites!` is the "I'm feeling lucky" mode, and will run the first rewrite it finds.
129+
-/
130+
syntax (name := rewrites') "rewrites" "!"? : tactic
131+
132+
open Elab.Tactic Elab Tactic in
133+
elab_rules : tactic |
134+
`(tactic| rewrites%$tk $[!%$lucky]?) => do
135+
let goal ← getMainGoal
136+
goal.withContext do
137+
let results ← rewrites (← rewriteLemmas.get) goal
138+
reportOutOfHeartbeats `rewrites tk
139+
if results.isEmpty then
140+
throwError "rewrites could not find any lemmas which can rewrite the goal"
141+
for r in results do
142+
let newGoal := if r.refl? = some true then Expr.lit (.strVal "no goals") else r.result.eNew
143+
addRewriteSuggestion tk (← mkConstWithFreshMVarLevels r.name) r.symm newGoal
144+
if lucky.isSome then
145+
match results[0]? with
146+
| some r => do
147+
replaceMainGoal
148+
((← goal.replaceTargetEq r.result.eNew r.result.eqProof) :: r.result.mvarIds)
149+
evalTactic (← `(tactic| rfl))
150+
| _ => failure
151+
152+
@[inherit_doc rewrites'] macro "rewrites!" : tactic =>
153+
`(tactic| rewrites !)

Mathlib/Tactic/TryThis.lean

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,22 @@ def addHaveSuggestion (origTac : Syntax) (t? : Option Expr) (e : Expr) :
3232
`(tactic| let this := $estx)
3333
addSuggestion origTac tac
3434

35+
/-- Add a suggestion for `rw [h]`. (TODO: this depends on code action support) -/
36+
def addRewriteSuggestion (origTac : Syntax) (e : Expr) (symm : Bool) (type? : Option Expr := none) :
37+
TermElabM Unit := do
38+
let estx ← delabToRefinableSyntax e
39+
let tac ←
40+
if symm then
41+
`(tactic| rw [← $estx])
42+
else
43+
`(tactic| rw [$estx:term])
44+
-- We resort to using `logInfoAt` here rather than `addSuggestion`,
45+
-- as I've never worked out how to have `addSuggestion` render comments.
46+
if let some type := type? then
47+
logInfoAt origTac m!"{tac}\n-- {← ppExpr type}"
48+
else
49+
logInfoAt origTac m!"{tac}"
50+
3551
/-- Add a `refine e` suggestion, also printing the type of the subgoals.
3652
(TODO: this depends on code action support) -/
3753
def addRefineSuggestion (origTac : Syntax) (e : Expr) : TermElabM Unit := do

test/rewrites.lean

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
import Mathlib.Tactic.Rewrites
2+
import Mathlib
3+
4+
example (f : α → β) (L M : List α) : (L ++ M).map f = L.map f ++ M.map f := by
5+
rewrites!
6+
7+
open CategoryTheory
8+
9+
example [Category C] {X Y Z : C} (f : X ⟶ Y) (g : Y ⟶ Z) : f ≫ 𝟙 _ ≫ g = f ≫ g := by
10+
rewrites!
11+
12+
example [Group G] (h : G) : 1 * h = h := by
13+
rewrites!
14+
15+
example [Group G] (g h : G) : g * g⁻¹ * h = h := by
16+
rewrites -- the right answer is not the first solution, so we can't use rewrites!
17+
/- Prints:
18+
rw [@Semigroup.mul_assoc]
19+
-- g * (g⁻¹ * h) = h
20+
rw [@mul_assoc]
21+
-- g * (g⁻¹ * h) = h
22+
rw [@mul_left_eq_self]
23+
-- g * g⁻¹ = 1
24+
rw [@mul_inv_self]
25+
-- 1 * h = h
26+
rw [@mul_right_inv]
27+
-- 1 * h = h
28+
rw [← @division_def]
29+
-- g / g * h = h
30+
rw [← @div_eq_mul_inv]
31+
-- g / g * h = h
32+
rw [← @DivInvMonoid.div_eq_mul_inv]
33+
-- g / g * h = h
34+
rw [@inv_eq_one_div]
35+
-- g * (1 / g) * h = h
36+
rw [← @mulRightEmbedding_apply]
37+
-- ↑(mulRightEmbedding h) (g * g⁻¹) = h
38+
-/
39+
rw [mul_inv_self]
40+
rw [one_mul]

0 commit comments

Comments
 (0)