Custom pre/post methods in simp

Stream: lean4

[Tomas Skrivan]

I would like to have the following simp theorem but it should be applied only when the lhs has the let binding:

theorem D_let (g : α → β) (f : α → β → γ)
  : D (λ x =>
         let tmp := g x
         f x tmp)
    = λ x dx =>
         let tmp  := g x
         let dtmp := D g x dx
         D (f x) tmp dtmp + D f x dx tmp

The simplifier ignores the let binding. The lhs could be also written as D (λ x => f x (g x)), it makes no difference. (Recently, I had a question related to this, but I was looking at it from the incorrect angle)

Internally, simplifier has pre and post methods of type Expr -> SimpM Step that are applied to each subterm before or after simp lemmas are applied. I can achieve what I need with a custom pre method. However, currently there is no mechanism how to define it.

I have modifed Lean source code to do that, my lean4 fork. Now, simpLocation has two additional arguments (pres posts : Array (Expr → Simp.SimpM Simp.Step) := #[]) and you can define your own simp tactic with custom pre and post methods:

import Lean.Meta
import Lean.Elab

open Lean Meta Simp

def myPre (e : Expr) : SimpM Step :=
  dbg_trace s!"Running pre  on {e}"
  pure (Step.visit (Result.mk e none))

def myPost (e : Expr) : SimpM Step :=
  dbg_trace s!"Running post on {e}"
  pure (Step.visit (Result.mk e none))

namespace Lean.Elab.Tactic

open Lean.Parser.Tactic

syntax (name := mysimp) "mysimp " (config)? (discharger)? (&"only ")? ("[" (simpStar <|> simpErase <|> simpLemma),* "]")? (location)? : tactic

@[tactic mysimp] def evalMySimp : Tactic := fun stx => do
  let { ctx, fvarIdToLemmaId, dischargeWrapper, pres, posts } ← withMainContext <| mkSimpContext stx (eraseLocal := false)
  dbg_trace "Running my simp!"
  dischargeWrapper.with fun discharge? =>
    simpLocation ctx discharge? fvarIdToLemmaId (expandOptLocation stx[5]) #[myPre] #[myPost]

end Lean.Elab.Tactic

@[irreducible] def foo (a b : Nat) := a + b
@[simp] theorem foo_simp (a b : Nat) : foo a b = a + b := by unfold foo; rfl

example (a b : Nat)
  : foo a b = a + b := by mysimp

Outputs:

Running my simp!
Running pre  on Eq.{1} Nat (foo _uniq.823 _uniq.824) (HAdd.hAdd.{0 0 0} Nat Nat Nat (instHAdd.{0} Nat instAddNat) _uniq.823 _uniq.824)
Running pre  on Eq.{1}
Running post on Eq.{1}
Running pre  on foo _uniq.823 _uniq.824
...

I have also attempted to modify default simp such that you can write simp [↓myPre, ↑myPost]. When parsing the tactic, I can get the name myPre : Name and check that a constant with that name has the desired type Expr -> SimpM Step but I'm having hard time converting to the actual value Expr -> SimpM Step. I have tried the command evalConst but that one is unsafe.

---

Would someone else find this useful? What would be the ideal way to specify custom pre/post methods? Would it make sense to add this change to the core repo?

[Leonardo de Moura]

This is a useful feature, but we don't have cycles to invest in it right now. We are currently trying to focus on bugs and features that we need for porting Mathlib to Lean 4. We do want to provide better support for custom simp extensions in the future. Right now, it is not very convenient as you noticed, and one has to implement their own simp command to be able to provide their own pre and post.

structure Methods where
  pre        : Expr → SimpM Step          := fun e => return Step.visit { expr := e }
  post       : Expr → SimpM Step          := fun e => return Step.done { expr := e }
  discharge? : Expr → SimpM (Option Expr) := fun e => return none

[Tomas Skrivan]

That is exactly what I did, I modified the default Methods inside Lean.Meta.Simp.simp:

  | none   => Simp.main e ctx (methods := Simp.DefaultMethods.methods
                                          |> pres.foldl  (λ m p => m.appendPre p)
                                          |> posts.foldl (λ m p => m.appendPost p))

Yeah I understand that this is not a priority right now, I will find some work around it.

[Leonardo de Moura]

Thanks for your patiencience. We will address this issue in the future.