Indexing lambdas with `DiscrTree`

Stream: lean4 dev

[Tomas Skrivan]

indexing lambdas: DiscrTree does not index lambdas properly, every lambda is just represented by a wildcard. This is problematic for tactics like continuity etc. Continuity of foo is usually stated as ∀ f, Continuous fun x => foo (f x) and DiscrTree stores this under the key Continuous _ i.e. there is no information about foo. The effect is that if we want to prove for example Continuous fun x => x + exp x, DiscrTree returns candidate theorems:

∀ (f : Nat → Nat), Continuous fun x => exp (f x)
∀ (f : Nat → Nat), Continuous fun x => sin (f x)
∀ (f g : Nat → Nat), Continuous fun x => f x + g x
...

Ideally only the third theorem would match. Returning all these theorems will be an issue once the number of continuity theorems become large.

---

eta reduction: Somewhat related problem is with eta reduction. Theorems like ∀ (f : Nat → Nat), Continuous fun x => exp (f x) do not match Continuous fun x => exp x. This is because Continuous fun x => exp x is eta reduced and and stored under the key Continuous exp which does not match Continuous _.

This becomes confusing with with functions with multiple arguments. Continuity of addition

∀ (f g : Nat → Nat), Continuous fun x => f x + g x

matches Continuous fun x => x + c but does not match Continuous fun x => c + x. Thus and additional theorem

∀ (c : Nat), Continuous fun x => c + x

is necessary. This might not be a problem as it is solvable with tooling that just generates these theorems automatically.

---

I would be interested in fixing DiscrTree to properly index lambdas. It would be very useful to me as I'm working on a general tactic for proving continuity, linearity, differentiability, ... and tactic for computing derivatives, adjoints, .... This is a also and issue for existing continuity tactic implemented with aesop. Would this be a desirable change? Should DiscrTree support this or should there be a variant of DiscrTree that supports this? I would be happy to implement this. Has anyone thought about it and know how to do it?

---

mwe demonstrating described behaviors:

code

import Lean
import Qq

open Lean Qq Meta

opaque Continuous (f : α → β) : Prop
opaque exp (n : Nat) : Nat
opaque sin (n : Nat) : Nat


def addTheorem (tree : DiscrTree Expr false) (e : Expr) : MetaM (DiscrTree Expr false) := do
  let (_, _, type) ← withReducible <| forallMetaTelescopeReducing e
  tree.insert type e

def printMatches (tree : DiscrTree Expr false) (e : Expr) : MetaM Unit := do
  let (_, _, type) ← withReducible <| forallMetaTelescopeReducing e
  let ms ← tree.getMatch type
  IO.println s!"Matches for: `{← ppExpr e}`"
  for m in ms do
    IO.println s!"  {← ppExpr m}"
  IO.println ""

#eval show MetaM Unit from do

  let mut t : DiscrTree Expr false := {}

  -- 1
  t ← addTheorem t q(∀ (f : Nat → Nat), Continuous fun x => exp (f x))
  -- 2
  t ← addTheorem t q(Continuous fun x => exp x)
  -- 3
  t ← addTheorem t q(∀ (f : Nat → Nat), Continuous fun x => sin (f x))
  -- 4
  t ← addTheorem t q(Continuous fun x => sin x)
  -- 5
  t ← addTheorem t q(∀ (f g : Nat → Nat), Continuous fun x => f x + g x)
  -- 6
  t ← addTheorem t q(∀ (c : Nat), Continuous fun x => c + x)


  -- matches 2
  printMatches t q(Continuous fun x => exp x)
  -- matchs 1,3,5
  printMatches t q(Continuous fun x => exp (id x))
  -- matchs 1,3,5
  printMatches t q(Continuous fun x : Nat => x + x)
  -- matches 1,3,5
  printMatches t q(∀ (c : Nat), Continuous fun x : Nat => x + c)
  -- matches 6
  printMatches t q(∀ (c : Nat), Continuous fun x : Nat => c + x)
  -- matches 1,3,5
  printMatches t q(Continuous fun x => x + exp x)

  IO.println t.format

[Joachim Breitner]

Possibly relevant research: The “Triemaps that match” paper by SPJ and @Sebastian Graf

[Joachim Breitner]

Although maybe you can do something much simpler where you go under lambdas, but without tracking the arguments; i.e. don’t distinguish fun x =. f x from fun x => f *. This will match more lemmas, but it’s an index only.
(all just vague thoughts so far)

[Tomas Skrivan]

