Propagate data downwards through the egraph

[Kixiron]

How would I go about propagating data "down" through the egraph, giving data from parents to their children?
As an example, I know the types of function parameters, but I need to get the type of each reference to the parameter, e.g.

(fun (x Int) x)

I know that the function has the parameter x which is an Int, but I need to propagate the variables in scope and their associated types downwards to the function body so that when looking at the x in the function body I know its correct type

[mwillsey]

That's a really tricky one, since (x Int) could appear in multiple places in different terms. I think you've hit on one of the hardest things to fit into the e-graph paradigm right now: context. You can't really pass things down because you don't "own" your children; someone else might have the same child and pass down something else.

[Kixiron]

What would you recommend as a "workaround" or alternative approach to the same goal? Encode types as part of the egraph?

[mwillsey]

Somehow do it 'outside the egraph'. If you are very careful to uniqify your variables, you can probably just use a custom pass over the e-graph to calculate the type of variables. You can then just manually add that to the analysis data of those terms if you wish.

[Kixiron]

I came up with this which I've got running as a hook on the runner, any tweaks you'd make to it? I think I may run into bad situations here where things get unified incorrectly but I'm not sure how to prevent it
(Note that the HashMap and HashSet are from im so while it looks like I'm cloning like there's no tomorrow I'm not actually doing so)

pub fn variable_propagation(graph: &mut EGraph, roots: &[Id]) {
    let variables = HashMap::new();
    let mut to_union = Vec::new();
    let visited = HashSet::new();

    for &root in roots {
        let class = &graph[root];

        for node in class.nodes.iter() {
            propagate(
                graph,
                variables.clone(),
                class.id,
                node,
                &mut to_union,
                visited.clone(),
            );
        }
    }

    for (id, new_id) in to_union {
        graph.union(id, new_id);
    }
}

fn propagate(
    graph: &EGraph,
    mut variables: HashMap<Id, Id>,
    node_id: Id,
    node: &Operator,
    to_union: &mut Vec<(Id, Id)>,
    visited: HashSet<Id>,
) {
    if let Operator::Let([var, val, _]) = *node {
        variables = variables.update(var, val);
    } else if let Operator::Var(var) = node {
        if let Some(&value) = variables.get(var) {
            to_union.push((node_id, value));
        }
    }

    for &child in node.children() {
        if !visited.contains(&child) {
            let visited = visited.update(child);

            for node in graph[child].nodes.iter() {
                propagate(
                    graph,
                    variables.clone(),
                    child,
                    node,
                    to_union,
                    visited.clone(),
                );
            }
        }
    }
}

[Bastacyclop]

What I have done myself to support types is to embed them into e-classes, types are encoded in the e-graph and I consider that (x: int) != (x: float), so the two are represented by different e-classes. The advantage is that you get very precise type information and rewrite rules can easily reason about types, the disadvantage is that you get less sharing. I have already encountered cases were rewriting the untyped program was much easier than rewriting the typed program due to this loss of sharing.

I would be delighted if there were other alternatives. I also have issues with bottom-up information, not just top-down information (actually type information may flow in both ways). For example, sometimes approximating the set of free variables with a union-based e-class analysis prevents valuable rewrites from being found.

I wonder if there could be hybrid approaches here to allow reasoning at different information levels, for example:

• only for the xs that have type int and are different from 0, I know that x / x = 1
• but for any x, whatever its type or value, I know that (app (fun x x) y) = y

Although I have no idea on how to make this work reasonably efficiently in practice.