tactical rewrite of expressions, outside of proofs

Stream: general

[Matthew Pocock]

Still playing with collaltz, and doing what I've seen maths people do on youtube when they don't understand a problem and creating a number system for it. Turns out that it's "almost" binary written backwards. Anyhoo ... to make a long story short I have this fragment:

-- a collatz numeral, decomposed into odd and even components
inductive ColN where
    | zero      : ColN                -- zero, also leading zero
    | of_even   : ColN → ColN   -- the even number that is twice the previous
    | of_odd    : ColN → ColN   -- the odd number that is one more than twice the previous

-- Evaluation at 0 recovers the natural number equivalent to the faux binary representation.
def ColN.eval (cn: ColN) (n : ℕ) : ℕ := match cn with
| ColN.zero => n
| ColN.of_even cm => 2 * (ColN.eval cm n)
| ColN.of_odd cm => 1 + 2 * (ColN.eval cm n)

#reduce ColN.eval (ColN.of_even $ ColN.of_odd ColN.zero) 0 -- 2
#reduce ColN.eval (ColN.of_even $ ColN.of_odd ColN.zero) 1 -- 6
#reduce ColN.eval (ColN.of_even $ ColN.of_odd ColN.zero) 2 -- 10

example (n : ℕ): (ColN.eval (ColN.of_even $ ColN.of_odd ColN.zero) $ n) = 4 * n + 2 := by
    repeat rw [ColN.eval]
    ring

#reduce ColN.eval (ColN.of_even $ ColN.of_odd ColN.zero) -- some ungodly eldritch terror

Clearly, rewriting ColN.eval and applying ring can simplify down my number to 4n+2, but is there a way to use this directly to automatically simplify the horrific lambda expression to λn, 4*n+2 outside of a proof and without having the target expression already written out? I guess I'm asking how to use tactics to rewrite/simplify code. I won't typically know the target equation and also don't really want to need to work it out on paper.

[Sky Wilshaw]

If you add the following lemmas, simp will prove your repeat rw [ColN.eval] expression.

@[simp]
lemma eval_zero (n : ℕ) : ColN.eval ColN.zero n = n := rfl

@[simp]
lemma eval_of_even (cn : ColN) (n : ℕ) : (ColN.of_even cn).eval n = 2 * cn.eval n := rfl

@[simp]
lemma eval_of_odd (cn : ColN) (n : ℕ) : (ColN.of_odd cn).eval n = 2 * cn.eval n + 1 := rfl

[Sky Wilshaw]

I did slightly change this line:

| ColN.of_odd cm => 2 * (ColN.eval cm n) + 1  -- not 1 + 2 * (ColN.eval cm n)

just because it's normally better style to add one on the right in Lean. That's just because natural addition is defined by recursion on the right.

[Sky Wilshaw]

You can write your own tactic to do that "simp then ring" process:

macro "eval_coln" : tactic => `(tactic| simp only [eval_zero, eval_of_even, eval_of_odd] <;> ring_nf)

then the proof becomes

example (n : ℕ): (ColN.eval (ColN.of_even $ ColN.of_odd ColN.zero) $ n) = 4 * n + 2 :=
  by eval_coln

[Sky Wilshaw]

It does a bit more than just evaluate ColN.eval expressions though - the ring_nf in the macro will try to turn all ring expressions in your goal into a normal form.

[Sky Wilshaw]

It's less readable to use a custom tactic though, so I think it's probably only better if you're using it lots.

[Matthew Pocock]

Thanks. I've tidied up my definitions. But my question isn't really about the example but the #reduce. I really want something that will eat my nasty #reduce ColN.eval ... expression and convince reduce to print out the equivalent of λn, 4*n+2 instead of the mess it currently does.

[Sky Wilshaw]

Ah, I see. In this case, we can change the simp lemmas to this form:

@[simp]
lemma eval_zero : ColN.eval ColN.zero = id := rfl

@[simp]
lemma eval_of_even (cn : ColN) : (ColN.of_even cn).eval = 2 * cn.eval := rfl

@[simp]
lemma eval_of_odd (cn : ColN) : (ColN.of_odd cn).eval = 2 * cn.eval + 1 := rfl

[Sky Wilshaw]

Note that the LHS and RHS of each lemma is now a function. For example, eval_of_even says that the function (ColN.of_even cn).eval is 2 * the function cn.eval. Scalar multiplication on functions is defined "pointwise", so (2 * f) x = 2 * (f x).

[Sky Wilshaw]

Then, calling simp on ColN.eval (ColN.of_even $ ColN.of_odd ColN.zero) gives 2 * (2 * id + 1). I wonder if we can make this any better.

[Sky Wilshaw]

If we write the simp lemmas in this form:

@[simp]
lemma eval_zero : ColN.eval ColN.zero = id := rfl

@[simp]
lemma eval_of_even (cn : ColN) : (ColN.of_even cn).eval = fun n => 2 * cn.eval n := rfl

@[simp]
lemma eval_of_odd (cn : ColN) : (ColN.of_odd cn).eval = fun n => 2 * cn.eval n + 1 := rfl

it introduces lambda abstractions as needed. Then simping ColN.eval (ColN.of_even $ ColN.of_odd ColN.zero) gives fun n => 2 * (2 * n + 1)). The issue with this is that ring is a bit worse when there's a fun in the way.

[Matthew Pocock]

I've done it manually for now:

/--!
This evaluates a `ColN` to a pair, ⟨m, c⟩ that can be used as in `m*x+c` to recover the equation of the colaltz number.
Unlike `ColN.eval`, it does not produce a nested lambda structure, but instead keeps everything flat.
-/
def ColN.eval_p (cn: ColN) : ℕ×ℕ := match cn with
| ColN.zero => ⟨1,0⟩
| ColN.of_even cm =>
    let ⟨ m, c ⟩ := cm.eval_p
    ⟨2*m, 2*c⟩
| ColN.of_odd cm =>
    let ⟨ m, c ⟩ := cm.eval_p
    ⟨2*m, 2*c + 1⟩

/--!
Evaluate a `ColN` using the parameterisation calculated by `ColN.eval_p`.
-/
def ColN.eval_p_n (cn: ColN) (n : ℕ) : ℕ :=
    let ⟨m, c⟩ := cn.eval_p
    m * n + c

Perhaps writing in out manually is how we're supposed to do things. I still wonder if there's a way to encourage tactics to rewrite code, outside of trying to satisfy a proof.

[Sky Wilshaw]

This is probably a good idea. Of course, you can also write simp lemmas to simplify these kinds of expressions, and the simplifier will have an easier job if things are parameters instead of inside lambdas. Incidentally, I think you can get away with splitting eval_p into two functions - one for the left output and one for the right. It might make manipulation a bit easier (but you may need more lemmas).

[Matthew Pocock]

Yeah. I can certainly do that. But what I really want to do is take a value like ColN.eval (ColN.of_even $ ColN.of_odd ColN.zero), and where #reduce currently prints out a complex lambda, be able to apply a tactic to that expression to reduce it to the simple one. So, to invent syntax.

#reduce {simp, ring_nf} ColN.eval (ColN.of_even $ ColN.of_odd ColN.zero) -- displays `fun x => 4 * x + 2`

[Sky Wilshaw]

You can see what would happen if you applied tactics to expressions by just putting them in a proof state:

example {x : ℕ} : 1 + 1 = x := by
  ring
-- x: ℕ
-- ⊢ 2 = x

I put the x in there to make the goal false, this stops the tactics from automatically solving the goal for me.

[Sky Wilshaw]

There's probably a better way to do this though.

[Sky Wilshaw]

I don't think most tactics really know how to act on "bare" terms (of types in Type not Prop).

[Kevin Buzzard]

Yeah tactic mode is for proving theorems, not constructing data in general. Some tactics can make awful terms but because they're supposed to be proofs it doesn't matter

[Kyle Miller]

I don't recommend using it for anything but exploration, but there's a reduce tactic that does what #reduce does.

[Matthew Pocock]

Thanks. I was hoping that I could convince the tactic system to do code optimizations, to rewrite verbose and inefficient execution paths into better ones. But seems the machinery isn't keen on that :D I realise the proofs themselves are a nightmare, but it's the target expression that I'm interested in tidying as an optimized expression equivalent to my original one. The proof doesn't matter to me. Just the rewritten term.

[Matthew Pocock]

I'll chalk it up as a "not ever intended to work"

[Yakov Pechersky]

Rewriting terms like this, for reduction purposes, is what norm_num and expression based "calculators" do. You could write a custom tactic or elaborator for your specific use case.

[Joachim Breitner]

Were you thinking of something like this:

import Lean.Meta.Tactic.Simp
import Mathlib.Tactic.RunCmd
import Std.Tactic.TryThis

open Lean

def pretty (e : Expr) : Elab.TermElabM Format := do
    let s ← Std.Tactic.TryThis.delabToRefinableSyntax e
    let f ← PrettyPrinter.ppCategory `term s
    pure f