What I'm doing right now is to index these theorems by two things 1. the property, Continuous, Differentiable, ... and 2. the function it talks about/head of the lambda body, sin, exp, HAdd.hAdd, .... This works fine but becomes problematic with bundled morphisms where the head of the lambda body becomes FunLike.coe.

[Jannis Limperg]

I'd be very keen on this, for obvious reasons. :sweat_smile:

[Tomas Skrivan]

Joachim Breitner said:

Although maybe you can do something much simpler where you go under lambdas, but without tracking the arguments; i.e. don’t distinguish fun x =. f x from fun x => f *. This will match more lemmas, but it’s an index only.
(all just vague thoughts so far)

Based on your idea I added

    | .lam _ t b _ =>
      let mx ← mkFreshExprMVar t
      let b' := b.instantiate1 mx
      pushArgs root todo b'

to this line and

  | .lam _ t b _ =>
    let mx ← mkFreshExprMVar t
    let b' := b.instantiate1 mx
    getKeyArgs b' isMatch root

to this line

and it works as expected

[Tomas Skrivan]

It effectively ignores lambdas and replaces bound variables with wildcards.

[Tomas Skrivan]

... and lean builds with this modification in core ... onto mathlib then :)

[Tomas Skrivan]

mathlib immediately fails to build because type class error in proofwidgets lakefile

