marcusrossel/lean-egg

A (WIP) equality saturation tactic for Lean based on egg.

Equality Saturation Tactic for Lean

This repository contains a (work-in-progress) equality saturation tactic for Lean based on egg. The tactic is useful for automated equational reasoning. Checkout the Lean/Egg/Tests directory for examples.

Setup

This tactic requires Rust and its package manager Cargo. They are easily installed following the official guide.

To use egg in your Lean project, add the following line to your lakefile.lean:

require egg from git "https://github.com/marcusrossel/lean-egg" @ "main"

Then, add import Egg to the files that require egg.

Basic Usage

The syntax of egg is very similar to that of simp or rw:

import Egg

example : 0 = 0 := by
  egg

example (a b c : Nat) (h₁ : a = b) (h₂ : b = c) : a = c := by
  egg [h₁, h₂]

open List in
example (as bs : List α) : reverse (as ++ bs) = (reverse bs) ++ (reverse as) := by
  induction as generalizing bs with
  | nil          => egg [reverse_nil, append_nil, append]
  | cons a as ih => egg [ih, append_assoc, reverse_cons, append]

But you can use it to solve some equations which simp cannot:

import Egg

variable (a b c d : Int)

example : ((a * b) - (2 * c)) * d - (a * b) = (d - 1) * (a * b) - (2 * c * d) := by
  egg [Int.sub_mul, Int.sub_sub, Int.add_comm, Int.mul_comm, Int.one_mul]

Sometimes, egg needs a little help with finding the correct sequence of rewrites. You can help out by providing guide terms which nudge egg in the right direction:

-- From Lean/Egg/Tests/Groups.lean
import Egg

example [Group G] (a : G) : a⁻¹⁻¹ = a := by
  egg [/- group axioms -/] using a⁻¹ * a

And if you want to prove your goal based on an equivalent hypothesis, you can use the from syntax:

import Egg

example (h : p ∧ q ∧ r) : r ∧ r ∧ q ∧ p ∧ q ∧ r ∧ p := by
  egg [and_comm, and_assoc, and_self] from h

For conditional rewriting, hypotheses can be provided after the rewrites:

import Egg

example {p q r : Prop} (h₁ : p) (h₂ : p ↔ q) (h₃ : q → (p ↔ r)) : p ↔ r := by
  egg [h₂, h₃; h₁]

Note that rewrites are also applied to hypotheses.

Related Work

• Guided Equality Saturation introduces the original prototype of the egg tactic.
• Bridging Syntax and Semantics of Lean Expressions in E-Graphs contains a high-level description of how the tactic handles binders and definitional equality.

lean-egg's Issues

`reduce` Configuration Option

Add a configuration option that reduces the terms sent to egg. This is useful for expressions like x + y, which would normally elaborate to HAdd.hAdd .... With reductions we get Nat.add which results in smaller, more debuggable terms.

Remove Universe Level Erasure

Universe level erasure doesn't work. E.g. if we take the theorem ∀ (a : Nat), ULift.down (ULift.up a) = a then with universe level erasure we can apply it in both directions. If we now obtain a proof of 10 = 10 like:

(1) 10
(2) ULift.down (ULift.up 10)
(3) 10

... then we will instantiate the universe params in (2) using mvars. But the first of these params will never be resolved, as the expression 10 only resolves the second one to 0.

`egg?`

Have an egg? variant of the tactic which extracts a small amount of information from the proof reconstruction procedure (like mvar assignment and rewrite positions) and suggests a proof by egg! which takes this information.

Type Class Rewrites

To address the problem that type classes can block rewrites from occurring, add a rewrite rule of the form TypeClass.op = reducedOp for every type class appearing in the goal term or selected rewrite rule set.

This issue is especially problematic when passing definitions to egg. For example, egg [Nat.add] is going to result in two rewrites Nat.add ?x Nat.zero = ?x and Nat.add ?x (Nat.succ ?y) = Nat.succ (Nat.add ?x ?y). These use Nat.add and Nat.zero instead of the type class versions HAdd Nat Nat Nat and OfNat Nat 0. If we add the rules HAdd.hAdd Nat Nat Nat = Nat.add and OfNat Nat 0 = Nat.zero though, this shouldn't be a problem.

Proof Compression

mkCongrOf can support multiple cHoles at once, so when we have multiple subsequent rewrites which do not overlap (that is, the subexpression position of one is not a prefix of the other), perform both steps at once. I'm not sure how this jibes with tracing as that relies on single steps. But perhaps tracing can force proof reconstruction back into using single steps.

Use for Explanation Length Optimization