elab "simp_expr " e:term : command => Lean.Elab.Command.liftTermElabM $ do
    let expr ← Lean.Elab.Term.elabTerm e none
    let ctxt ←  Lean.Meta.Simp.Context.mkDefault
    let (r, _) ← Lean.Meta.simp expr ctxt
    logInfo $ m!"Before simplifying:\n{← pretty expr}\n" ++
              m!"After simplifying:\n{← pretty r.expr}"
    -- Std.Tactic.TryThis.addTermSuggestion e r.expr

simp_expr (id ())

This prints

Before simplifying:
id ()
After simplifying:
()

[Joachim Breitner]

(I was actually pondering to (ab)use the simplifier like that for a little experiment of mine. I don’t see why it shouldn’t work.)

[Alex J. Best]

In mathlib3 there was a simp_result tactic combinator that simplified the results of other tactics before assigning them. But as far as I know it wasn't ported yet

[Matthew Pocock]

Joachim Breitner said:

Were you thinking of something like this:

Yes, that's a lot more like what I was thinking of. Thanks.

[Mario Carneiro]

Isn't that the #simp command?

[Joachim Breitner]

It seems that for what I wanted to do, it doesn’t quite work: I was hoping to (ab)use the simp machinery as a convenient way to perform recursive term rewrites to Expr. But my expressions are sometimes Prop, and then simp will simply say “done” because two proofs are definitionally equal. Too bad.

[Joachim Breitner]

Isn't that the #simp command?

Indeed! Didn’t even know about that :-)

[Joachim Breitner]

I assume this #simp doesn’t seem to do do anything because the expression is already a Prop:

import Mathlib.Tactic.Conv

set_option trace.Meta.Tactic.simp true
set_option trace.Meta.Tactic.simp.heads true
set_option trace.Debug.Meta.Tactic.simp true

def other_eq_self (a : Prop) : (a = a) = True := eq_self a
theorem other_eq_self_eq_eq_self (a : Prop) :
    other_eq_self a = eq_self a := rfl
#simp [other_eq_self_eq_eq_self] => other_eq_self True

Ah, here it is, from the core of simp:

partial def simp (e : Expr) : M Result := withIncRecDepth do
  checkMaxHeartbeats "simp"
  let cfg ← getConfig
  if (← isProof e) then
    return { expr := e }

So no chance using simp to rewrite proof terms :-)

[Tomas Skrivan]

I'm have not read carefully this thread so I'm not sure if I'm relevant with this but I have macro term rewrite_by conv that rewrite term by conv tactic.

For example

#check (id 42) rewrite_by simp

prints 42 : Nat.

code

import Lean

open Lean.Parser.Tactic.Conv
open Lean.Elab.Tactic.Conv

open Lean Elab Term Meta

def elabConvRewrite (e : Expr) (stx : TSyntax `conv) : TermElabM (Expr × Expr) := do
  let (rhs, eq) ← mkConvGoalFor e

  let goals ← Tactic.run eq.mvarId! do
    let (lhsNew, proof) ← convert e (Tactic.evalTactic stx)
    updateLhs lhsNew proof
    return ()

  if goals.length = 0 then
    throwError "this is a bug in rewriteByConv"

  if goals.length > 1 then
    throwError s!"error in `rewriteByConv`, unsolved goals {← goals.mapM (fun g => do ppExpr (← g.getType))}"

  (goals.get! 0).refl

  return (← instantiateMVars rhs, ← instantiateMVars eq)

def rewriteByConv (e : Expr) (stx : TSyntax `conv) : MetaM (Expr × Expr) := do
  let (r,_) ← (elabConvRewrite e stx).run {} {}
  return r

