Add method for extracting terms that match a pattern in python

[goombamaui]

I am using egglog to implement a few compiler simplification rules and one thing I want to do is use extraction based substitution. As a simplified example, if I have Expression.Let("x", 3, exp2) which corresponds to let x=3 in exp2, I could use egglog rules to simplify this using rewrites: rewrite(Expression.Let(vname, Expression.Int(ival), exp2)).to(...), but I want to instead defer this to a python function (IE: rule(e1 == <something>).then(union(e1).with_(f(e1))) where f is a python function). Essentially, I want to be able to have a rule rule(e1==<something>) and be able to passed the matched expression e1 (rebuilt in python) into a python function (and use methods like get_callable_args, etc. on it).

However, I do not think this can be done currently; py_eval_fn and PyObject.from_string work when trying to have a rule execute a python function that depends on string (or other primitive) values in the matched expression:

from egglog import *

class A(Expr):
    def __init__(self, v:StringLike): ...
def say_hi(x: String):
    print("Hi, " + str(x))

egraph = EGraph()

expr = egraph.let("expr", A("x"))
svar = var("svar", String)
avar = var("avar", A)
egraph.register(rule(avar==A(svar)).then(py_eval_fn(say_hi)(PyObject.from_string(svar))))
print("Starting")
egraph.run(1)

However, replacing PyObject.from_string(svar) with PyObject(avar) and modifying say_hi to use the matched object seems to fail. Perhaps one potential solution would be to have a PyObject.from_expr that translates a variable from its egraph representation back into python.

[saulshanabrook]

Could you give a little bit more of an example of what you want with:

I want to instead defer this to a python function (IE: rule(e1 == ).then(union(e1).with_(f(e1))) where f is a python function)

What kind of Python function? In this example you are trying to print in a rule. Is this the kind of thing you want to do?

Currently yes, there is very limited surface area for Python primitives, basically only interacting with PyObjects. These can be turned back and forth into primitives, but this would be the only way to really interact with egglog.

If you just want to eval, you could use py_eval:

from egglog import *

egraph = EGraph()

egraph.register(rule().then(py_eval("print('hello')")))
egraph.run(1)
# prints 'hello'

The title of the issue is "Add method for extracting terms that match a pattern in python". extract used to have the ability to match on "holes", so you could add variables into the value you wanted to extract and it would get the first one. This was removed from upstream egglog, but can be re-implemented. If this is what you are looking for, I can show you how to get the same behavior.

[goombamaui]

Thank you for the help!

py_eval is insufficient because the idea is I want access the expression matched in the rule when the rule fires.

In the rule I use printing as an example; ideally, I want to merge two expressions (so f(e1) does some processing to create a new expression e2 and merges e1 with e2). It would be nice to have something of the form
rule(e1 == <something>).then(union(e1).with_(f(e1))) where f(e2) calls the python function. It is sufficient to just add the element to a list and then after the ruleset finishes running, loop through the list, calculate f(e) for all e in the list, and then use egraph.register(union(...).with_(...)) to merge elements afterwards.

extract used to have the ability to match on "holes", so you could add variables into the value you wanted to extract and it would get the first one. If this is what you are looking for, I can show you how to get the same behavior.

This would be very helpful; I think a rule can simply flag an expression with a fact that is then extracted afterwards for the python function to run on.

[saulshanabrook]

Could you extrapolate on what f does here that is not possible to do with current API?

If you want union a list of expressions, you should be able to do this with the builtin Vec sort or with a custom linked list implementation, by recursing on the list and unioning each element till it is empty.

[goombamaui]

f can be theoretically implemented using rules, but the idea is to have f be a python function instead because it could be faster (in particular, skimming through this: https://arxiv.org/pdf/2111.13040, implementing de bruijn index shifting by extracting / inserting terms from the graph instead of using rules might be of interest)
Here is a simple example; the idea is instead of implementing the simplify function using egglog rules, the goal is to use the python function instead.

from __future__ import annotations
from egglog import *

class A(Expr):
    def __init__(self, v: i64Like) -> None: ...
    @classmethod
    def of(cls, left: A, right: A) -> A: ...
    @classmethod
    def simplify(cls, elem: A) -> A: ...

def simplify(elem: A):
    match get_callable_args(elem):
        case (i64(value),):
            return A(value+1)
        case (left, right):
            return A.of(simplify(left), simplify(right))

egraph = EGraph()
expr = egraph.let("expr", A.simplify(A.of(A(3), A(4))))
e1 = var("e1", A)
e2 = var("e2", A)
egraph.register(rule(e1==A.simplify(e2)).then(union(e1).with_(py_eval_fn(simplify)(PyObject(e2))))) # doesn't work
egraph.run(1)

Alternatively, we could extract all terms of the form A.simplify(e1) and then merge each one with simplify(e1) using the "match on holes" concept you mentioned earlier.

[saulshanabrook]

Yeah I can under conceptually why you might want to do this. However it's currently not supported. It might require changes as well upstream to expose enough. It also opens up the surface area quite a bit of the API.

I think what would be most helpful to drive a feature like this is a real use case showing how the current way of doing things is too slow. There might be ways to optimize it as well without making this large of a change. So like if you have an end to end example of replacing with de Bruijn indices with the current API, and then show that either it's too slow or not ergonomic enough for whatever reason, I think that would be really helpful for looming at the design space.

[goombamaui]

I wrote a simple implementation to test this; assuming only substitution and no other rules, python is faster than egglog for substituting in smaller expressions but egglog beats python when dealing with larger expressions (>256 terms in a binary tree like structure). It is possible that other rewrite rules might interact with substitution in a way that makes egglog much worse but in isolation, using egglogs rules are a more optimal solution.

I will use stick to implementing these functions using egglog's rules in the main implementation and will close this issue for now; if the substitution rules cause issues with the other rules and speed or memory becomes an issue, I will re-open it later. Thanks for the help!

[saulshanabrook]

Cool! That's interesting to hear. I have tried to keep things OK performance wise, but there is also probably lots of places things could be improved!

If you have an example that stretches things, feel free to add it as a benchmark here so we can keep track of it and maybe improve it. Especially if a significant amount of the performance is spent in Python, not in egglog, that would be important to fix here.

More examples here are also welcome for the docs.

The goal for the Python bindings is to at least not overshadow the egglog time on performance sensitive tasks.