Make Analysis::make take in the E-graph mutably

[imbrem]

I have a use case where I want to tag each value in the E-graph with type information, which itself should be stored either in the E-graph or in another E-graph stored in the Analysis. This requires me to mutate the E-graph or at least the Analysis from inside Analysis::make (the alternative would be to wrap a nested E-graph in the analysis in RefCell or RwLock); thankfully, just making the E-graph reference passed in mutable is enough to get this to work with only minimal changes (sticking an &mut on implementations of Analysis) to anything else. Unfortunately, this is still a breaking change, so not sure if you guys would be interested in it.

[yihozhang]

Does the solution proposed in #276 work for your setting?

[imbrem]

Not particularly well, since I'd like my type to be Sync since once the e-graph rewriting is done many threads might inspect it. I actually only really need mutable access to Analysis, so a compromise would be some kind of wrapper around &mut EGraph, but I feel like there's no reason not to just allow mutable access to the EGraph with a suitable warning

[mwillsey]

@imbrem I think wrapping your Analysis type in a mutex should work, right? I'm a little hesitant to break the API if we don't need to.

[imbrem]

@mwillsey well we'd need an RwLock for acceptable performance, but even that introduces overhead, not to mention hassles like e.g. manual Clone implementation, and to a degree having "too much freedom." Besides, it seems inconsistent with the rest of the API, which is passed a mutable Analysis (e.g., merge), to have make not receive a mutable analysis, especially since it doesn't break any of the test cases.

One potential non-breaking option is to make a new AnalysisMut trait and have a blanket impl for things which implement Analysis, but I don't think this is very clean.

[mwillsey]

Sure, that all makes sense. I still don't have a great understanding of your use case though. It seems like perhaps a "hook" of some kind on e-node insertion could suffice (and would not break the API).

[imbrem]

A simplified version of my use case is as follows: there's some kind of map in my analysis, and I want every E-node analysis update to be able to not only look at what's currently in the map, but also to be able to change the map so everything points to the right value. Doing this on insertion would not allow changing the value in the map; for my current needs this is possible to work around, but again it requires a bit of "yoga" (I need to, for example, have a guarantee that I did the insertion in the hook which runs before make, or trying to run the make function will panic), and it rules out some potential future designs as well.

[mwillsey]

To make that a big more concrete: we could add something like

fn enode_just_added(egraph: &mut EGraph<L, Self>, enode: &L, data: Self::Data);

or perhaps a hook that runs before insertion. Would that suffice? I'm trying to disentangle what you need from make, since there are other uses of make outside of e-node insertion (like this one in rebuilding).

[imbrem]

The issue is that I need to mutate the EGraph, or specifically the analysis, to generate the Data in the first place, since it's an index into a map which may not yet exist.

[imbrem]

I think actually some sample code would make this much clearer:

Let's say my language looks like this

"pair" = Pair([Id; 2]),
Integer(u32)

I'd like to tag each value with a type, which is going to be a big tree. Since there's going to be a lot of values with the same type, I'd like to avoid allocations and do a bit of hash-consing. So in my analysis, I have Data be a TypeId, and I have a TypeSet which has methods

pub fn pair(&mut self, left: TypeId, right: TypeId) -> TypeId;
pub fn integer(&mut self) -> TypeId

So I can't get the TypeId without having mutable access to the Analysis. Why this needs to be an Analysis in the first place is (A): there's no other clean way to tag every member of an enode with some data on creation, though I could do this with some kind of wrapper struct, but more importantly (B): on unioning things, I can detect type errors and, if I instead have an EGraph of types which may include type variables, do unification and type inference that way, which is what I want to experiment with.

[mwillsey]

Ok, I see your point. @oflatt @yihozhang any thoughts?

In the documentation, it would be nice to clarify that it is not make's responsibility to insert the e-node; that has already been done.

(ps: your use case sounds like it might be nicer in egglog; you could express this natively with a typeof function or something. caveat: egglog is not quite as polished of a tool as egg right now)

[yihozhang]

The example makes sense to me. I think to be consistent with other functions like merge, what we really want is mutable access to Analysis and immutable access to EGraph at the same time. But there's no trivial way of separating them.

[imbrem]

@yihozhang well, actually, mutable access to the E-graph as well allows us to be even more flexible so long as we don't trigger an infinite loop, as in this case we can do dependent-type-style things (e.g. generate Ids for enodes while making data for the analysis). This should be fine so long as the enodes passed in are of a different kind which will not lead to a recursive call, which shouldn't be too hard to guarantee.

[yihozhang]

Yeah, though I think we want E-class analysis to strike the right balance between being flexible and being a good, sound abstraction. In general, however, unrestrained access to E-graphs seems ad-hoc and dangerous (as you said, now you need to reason about how the internals of egg work). In the egg paper, we just require an e-class analysis to be a lattice and everything just works.

Here is another proposal to make's interface:
before:

    fn make(egraph: &EGraph<L, Self>, enode: &L) -> Self::Data;

after:

    fn make(&mut self, enode: &L, children: &[Data]) -> Self::Data;

This is "more faithful" to the idea that the e-class analysis of a parent is the composition of the e-class analyses of its children, and it seems this works for your use case. This seems to be a bigger change, though.

Either way, I'm fine with the current change (as well as my proposed change).

[imbrem]

@yihozhang I still favor the proposed change, since the requirement is not only relatively minor but also only relevant to advanced users. There could even be a second make method implementing my API, with a default implementation in terms of your simplified version; but again it is useful to be able to look at the EGraph in make so... eh, I say just provide the flexibility if it's possible without any real additional complexity right?

[mwillsey]

Ok, I think this change is worth the cost; it definitely opens up some more opportunities to do things in the e-graph that would be easier in egglog, but it worth supporting them here in egg I think.