syntax:1 term "rewrite_by" convSeq : term

elab_rules : term
| `($x:term rewrite_by $rw:convSeq) => do
  let x ← elabTerm x none
  let (x',_eq) ← elabConvRewrite x (← `(conv| ($rw)))
  return x'

[Tomas Skrivan]

Matthew Pocock said:

Thanks. I was hoping that I could convince the tactic system to do code optimizations, to rewrite verbose and inefficient execution paths into better ones. But seems the machinery isn't keen on that :D I realise the proofs themselves are a nightmare, but it's the target expression that I'm interested in tidying as an optimized expression equivalent to my original one. The proof doesn't matter to me. Just the rewritten term.

I'm definitely using Lean for this purpose. I had to modify the simplifier a bit but doing this kind of stuff is totally possible.

I'm planning on modifying the simplifier even more such that it handles let bindings a bit better and I will add an option to turn off proof generation such that on massive programs is runs reasonably fast.

[Tomas Skrivan]

An example of code optimization, I have conv tactic let_normalize that cleans up let bindings and moves them out of lambdas if possible

#check
    (fun (x y : Nat) =>
      let a := x + 1
      let b := x + a
      let c := x + y
      a+b+c)
    rewrite_by
      let_normalize

prints

fun x =>
  let a := x + 1;
  let b := x + a;
  fun y =>
  let c := x + y;
  a + b + c

[Kyle Miller]

@Tomas Skrivan mathlib has a lift_lets tactic, but I don't recall a conv version yet. There are conv tactics at least for reduce, whnf, projection unfolding, and a few others.

[Kyle Miller]

@Joachim Breitner You might be interested in the elementwise attribute in mathlib (look in the category theory folder in the tactic folder). It uses a function to simp the lhs and rhs independently when generating a lemma statement.

[Matthew Pocock]

@Tomas Skrivan That looks interesting. It works for me in some simple examples. I tried to capture the rewritten form inside a def, but failed.

inductive ColN where
    | zero      : ColN
    | even_of   : ColN → ColN
    | odd_of    : ColN → ColN
deriving Repr

@[simp]
def ColN.eval (cn: ColN) (n : ℕ) : ℕ := match cn with
| ColN.zero => n
| ColN.even_of cm => 2 * (ColN.eval cm n)
| ColN.odd_of cm => 2 * (ColN.eval cm n) + 1

def z10 := ColN.even_of $ ColN.odd_of ColN.zero

#reduce z10.eval -- bleutch!
#check z10.eval rewrite_by simp ; ring -- gives 2+n*4, which is exactly what we would want :)

-- try to capture the rewritten form
def z10eval := (z10.eval rewrite_by simp ; ring)

#reduce z10eval -- bleuch again!
#check z10eval

[Tomas Skrivan]

You should not use #reduce z10eval, calling #reduce fun n => 2 + n * 4 gives you the same result as you are reducing through multiplication on Nat. To see the value of z10eval do #print z10eval and you get

def z10eval : ℕ → ℕ :=
fun n => 2 + n * 4

as desired, you can also play with it using rewrite_by as

#check
  z10eval
  rewrite_by
    unfold z10eval

[Tomas Skrivan]

Kyle Miller said:

Tomas Skrivan mathlib has a lift_lets tactic, but I don't recall a conv version yet. There are conv tactics at least for reduce, whnf, projection unfolding, and a few others.

Nice! Much cleaner implementation then mine. For my use case I need few more reductions like splitting let a := (x,y); f a into let a1 := x; let a2 := y; f (a1,a2) but that should be easy to adopt.

[Tomas Skrivan]

@Kyle Miller Ideally I would like lift_lets simplifications happen when I run simplifier, any idea what would be the easiest way to do this?

[Tomas Skrivan]

I have lots of simp rules that create nested let bindings or let bindings like let a := (x, 0) which I want to split and eliminate the zero.