Implementing a tactic modifier

Stream: lean4

[François G. Dorais]

I was thinking about implementing a "tactic modifier" like clean tac that would run the tactic tac with some additional cleanup afterwards to get rid of common crud (e.g, superfluous lambdas, Nat.add instead of +, etc.). I got a partial implementation working using macros but I'm struggling with some desirable features. Here is the most desirable feature that I don't have a clue how to do.

Here's the base case macro I'm using

syntax (name := clean) "clean " tactic (location)? : tactic
macro_rules
| `(tactic| clean $tac:tactic $[$loc:location]?) => `(tactic| $tac; simp only [clean] $[$loc]?)

where clean is a simp attribute for this task. (It doesn't matter what clean does, so replace with simp only [] if you want to experiment on your own machine.)

The first issue that I'm running into is that I need to write clean rw [blah] at h at h instead of just clean rw [blah] at h where clean would somehow detect that the rewrite has a location parameter and only clean at that location. Is there a way to do that?

[Alex Keizer]

What's happening is that the location parameter is "consumed" by clean with a higher precedence, so clean rw[x] at h is parsed as clean (rw[x]) at h, i.e., $tac = rw[x]. You want to somehow pass that location parameter onto the tactic you're executing.

[Henrik Böving]

I think they are aware of what is happening, the question is how to do this. I guess one could have parenthesis around the tactic and if a location modifier is passed to the clean one also pass it to the parenthesized tactic.

[Alex Keizer]

One way is by having an explicit case for each tactic with a location you want to support, with the downside of it being a bit brittle, you have to duplicate the syntax of every such location-bearing tactic
Something like this?

open Lean.Parser.Tactic (location)
syntax (name := clean) "clean " tactic (location)? : tactic
macro_rules
| `(tactic| clean rw $[$cfg:config]? $rw:rwRuleSeq $[$loc:location]?) => `(tactic| rw $[$cfg]? $rw $[$loc]?; simp only [clean] $[$loc]?)
| `(tactic| clean $tac:tactic $[$loc:location]?) => `(tactic| $tac; simp only [clean] $[$loc]?)

[François G. Dorais]

It's actually the opposite: clean rw [blah] at h parses as clean (rw [blah] at h) Anyway, I got something working! (It's funny how often merely asking for help leads to a sudden breakthrough!)

syntax (name := clean) "clean " tactic (location)? : tactic
macro_rules
| `(tactic| clean $tac:tactic) => do
  match tac.find? (λ s => s.getKind == ``location) with
  | some loc => `(tactic| $tac; simp only [] $loc)
  | none => `(tactic| $tac; simp only [])
| `(tactic| clean $tac:tactic $loc) => `(tactic| $tac; simp only [] $loc)

Right now this doesn't do exactly the right thing for clean (rw [foo] at h1; rw [bar] at h2) because it only finds the location at h1.

Is there an analogue of Syntax.find? that returns a list of all occurrences? (I'll roll my own if necessary but I generally try to avoid yet another implementation of the wheel.)

The next issue is how to correctly unify locations. Maybe there's a neat trick for that?

[François G. Dorais]

I've solved it! Improvements are welcome!

-- could be a useful addition to Prelude.lean
partial def Lean.Syntax.filter (stx : Syntax) (p : Syntax → Bool) : Array Syntax :=
  match stx with
  | stx@(Syntax.node _ _ args) =>
    let as := args.concatMap (λ stx => filter stx p)
    if p stx then as.push stx else as
  | stx => if p stx then #[stx] else #[]

-- could be a useful addition to Lean/Elab/Tactic/Location.lean
def Lean.Elab.Tactic.joinLocation : Location → Location → Location
| .wildcard, _ => .wildcard
| _, .wildcard => .wildcard
| .targets hs₁ g₁, .targets hs₂ g₂ => .targets (hs₁ ++ hs₂) (g₁ || g₂)

open Lean.Parser.Tactic (location)
open Lean.Elab.Tactic (Location expandLocation joinLocation)

syntax (name := clean) "clean " tactic (location)? : tactic

macro_rules
| `(tactic| clean $tac $[$loc]?) => do
  let mut loc : Location := match loc with
  | some loc => expandLocation loc
  | none => .targets #[] false
  for stx in tac.filter fun stx => stx.getKind == ``location do
    loc := joinLocation loc (expandLocation stx)
  match loc with
  | .wildcard => `(tactic| $tac; simp only [clean] at *)
  | .targets hs _ => `(tactic| $tac; simp only [clean] at $[$hs]* ⊢)