feat: add pretty run report

[kaeun97]

Resolves #398.

Here is an example code:

from __future__ import annotations
from egglog import *

egraph = EGraph()

class Num(Expr):
    def __init__(self, n: i64Like) -> None: ...
    def __add__(self, other: Num) -> Num: ...
    def __mul__(self, other: Num) -> Num: ...

x, y = vars_("x y", Num)
egraph.register(rewrite(x + y).to(y + x))
egraph.register(Num(1) + Num(2))
report = egraph.run(10)
print(report)

Output before:

RunReport { iterations: [IterationReport { rule_set_report: RuleSetReport { changed: true, rule_reports: {"(rewrite (__main___Num___add__ _x _y) (__main___Num___add__ _y _x))": [RuleReport { plan: None, search_and_apply_time: 2.625µs, num_matches: 1 }]}, search_and_apply_time: 5.375µs, merge_time: 583ns }, rebuild_time: 1.125µs }, IterationReport { rule_set_report: RuleSetReport { changed: false, rule_reports: {"(rewrite (__main___Num___add__ _x _y) (__main___Num___add__ _y _x))": [RuleReport { plan: None, search_and_apply_time: 1.125µs, num_matches: 1 }]}, search_and_apply_time: 2.75µs, merge_time: 1.041µs }, rebuild_time: 0ns }], updated: true, search_and_apply_time_per_rule: {"(rewrite (__main___Num___add__ _x _y) (__main___Num___add__ _y _x))": 3.75µs}, num_matches_per_rule: {"(rewrite (__main___Num___add__ _x _y) (__main___Num___add__ _y _x))": 2}, search_and_apply_time_per_ruleset: {"": 8.125µs}, merge_time_per_ruleset: {"": 1.624µs}, rebuild_time_per_ruleset: {"": 1.125µs} }

Output after:

PrettyRunReport(iterations=[PrettyIterationReport(rule_set_report=PrettyRuleSetReport(changed=True, rule_reports={'rewrite(x + y).to(y + x)': [PrettyRuleReport(plan=None, search_and_apply_time=datetime.timedelta(0), num_matches=1)]}, search_and_apply_time=datetime.timedelta(0), merge_time=datetime.timedelta(0)), rebuild_time=datetime.timedelta(0)), PrettyIterationReport(rule_set_report=PrettyRuleSetReport(changed=False, rule_reports={'rewrite(x + y).to(y + x)': [PrettyRuleReport(plan=None, search_and_apply_time=datetime.timedelta(0), num_matches=1)]}, search_and_apply_time=datetime.timedelta(0), merge_time=datetime.timedelta(0)), rebuild_time=datetime.timedelta(0))], updated=True, search_and_apply_time_per_rule={'rewrite(x + y).to(y + x)': datetime.timedelta(0)}, num_matches_per_rule={'rewrite(x + y).to(y + x)': 2}, search_and_apply_time_per_ruleset={'': datetime.timedelta(0)}, merge_time_per_ruleset={'': datetime.timedelta(0)}, rebuild_time_per_ruleset={'': datetime.timedelta(0)})

[saulshanabrook]

Thank you for this! Added a few comments. Could you also add this to the changelog file with a link to this PR?

[codspeed-hq]

Merging this PR will improve performance by 67.74%

⚡ 1 improved benchmark
✅ 5 untouched benchmarks
⏩ 8 skipped benchmarks¹

Performance Changes

	Mode	Benchmark	BASE	HEAD	Efficiency
⚡	Simulation	test_jit[lda]	11.6 s	6.9 s	+67.74%

Tip

Curious why this is faster? Comment @codspeedbot explain why this is faster on this PR, or directly use the CodSpeed MCP with your agent.

---

_{Comparing kaeun97:kaeun97/pretty-report (01802ec) with main (8812ec9)}

Footnotes

1. 8 benchmarks were skipped, so the baseline results were used instead. If they were deleted from the codebase, click here and archive them to remove them from the performance reports. ↩

[saulshanabrook]

Thanks for the fixes, I left a few small comments. There are also some mypy and formatting issues I think.

There is a bigger question about performance, if the codspeed is correct it looks like this slows things down by a ton!

Taking almost 40% of the time in a bigger benchmark just to translate bindings.

It makes me wonder about a different approach, where we set each rewrite and rule with a manual name like 1, 2, 3, ... and then we don't have to do the name searching and mangling and can just parse the name as an int then look it up? And if it's a birewrite just take off the <= or >=?

It would make the egglog file a bit more verbose, but makes parsing the reports more straightforward and more performant which seems like a good tradeoff?

I was also going back and forth on whether the RunReport should store a RewriteOrRule or the decl? If we just store the RewriteOrRule it's easier to pretty print, can just use the builtin one, and it's easier for users to grab that off and compare it or use it... But most of the other exposed objects just store the decls, so I will leave it up to you!

EDIT: It looks like the docs failures also highlight some other exceptions from this. I imagine also if we name the rules here that might also help since it seems like it's hitting on looking up the string?

[kaeun97]

@saulshanabrook Thanks for the thorough review! I do agree that the performance looks concerning. The numeric name approach you mentioned would work for bindings with a "name" field - so not for, RewriteDecl nor BiRewriteDecl. That would require the rust side change. Happy to prioritize that before continuing on with this PR. Also, we can do lazy loading (translate when the user needs it) to have minimal impact to performance.

[saulshanabrook]

The numeric name approach you mentioned would work for bindings with a "name" field - so not for, RewriteDecl nor BiRewriteDecl. That would require the rust side change. Happy to prioritize that before continuing on with this PR.

Ah yeah I kept forgetting about this! I just talked to some other folks on the egglog team and they said that sounds like a great feature to add, just something we hadn't gotten around to yet. It should also I think be relatively straightforward so a good first PR to egglog core if you don't mind doing that...

Then once that is merged hopefully should just be able to update the pin here and can use that feature. I believe the version of egglog we depend on here is pretty recent, so hopefully won't be other changes we have to adapt to.

[kaeun97]

Hey @saulshanabrook , thank you again for your feedback. The benchmark seems much better now. Let me know how it looks!

[saulshanabrook]

Thanks again for your continued updates on this!

I have some additional cleanup feedback, to try and keep the data structures a bit more minimal and specific, raise any errors earlier, and make sure bi-rewrite preserves both times.

[saulshanabrook]

Since we now support name for rewrite/birewrite, could we expose this to the user level as well? And then this logic here would be similar to the RuleDecl handling, where it checks for an explicit name and if it doesn't have one generates one. This would entail adding the name to pretty.py, declarations.py and egraph.py I believe.

This isn't strictly necessary for this PR though so if you don't feel like doing this here that's fine.

[saulshanabrook]

Can we instead just use rule_name_to_command_decl, so we can remove this additional mapping and there is just one source? We will know which ones are named, because we can see if the CommandDecl has a name or not. We can also update it to be more specific and just go from str to RuleDecl | BiRewriteDecl | RewriteDecl I believe.

[saulshanabrook]

What if we remove this, and instead store in the rule_name_to_command_decl version for <= and => when adding a bi-rewrite? Then that structure should always include all egglog rules we output, so we can do a lookup and if it's missing the exception just percolates up, avoiding a silent failure?

[saulshanabrook]

What if we just store CommandDecl's here regardless of if it has a name or not, then just change the repr/str to display it the name as a string if has one, otherwise pretty print the full command?

[saulshanabrook]

When we build these dictionaries from the bindings, could we check for duplicate keys (either named or unnamed) and combine the values for them? So that for BiRewrite, we don't lose the first one?

[saulshanabrook]

Can you add RunReport to __all__