Skip to main content

chomsky_uir/
analysis.rs

1use crate::constraint::{ConstraintSet, Effect};
2use crate::egraph::{Analysis, EGraph, Language};
3use crate::intent::IKun;
4
5#[derive(Default, Clone, Debug)]
6pub struct ConstraintAnalysis;
7
8impl Analysis<IKun> for ConstraintAnalysis {
9    type Data = ConstraintSet;
10
11    fn make(egraph: &EGraph<IKun, Self>, enode: &IKun) -> Self::Data {
12        let mut set = ConstraintSet::default();
13        match enode {
14            IKun::EffectConstraint(e) => set.effect = *e,
15            IKun::OwnershipConstraint(o) => set.ownership = Some(*o),
16            IKun::TypeConstraint(t) => set.r#type = Some(t.clone()),
17            IKun::AtomicConstraint => set.is_atomic = true,
18
19            // Propagation logic:
20            // WithConstraint(expr, constraint)
21            IKun::WithConstraint(expr, constraint) => {
22                let expr_data = egraph.get_class(*expr).data.clone();
23                let constraint_data = egraph.get_class(*constraint).data.clone();
24                set.merge(&expr_data);
25                set.merge(&constraint_data);
26            }
27
28            // Map(f, input) -> inherits constraints from f and input
29            IKun::Map(f, input) => {
30                let f_data = egraph.get_class(*f).data.clone();
31                let input_data = egraph.get_class(*input).data.clone();
32                set.merge(&f_data);
33                set.merge(&input_data);
34            }
35
36            // Seq(actions) -> union of all constraints
37            IKun::Seq(actions) => {
38                for &action in actions {
39                    let action_data = egraph.get_class(action).data.clone();
40                    set.merge(&action_data);
41                }
42            }
43
44            IKun::StateUpdate(target, val) => {
45                let target_data = egraph.get_class(*target).data.clone();
46                let val_data = egraph.get_class(*val).data.clone();
47                set.merge(&target_data);
48                set.merge(&val_data);
49                set.effect = set.effect.join(&Effect::WriteOnly);
50            }
51
52            IKun::ResourceClone(target) | IKun::ResourceDrop(target) => {
53                let target_data = egraph.get_class(*target).data.clone();
54                set.merge(&target_data);
55                // RC operations are side-effecting (modifying reference counts)
56                set.effect = set.effect.join(&Effect::ReadWrite);
57            }
58
59            IKun::CrossLangCall(call) => {
60                for &arg in &call.arguments {
61                    let arg_data = egraph.get_class(arg).data.clone();
62                    set.merge(&arg_data);
63                }
64                // Cross-language calls and intrinsics are assumed to be non-pure by default
65                set.effect = set.effect.join(&Effect::ReadWrite);
66            }
67
68            // Function application: inherits effect from the function and arguments
69            IKun::Apply(func, args) => {
70                let func_data = egraph.get_class(*func).data.clone();
71                set.merge(&func_data);
72                for &arg in args {
73                    let arg_data = egraph.get_class(arg).data.clone();
74                    set.merge(&arg_data);
75                }
76            }
77
78            // Lambda definition itself is pure (it's just a value)
79            IKun::Lambda(_, _) | IKun::Closure(_, _) => {
80                set.effect = Effect::Pure;
81            }
82
83            // Default: most atoms are Pure by default?
84            IKun::Constant(_)
85            | IKun::FloatConstant(_)
86            | IKun::BooleanConstant(_)
87            | IKun::StringConstant(_)
88            | IKun::Symbol(_) => {
89                set.effect = Effect::Pure;
90            }
91
92            _ => {
93                // For other nodes, we collect constraints from children
94                for &child in enode.children().iter() {
95                    let child_data = egraph.get_class(child).data.clone();
96                    set.merge(&child_data);
97                }
98            }
99        }
100        set
101    }
102
103    fn merge(&mut self, to: &mut Self::Data, from: Self::Data) -> bool {
104        to.merge(&from)
105    }
106
107    fn is_compatible(&self, data1: &Self::Data, data2: &Self::Data) -> bool {
108        data1.can_merge(data2)
109    }
110}