error: typeclass instance problem is stuck, it is often due to metavariables
   FamilyOut BuildData (BuildInfo.key (Package.target `widgetPackageLock pkg.toPackage)) (BuildJob FilePath)

[Joachim Breitner]

Hmm, not sure it works that simply; I'd expect you need some new Key for that as well. And then maybe eta-expand when matching something that isn't a lambda against a lambda pattern.

Does anyone know what other systems use here?

[Tomas Skrivan]

I added a new key, commit, but there are still issues.

For example, the theorem id_map' itself is not provable by simp anymore.

[Jovan Gerbscheid]

(deleted)

[Jovan Gerbscheid]

I've made my own version of DiscrTree which is able to index lambdas. The new definition of Key is

  | const  : Name → Nat → Key
  | fvar   : Nat → Nat → Key
  | bvar   : Nat → Nat → Key
  | star   : Nat → Key
  | lit    : Literal → Key
  | other  : Key
  | lam    : Key
  | forall : Key
  | proj   : Name → Nat → Nat → Key

I added the lambda key just like you, and I replaced the arrow key by the forall key to allow matching with dependent arrows as well. I also added the bvar key for indexing bound variables (using De Bruijn indexing). It is for example able to match the left side of Finset.sum_range_id, and it's able to match the left side of Classical.skolem.

[Jovan Gerbscheid]

I use this version of DiscrTree for point & click library search, which means that you can click on a subexpression, and then choose a library result from a shortlist to apply it or rewrite with it.

[Tomas Skrivan]

Cool, would you mind sharing your code? Did you try to modify lean to use your version off DiscrTree, does it work?

[Jovan Gerbscheid]

I was a bit scared of trying to get this into Lean, because I though many other things could break due to the changes. Also, the code isn't the cleanest it could be.

Here's the code:
https://github.com/Human-Oriented-ATP/lean-tactics/tree/main/MotivatedMoves/LibrarySearch

[Jovan Gerbscheid]

I made my own version of DiscrTreeTypes.lean and DiscrTree.lean. The code works well. Another thing I've added to the unification is a score that indicates how good the match is. So for example Nat.add_comm gets a higher score than add_comm when matching with a sum of two natural numbers. This is very useful for giving an order to the list of search results, and then you can limit yourself to a fixed length initial segment of this list.

[Jovan Gerbscheid]

I also noticed the eta reduction problem, and I thought it could be solved by putting all functions into eta unreduced form. But it didn't seem like a big problem to me so I hadn't thought about it more.

[Jovan Gerbscheid]

I also made it so that star patterns are now indexed by an integer, so that a metavariable that appears multiple times can be given the same star. In particular this means that the right side of two_mul will only match with a sum of two equal terms, and if it does, then it gets a higher matching score than add_comm.

[Tomas Skrivan]

Very nice! I will be playing with your version of DiscrTree for some time :)

[Tomas Skrivan]

I think it would be great to either figure out if core DiscrTree can be modified this way or if it is fundamentally incompatible then put it in std or mathlib as it would be really useful to other people.

[Jovan Gerbscheid]

There's multiple things that I changed, so we could also try adding them to DiscrTree one by one to see which ones work and which ones don't.

[Jovan Gerbscheid]

I also changed the fvar key to take a Nat index (just like star) instead of an FVarId, which I would use for existentially bound variables in lemmas like Dense.exists_between. The FVarId indexing is only useful when you add local hypotheses to a (temporary) DiscrTree, since the FVarId is generated just for this local context. I don't do this, but I guess this might be used somewhere in Lean.

[Tomas Skrivan]

Just tested your version of DiscrTree and it indexes lambdas as expected in the examples I posted!

[Jovan Gerbscheid]

What would be the best way to go forward with this? Should I try PRing changes in DiscrTree to gradually build in all of the new features that I have in my version?

[Tomas Skrivan]

We should check with @Leonardo de Moura if indexing of lambdas is even a desirable change to DiscrTree in core. It might be desirable but would require a lot of care/work and it is probably low priority.

If it is not desirable then probably add an alternative DiscrTree to the standard library. It has its own version of HashMap so why not DisrcTree.

[Leonardo de Moura]

@Tomas Skrivan Thanks so much for thinking of me for this discussion! You’re absolutely right that higher-order matching is a complex challenge. Currently, the DiscrTree is not designed to handle this complexity. However, I do appreciate the potential benefits such changes could bring to certain projects. While I believe that these types of workarounds described above might not be the best fit for core or the standard library, it's great to explore these ideas, and include them in other packages. Thanks again for including me in this conversation!

[Jovan Gerbscheid]

Ok, so in that case I feel like the easiest way forward might be to make a PR for my version of DiscrTree in Mathlib. In particular, I want to have my library search tactic in Mathlib which uses my DiscrTree, but that will come later.

[Scott Morrison]

@Jovan Gerbscheid, note that the existing exact? is moving to Std. You might want to target Std if you are planning on having an alternative DiscrTree which can be dropped in as a replacement.

[Jovan Gerbscheid]

For the eta-reduction problem I had two solutions in mind, and I was wondering which is considered better.

The problem is that a function fun x => f x sometimes can and sometimes cannot be eta-reduced. e.g. fun x => 1 + x is just 1 +, but fun x => x + 1 cannot be reduced. Ideally both would match with the pattern fun x => f x + g x from the library, but that doesn't happen.

This example suggests it might be a good idea to always introduce lambda's whenever possible, but for example if f g : Nat → Nat, then we want f + g to still match library results about addition, which wouldn't work if this was automatically turned into fun x => (f + g) x.

So a way we can solve it instead, is at every point during the DiscrTree search, when our expression is a function that isn't in a lambda form, in addition to the usual search, we turn the expression into a lambda function (the opposite of eta-reduction), and then try finding this in the DiscrTree.

However, instead of doing more work during the DiscrTree search, we could do more work while building the DiscrTree. For a lemma involving fun x => f x + g x, we should realize that certain expressions that should match with this will get eta-reduced, so we can create two different DiscrTree paths, one corresponding to fun x => f x + g x and one to f +. Doing this ensures that even if our search query is "accidentally" eta-reduced, we still find this lemma.

What do you think?

[Tomas Skrivan]

Do I understand it correctly that if you add an entry with a path:

#[λ,x₁,...,xₙ,*]

then you add an additional entry with the path

#[x₁,...,xₙ]

Does this capture the eta reduction? I think this might be creating extra false positives?

[Jovan Gerbscheid]

Yes that's the idea, but only if the * at the end is an argument of x₁, and not if e.g. * is the last argument of the last argument of x₁.

[Jovan Gerbscheid]

I don't directly see how a false positive would appear.

[Tomas Skrivan]

Sounds good and should fix the issues I had with eta reduction. Let me know if you implement it in your DiscrTree I would like to try it out.

[Scott Morrison]

@Jovan Gerbscheid, I would love to see this integrated with exact?!

[Tomas Skrivan]

Wouldn't the statement Continuous fun x => f x + f x create false positives? I think the path is #[Continuous, λ, HAdd.hAdd, *, *] and it can be 'eta reduced' to #[Continuous, HAdd.hAdd, *] but the statement does not unify with Continuous (HAdd.hAdd c)

[Eric Wieser]

That eta reduction sounds like nonsense to me

[Eric Wieser]

Do things become messy if we introduce a key for the i-th-most-recent bound variable?

[Tomas Skrivan]

Eric Wieser said:

That eta reduction sounds like nonsense to me

Can you be more specific? Is the problem or the solution of eta reduction nonsense?

[Tomas Skrivan]

Eric Wieser said:

Do things become messy if we introduce a key for the i-th-most-recent bound variable?

This would not work for the continuity lemma because when you are constructing the key/path f and g are metavariables and are replaced with a wildcard. Thus you can't see those bound variables through it.

[Jannis Limperg]

This might be somewhat relevant: https://dl.acm.org/doi/10.1145/1614431.1614437

[Jovan Gerbscheid]

I've implemented the feature that the DiscrTree encoding of expressions now gives all possible eta equivalent forms. For small imports it seems to work fine. Looking at heartbeats, building the DiscrTree runs about 1.5x longer. When I tried building the DiscrTree for all of Mathlib, my VScode kept crashing. It also reaches the 1 000 000 heartbeats which is the default maximum for building it (the previous version got up to around 820 000)

I put the code in a new branch:
https://github.com/Human-Oriented-ATP/lean-tactics/tree/DiscrTree/MotivatedMoves/LibrarySearch

[Joachim Breitner]

You could try this Trie implementation, which adds path compression, and may help with some of the performance issues: https://github.com/leanprover/lean4/pull/2577

[Jovan Gerbscheid]

I've implemented a path structure in my version of DiscrTree as well, thanks to your earlier suggestion. I'm afraid though that this doesn't help much with the time required for building the DiscrTree, more with using it.

[Jovan Gerbscheid]

I notice that your version still has an empty constructor and can have values anywhere along the tree, but this is unnecessary since all values are at the leaves of the tree

[Scott Morrison]

Jovan Gerbscheid said:

I've implemented the feature that the DiscrTree encoding of expressions now gives all possible eta equivalent forms. For small imports it seems to work fine. Looking at heartbeats, building the DiscrTree runs about 1.5x longer. When I tried building the DiscrTree for all of Mathlib, my VScode kept crashing. It also reaches the 1 000 000 heartbeats which is the default maximum for building it (the previous version got up to around 820 000)

I put the code in a new branch:
https://github.com/Human-Oriented-ATP/lean-tactics/tree/DiscrTree/MotivatedMoves/LibrarySearch

1.5x doesn't sound too bad. Remember for exact? this DiscrTree is built in CI, not by the user.

You can just bump the 1000000 heartbeat limit for building the tree!

[Jovan Gerbscheid]

(deleted)

[Jovan Gerbscheid]

It turned out that the implementation I had was a lot slower than I said, and had some bugs, but I after trying lots of things, I now have a version that is able to build on my device and manages to barely stay below that 1000000 heartbeat limit with 970000 heartbeats. Note that this heartbeat count is actually for 5 different DiscrTrees that I build at the same time. the biggest two are for rewriting and for applying. There is also one for applying in negative position (e.g. matching with p in p ⇨ q) and the final two are for ordered rewriting, one for positive and one for negative position. I tried my apply library search on Continuous Real.exp, and got the following results:

[(4, [Real.continuous_exp : Continuous Real.exp]),
 (3,
  [Continuous.exp : ∀ {α : Type u_1} [inst : TopologicalSpace α] {f : α → ℝ},
     Continuous f → Continuous fun y => Real.exp (f y)]),
 (1,
  [Continuous.congr : ∀ {α : Type u_1} {β : Type u_2} [inst : TopologicalSpace α] [inst_1 : TopologicalSpace β]
     {f g : α → β}, Continuous f → (∀ (x : α), f x = g x) → Continuous g,
....

As you can see it matches the obvious Real.continuous_exp with a score of 4. Importantly it also finds Continuous.exp, which gets a score of 3, and this is 1 lower because the domain of the Continuity statement isn't necessarily Real. So it successfully matched an eta reduced expression.

In finding the keys for DiscrTree, an Expr is now first turned into a discrimination tree expression which I defined as

inductive DTExpr where
  | const  : Name → Array DTExpr → DTExpr
  | fvar   : FVarId → Array DTExpr → DTExpr
  | bvar   : Nat → Array DTExpr → DTExpr
  | star   : MVarId → DTExpr
  | lit    : Literal → DTExpr
  | sort   : DTExpr
  | lam    : DTExpr → DTExpr
  | forall : DTExpr → DTExpr → DTExpr
  | proj   : Name → Nat → DTExpr → Array DTExpr → DTExpr

And I made the function flatten : DTExpr → Array Key for turning this into a sequence of Keys.

For obtaining the DTExpr, because the eta reduction gives multiple results, we have mkDTExprs : Expr → MetaM (List DTExpr). To do some of the computations that have multiple outcomes, I locally defined the Monad transformer

abbrev ListT (m : Type u → Type v)  a := m (List a)

With instances for Monad, MonadLiftT and MonadControl. It would be handy if something like this already existed in Lean, but I've heard that List is intentionally not a monad (maybe because of performance, but in this use case the lists are very short, so the performance is fine).

There is one lemma that I found that stood out because it has 256 possible translations to DTExpr, which isMeasureTheory.withDensitySMulLI_apply. I decided to specifically exclude it from the DiscrTree.

The code is now on the main branch again:
https://github.com/Human-Oriented-ATP/lean-tactics/tree/main/MotivatedMoves/LibrarySearch

[Kevin Buzzard]

This is https://leanprover-community.github.io/mathlib4_docs/Mathlib/MeasureTheory/Function/L1Space.html#MeasureTheory.withDensitySMulLI_apply . Is there really supposed to be two consecutive coercions in that statement?

[Jireh Loreaux]

I think so. The inner one should be the Subtype.val to the space of equivalent classes of strongly measurable functions, and the outer is docs#MeasureTheory.AEEqFun.cast

[Jovan Gerbscheid]

I realize that the number 256 comes from doing a case split on the same lambda expression over and over, but once the choice has been made in one place, it should be the same everywhere else. If it's not too expensive I could locally cache the choices that have been made.

[Jovan Gerbscheid]

I've implemented this now, and optimized the implementation further, so building the DiscrTree now takes very close to the amount of time I had before. :smiley:

[Jovan Gerbscheid]

How should I go about getting this into Std? Do I simply make a pull request?

I also use a monad transformer called StateListT, which I defined in a separate file. It is equivalent to wrapping StateT around ListT (which I also defined separately), but turns out to be more efficient (saving 5% on heartbeats in building the DiscrTree). I guess I would have to add this to Std as well then. Unless there is some better alternative.

[Jovan Gerbscheid]

How do I get permission for pushing into Std?

[Scott Morrison]

You just open a PR.

[Scott Morrison]

(from your fork)

[Scott Morrison]

I would really like to see this integrated into exact?. Is that your plan?

[Scott Morrison]

(exact? is currently still in Mathlib, but may be moving as soon as this week.)

[Jovan Gerbscheid]

I have some other library search functionality. I mainly want rewrite search in combination with a point & click system. I also have an apply search which is more general that lean's apply. But for now plugging it in into exact? is probably the easiest first step

[Jovan Gerbscheid]

I opened the PR

[François G. Dorais]

@Jovan Gerbscheid I didn't mean to overwhelm you with my PR comments. This is a great PR, it just needs a lot of menial work to get up to merging into Std. Please ask if you have any questions!

[Jovan Gerbscheid]

Thanks for the comments. I'm still a bit unsure about the StateListT situation. As you said, ideally a list monad structure is lazy, but simply replacing it with MLList also doesn't work. I'm curious what the general view is on list monads in lean, since it's such a key example of monads, and yet the instance Monad List isn't present in Lean, and ListT also doesn't exist in Lean.

[François G. Dorais]

The problem is that ListT and StateListT, as you defined them, are not a valid monad transformers. The bind_assoc rule may fail when the base monad is not commutative (e.g. IO and many others). It is fine as a "functor transformer" but that's not that useful. Specialized versions are totally fine though, even if the inner monad is not commutative, so long as consequences are understood and managed. In your use case, the inner monad is MetaM which I'm not sure is commutative so this might lead to unexpected behavior where actions are reordered. The alternative is to use MLList which is safer but harder to use. :/

[François G. Dorais]

PS: Since ListT is not actually used in your PR, it should be a separate PR. However, it will likely be rejected for reasons I mentioned.

[François G. Dorais]

PPS: Please reply to the PR in comments since there are no notifications for edits of the PR description.

[Jovan Gerbscheid]

That's actually very interesting, and I hadn't realized this problem. I made a minimal example of the non-commutativity:

def problem : ListT (StateM (Array Nat)) Unit := do
  ListT.mk $ pure [(),()]
  modify (·.push 0)
  modify (·.push 1)

#eval problem #[]

depending on the bracketing of the do block, you either get #[0,1,0,1] or #[0,0,1,1]. However with MLList, it always gives #[0, 1, 0, 1].

[Jovan Gerbscheid]

Also, I don't think I edited the PR description

[François G. Dorais]

(I probably was mistaken. No worries.)