Skip to main content

slotted_egraphs/rewrite/
pattern.rs

1use crate::*;
2
3#[derive(Clone, Hash, PartialEq, Eq)]
4/// A Pattern to match against, or as the rhs of a rewrite rule.
5///
6/// - It supports pattern-variables `?x` to match against anything.
7/// - It supports (on the rhs) substitutions `b[x := t]` to substitute natively.
8pub enum Pattern<L: Language> {
9    ENode(L, Vec<Pattern<L>>),
10    PVar(String),                                             // ?x
11    Subst(Box<Pattern<L>>, Box<Pattern<L>>, Box<Pattern<L>>), // Subst(b, x, t) means `b[x := t]`
12}
13
14// We write this as pattern[subst] for short.
15pub fn pattern_subst<L: Language, N: Analysis<L>>(
16    eg: &mut EGraph<L, N>,
17    pattern: &Pattern<L>,
18    subst: &Subst,
19) -> AppliedId {
20    match &pattern {
21        Pattern::ENode(n, children) => {
22            let mut n = n.clone();
23            let mut refs: Vec<&mut _> = n.applied_id_occurrences_mut();
24            if CHECKS {
25                assert_eq!(children.len(), refs.len());
26            }
27            for i in 0..refs.len() {
28                *(refs[i]) = pattern_subst(eg, &children[i], subst);
29            }
30            eg.add_syn(n)
31        }
32        Pattern::PVar(v) => subst
33            .get(v)
34            .unwrap_or_else(|| {
35                panic!("encountered `?{v}` in pattern, but it is missing in the `subst`")
36            })
37            .clone(),
38        Pattern::Subst(b, x, t) => {
39            let b = pattern_subst(eg, &*b, subst);
40            let x = pattern_subst(eg, &*x, subst);
41            let t = pattern_subst(eg, &*t, subst);
42
43            // temporary swap-out so that we can access both the e-graph and the subst-method fully.
44            let mut method = eg.subst_method.take().unwrap();
45            let out = method.subst(b, x, t, eg);
46            eg.subst_method = Some(method);
47            out
48        }
49    }
50}
51
52// TODO maybe move into EGraph API?
53pub fn lookup_rec_expr<L: Language, N: Analysis<L>>(
54    re: &RecExpr<L>,
55    eg: &EGraph<L, N>,
56) -> Option<AppliedId> {
57    let mut n = re.node.clone();
58    let mut refs: Vec<&mut AppliedId> = n.applied_id_occurrences_mut();
59    if CHECKS {
60        assert_eq!(re.children.len(), refs.len());
61    }
62    for i in 0..refs.len() {
63        *(refs[i]) = lookup_rec_expr(&re.children[i], eg)?;
64    }
65    eg.lookup(&n)
66}
67
68pub fn pattern_to_re<L: Language>(pat: &Pattern<L>) -> RecExpr<L> {
69    let Pattern::ENode(n, children) = &pat else {
70        panic!()
71    };
72    let children: Vec<RecExpr<L>> = children.iter().map(|x| pattern_to_re(x)).collect();
73    RecExpr {
74        node: n.clone(),
75        children,
76    }
77}
78
79pub fn re_to_pattern<L: Language>(re: &RecExpr<L>) -> Pattern<L> {
80    let children: Vec<Pattern<L>> = re.children.iter().map(|x| re_to_pattern(x)).collect();
81    Pattern::ENode(re.node.clone(), children)
82}