When rewriting in an untyped setting, the hope is that even though invalid rewrites can occur, since both the start and end terms are well-typed, it's unlikely that we'll find a bridge from one to the other via ill-typed terms. A potential problem is that when we do find such a bridge, it might be part of the explanation returned be egg, even when it may also have been possible to find a valid rewrite sequence. In that case proof reconstruction will fail, even though the given equality is actually provable. As a simple measure against this, then proof reconstruction fails, we can try to rerun the task with explanation length optimization enabled, which might prune the unnecessary ill-typed rewrites.

Special Handling for Eq Conditions

If a rewrite as a condition of the form a = b, then it would be nice if we could check if it is satisfied, without explicitly requiring a fact of this form. For example, this is useful for automated satisfaction of rfl-conditions like 0 = 0, which can pop up in if-then-else expressions.

More Efficient `validDirs`

The implementation of Rewrite.validDirs could be more efficient:

1. Extract mvars and level mvars in only one traversal of the expression.
2. When comparing the resulting sets, traverse one set, removing all matching vars from both sets. Then compare the resulting sets. If both are empty we get Directions.both. If the lhs is empty but the rhs isn't, we get .backward. Etc...

Non-Equivalence Goals (`egg [...] from`)

Currently, the egg tactic only works if the proof goal has the form a = b or p ↔ q. We can also support proving goals which do not have this form, by proving that a given hypothesis is equivalent to the goal. For example:

... 
h : p ∧ q ∧ r
⊢ r ∧ q ∧ p

That is, we ask egg to find the sequence of rewrites that gets us from h to the goal.
A syntax for this could be egg [...] from h.

E-Graph Visualization

Egg can generate dot/svg/png-files for its e-graphs. Try using the proofwidgets package to display this in the info-view.

Explosion for Inhabited Types

If explosion fails, check if the given type is inhabited and use its default value instead. Or if that fails, check the terms of the goal and given rewrites if they contain any members of the required type.

A very constructed example of this is:

def toTrue : Nat → Bool := fun _ => true
theorem toTrue_eq : ∀ n m, toTrue n = toTrue m := fun _ _ => rfl 

example : toTrue 0 = toTrue 1 := by
  egg [toTrue_eq]

Explosion will normally work here, as the proof context is empty.

More Natlit Defeqs

In Sebastian Ullrich's dissertation on page 36 he also mentions that rules like Nat.mul (nat_lit x) (nat_lit y) => (nat_lit x * y) hold definitionally. Find out the full set of these and implement them.

Level MVar Unification

When calling isDefEq on expressions containing level mvars, those level mvars aren't necessarily assigned. This causes problems, e.g. when using Config.Encoding.eraseConstLvls.
Using a syntactic unification like the example below won't work as we still need to be able to unify syntactically different expressions if we want to (a) use natlit-conversions and (b) skip proofs by refl in proof reconstruction.

-- Note: This function expects its input expression to be normalized (cf. `Egg.normalize`).
def unify : Expr → Expr → MetaM Bool
  | e₁@(.mvar _), e₂ | e₁, e₂@(.mvar _) => isDefEq e₁ e₂
  | e₁@(.bvar _), e₂@(.bvar _) | e₁@(.fvar _), e₂@(.fvar _) | e₁@(.lit _), e₂@(.lit _) =>
    return e₁ == e₂
  | .app e₁₁ e₁₂, .app e₂₁ e₂₂ | .lam _ e₁₁ e₁₂ _, .lam _ e₂₁ e₂₂ _
  | .forallE _ e₁₁ e₁₂ _, .forallE _ e₂₁ e₂₂ _ =>
    unify e₁₁ e₂₁ <&&> unify e₁₂ e₂₂
  | .sort lvl₁, .sort lvl₂ => isLevelDefEq lvl₁ lvl₂
  | .const name₁ lvls₁, .const name₂ lvls₂ => do
    unless name₁ == name₂ && lvls₁.length == lvls₂.length do return false
    let mut result := true
    for lvl₁ in lvls₁, lvl₂ in lvls₂ do
      result ← (return result) <&&> isLevelDefEq lvl₁ lvl₂
    return result
  | .mdata .., _ | _, .mdata .. | .proj .., _ | _, .proj .. | .letE .., _ | _, .letE .. =>
    panic! "'Egg.unify' received non-normalized expression"
  | _, _ => return false

Handle Invalid Rewrites in Rust

Currently when the Rust component finds an invalid rewrite it just crashes as we're using .unwrap(). Have Rust communicate the failure back to Lean instead.

Nat Literals

Literals for natural numbers (e.g. (lit 42)) might be problematic for some rewrites which expect an application of Nat.succ or Nat.zero. To handle this, we could add special conditional rewrite rules of the form:

(lit 0)                        <=> (const `Nat.zero)
(lit x)               [if x > 0]=> (app (const `Nat.succ) (lit x-1)) 
(app (const `Nat.succ) (lit x)) => (lit x+1)

Note that this way rewrite rules which match repeated occurrences of Nat.succ will also work, as the second rewrite can be applied multiple times. Also, when a rewrite expects to find a lit we can apply the reverse of the first rule or the third rule.

A potential problem with this approach is that it might blow up the e-graph.

Non-Recursive Equations

The equations produced for non-recursive definitions aren't useful. For example:

def f : Bool → Nat
  | false => 0
  | true => 1

produces the single equation:

(app (const f) ?106) = (app (app (app (app (const f.match_1 (succ 0)) _) ?106) _) _) 

As a result, this fails:

example : f false = 0 := by
  egg [f]

Cf. the related Zulip discussion for more info.

Could Julia help?

We have higher-level constructs, Julia has a nice FFI C interface, should be quite easy to call from C. Performance is not the same as egg but comparable (faster than egg compiled in debug mode, slower than --release mode, sits in between),

This branch is latest release candidate:
JuliaSymbolics/Metatheory.jl#185

We have stuff that could help for beta reduction. See tests/integration folder

Stack Overflow from β-reduction

import Mathlib
import Egg

set_option egg.betaReduceRws false
set_option egg.etaReduceRws false

set_option egg.genBetaRw false in
example (n : Real) : n / n = 1 := by
  egg -- does not crash

set_option egg.genBetaRw true in
example (n : Real) : n / n = 1 := by
  egg -- crashes with `Stack overflow detected. Aborting`

I'm guessing this is a stack overflow from too many recursive calls during β-reduction.
Perhaps we should implement a non-recursive version once explanations are handled correctly.

β-Reduction

In Mathlib's to_digits_core_lens_eq_aux, we run into:

simp only [hx, if_true, List.length, congrArg (fun l ↦ l + 1) hlen]

This works well with egg up to the congrArg (fun l ↦ l + 1) hlen whose type is:

(fun l ↦ l + 1) (List.length l1) = (fun l ↦ l + 1) (List.length l2)

... while the goal is:

List.length l1 + 1 = List.length l2 + 1

If we had a rule for β-reduction, this wouldn't be a problem. But we might also be able to solve this by checking for terms of the form (λ x, ...) y in the egg frontend and reducing them.

Explosion

When a rewrite is not allowed in both directions, try to generate rewrites for the missing directions by populating non-overlapping mvars with terms from the local context. While this acts as a simple heuristic for guessing rewrite steps like 0 => a + -a, it's extremely important in the context of rewrites involving typeclasses. For example, in the context of a fixed group G, the rewrite rule add zero a = a is applicable in both directions as the only mvar is a. But if we define it as an axiom of a group class (i.e. parameterized over a group class instance), then it is only applicable in the forward direction, as the LHS also contains mvars for the typeclass instance an class's type parameter (the alpha). Explosion would fix this by filling in the correct typeclass related mvars whose values should be present in the local context.

Note, the information of which group a belongs to is technically present on the RHS, as a will have some type that has a group instance. We just don't represent this syntactically in the untyped setting. Note that in a typed setting we could almost recover the bidirectionality of this rule as the type annotations would provide some of the missing mvars on the RHS. I think this wouldn't include the required typeclass instance though, so we'd still have to guess that by explosion.

Note also that the previous paragraph requires that we annotate rewrites with type tags before we determine their direction. Not sure how that would work in practice.

Tests by `simp only`

To test the egg tactic, we can try to replace simp only [...] calls in Mathlib with egg [...].
The following code snippet overrides the behavior simp only [...], such that we can call egg in its place.

-- "simp " (config)? (discharger)? ("only ")? ("[" (simpStar <|> simpErase <|> simpLemma),*,? "]")? (location)?
elab stx:simp : tactic => do
  let stx := stx.raw
  if stx[3].isNone then evalSimp stx else 
    if stx[4].isNone then 
      logInfo "1"
    else
      let args := stx[4][1].getSepArgs 
      if args.any (·.getKind == ``simpErase) then
        evalSimp stx
      else if args.any (·.getKind == ``simpStar) then
        logInfo "egg [*]"
      else if args.any (!·[0].isNone) then -- In this case we only have simpLemmas, but at least one of them has a leading ↑ or ↓.
        evalSimp stx
      else
        let terms := args.map (·[2])
        logInfo s!"egg {terms}"

To get the goal before and after the call to simp only [...], we can use the following scheme:

let s ← saveState
evalSimp stx
let goalAfterSimp ← getMainGoal
s.restore
let goalBeforeSimp ← getMainGoal