Feature request: High level Print-function

[adrianleh]

Requesting a high-level feature to run print-function on an EGraph and get the result as string or preferably high-level python objects

[saulshanabrook]

Yeah that sounds good! I agree that returning as high level Python objects would be ideal and I think doable.

It looks like internally, print-function, uses the function_to_dag method.

It's already public, so I think to implement this in Python we could:

1. Expose the function_to_dag in the low-level bindings. We already expose the TermDag elsewhere, so all the conversions for that are already taken care of.
2. Add a high level function to EGraph that uses this lowl level call and returns a proper Python object.

I was thinking the definition could be something like this:

class EGraph:
    [...]
    def extract_function(self, fn: Callable[..., T], n: int) -> list[tuple[T, T]]:
        """
        Returns a list of the first `n` calls to the `fn` in the EGraph.
        Each item in the list will be a pair of the function call and the lowest cost extract equivalent value.
        Analogous to the `print-function` command in egglog.
        """
        ...

Do you think that would work for you?

[adrianleh]

I think this would be a great idea to do! As for the specific implementation, the idea sounds great to be but I'm not sure if having the tuple have the same type is a good idea, as functions might have different return values compared to function call inputs (perhaps I'm misunderstanding here)?

[saulshanabrook]

Ah, both items in the tuple would be the return value. One would just be the function call itself, the other would be the lowest cost extracted version of it.

For example, this egglog code, which prints:

(
   (Mul (Num 2) (Add (Var "x") (Num 3))) -> (Mul (Num 2) (Add (Var "x") (Num 3)))
   (Mul (Num 2) (Var "x")) -> (Mul (Num 2) (Var "x"))
   (Mul (Num 2) (Num 3)) -> (Num 6)
)

would translate to this output of extract_function:

[
    (Math(2) * (Math.var("x") + Math(3)), Math(2) * (Math.var("x") + Math(3))),
    (Math(2) * Math.var("x"), Math(2) * Math.var("x")),
    (Math(2) * Math(3), Math(6)),
]

[adrianleh]

I see. How would this work with unextractable functions as in this example?

[saulshanabrook]

It would work similarily to the Rust output, since it uses the same underlying function. i.e. the first item in the tuple would be the unextractable function call, then the second item would be an equivalent extractable value.