logicaffeine_compile/optimize/egraph/
mod.rs

1//! The ARCHITECT — equality saturation over a compiler e-graph
2//! (EXODIA Phase 4, sprints 17–22).
3//!
4//! Classic egg-style design: a hash-consed term bank over the shared
5//! [`UnionFind`] (the same engine the kernel's congruence closure uses),
6//! a worklist `rebuild` restoring the congruence invariant after unions,
7//! budgeted saturation with a deterministic rule order, and bottom-up
8//! cost extraction.
9//!
10//! Class facts (scalar kind + integer interval) come from the Oracle and
11//! gate the conditional rewrites: `x / 2^n → x >> n` fires only with a
12//! non-negativity proof, the Group-2 boolean laws fire only on proven
13//! Bools, and EVERYTHING fails closed when no fact is present.
14
15pub mod convert;
16pub mod enode;
17pub mod extract;
18pub mod rules;
19
20pub use enode::CompilerENode;
21
22use std::collections::HashMap;
23
24use logicaffeine_base::union_find::UnionFind;
25
26use crate::optimize::ScalarKind;
27
28pub type NodeId = usize;
29
30/// Saturation stops after this many fixpoint iterations…
31pub const SATURATION_ITERS: usize = 8;
32/// …or when the term bank would exceed this many nodes.
33pub const MAX_NODES: usize = 10_000;
34
35/// Per-class knowledge, seeded from literals and the Oracle.
36///
37/// Facts describe the VALUE of every member of the class (if evaluation
38/// succeeds); merging classes therefore INTERSECTS intervals.
39#[derive(Debug, Clone, Copy, Default)]
40pub struct ClassFact {
41    pub scalar: Option<ScalarKind>,
42    pub range: Option<(i64, i64)>,
43    /// The value (if evaluation succeeds) is a collection — list, text,
44    /// map, or set. Kind facts are mutation-immune (contents change,
45    /// kinds never do), so they merge by OR.
46    pub collection: bool,
47    /// The value is specifically a LIST (implies `collection`).
48    pub list: bool,
49}
50
51pub struct CompilerEGraph {
52    uf: UnionFind,
53    /// Node id → the node as added (children canonical at insert time).
54    nodes: Vec<CompilerENode>,
55    /// Canonical node → its id. Stale keys (whose children later merged)
56    /// are either still-correct or unreachable by canonical lookups.
57    memo: HashMap<CompilerENode, NodeId>,
58    /// Class root → member node ids.
59    class_nodes: HashMap<NodeId, Vec<NodeId>>,
60    /// Child class root → parent node ids (for congruence repair).
61    uses: HashMap<NodeId, Vec<NodeId>>,
62    /// Class root → facts.
63    facts: HashMap<NodeId, ClassFact>,
64    /// Roots whose parents need congruence repair.
65    dirty: Vec<NodeId>,
66}
67
68impl CompilerEGraph {
69    pub fn new() -> Self {
70        CompilerEGraph {
71            uf: UnionFind::new(),
72            nodes: Vec::new(),
73            memo: HashMap::new(),
74            class_nodes: HashMap::new(),
75            uses: HashMap::new(),
76            facts: HashMap::new(),
77            dirty: Vec::new(),
78        }
79    }
80
81    pub fn node_count(&self) -> usize {
82        self.nodes.len()
83    }
84
85    pub fn find(&mut self, id: NodeId) -> NodeId {
86        self.uf.find(id)
87    }
88
89    /// The stored node with children mapped to their CURRENT class roots.
90    pub fn canonical_node(&mut self, id: NodeId) -> CompilerENode {
91        let node = self.nodes[id];
92        node.map_children(|c| self.uf.find(c))
93    }
94
95    /// Member node ids of `id`'s class, in insertion order (deterministic).
96    pub fn class_members(&mut self, id: NodeId) -> Vec<NodeId> {
97        let root = self.find(id);
98        self.class_nodes.get(&root).cloned().unwrap_or_default()
99    }
100
101    /// Hash-consing insert: an existing congruent node is returned instead
102    /// of allocating. Literal leaves seed their own class facts.
103    pub fn add(&mut self, node: CompilerENode) -> NodeId {
104        let canonical = node.map_children(|c| self.uf.find(c));
105        if let Some(&existing) = self.memo.get(&canonical) {
106            return self.uf.find(existing);
107        }
108        let id = self.uf.make_set();
109        debug_assert_eq!(id, self.nodes.len());
110        self.nodes.push(canonical);
111        self.memo.insert(canonical, id);
112        self.class_nodes.insert(id, vec![id]);
113        for child in canonical.children() {
114            self.uses.entry(child).or_default().push(id);
115        }
116        match canonical {
117            CompilerENode::Int(k) => {
118                self.facts.insert(
119                    id,
120                    ClassFact {
121                        scalar: Some(ScalarKind::Int),
122                        range: Some((k, k)),
123                        ..ClassFact::default()
124                    },
125                );
126            }
127            CompilerENode::Bool(_) => {
128                self.facts.insert(
129                    id,
130                    ClassFact { scalar: Some(ScalarKind::Bool), ..ClassFact::default() },
131                );
132            }
133            CompilerENode::Float(_) => {
134                self.facts.insert(
135                    id,
136                    ClassFact { scalar: Some(ScalarKind::Float), ..ClassFact::default() },
137                );
138            }
139            // A slice that evaluates successfully ALWAYS yields a list.
140            CompilerENode::Slice(..) => {
141                self.facts.insert(
142                    id,
143                    ClassFact { collection: true, list: true, ..ClassFact::default() },
144                );
145            }
146            // A copy that evaluates preserves its operand's kind; the
147            // converter upgrades this when the operand is a proven list.
148            _ => {}
149        }
150        id
151    }
152
153    /// Merge two classes. Facts intersect (both described the same value).
154    /// Returns true if a merge actually happened.
155    pub fn union(&mut self, a: NodeId, b: NodeId) -> bool {
156        let ra = self.uf.find(a);
157        let rb = self.uf.find(b);
158        if ra == rb {
159            return false;
160        }
161        let fa = self.facts.remove(&ra);
162        let fb = self.facts.remove(&rb);
163        self.uf.union(ra, rb);
164        let root = self.uf.find(ra);
165        let loser = if root == ra { rb } else { ra };
166
167        let mut members = self.class_nodes.remove(&loser).unwrap_or_default();
168        self.class_nodes.entry(root).or_default().append(&mut members);
169        let mut moved_uses = self.uses.remove(&loser).unwrap_or_default();
170        self.uses.entry(root).or_default().append(&mut moved_uses);
171
172        let merged = match (fa, fb) {
173            (Some(x), Some(y)) => ClassFact {
174                scalar: x.scalar.or(y.scalar),
175                range: match (x.range, y.range) {
176                    (Some((al, ah)), Some((bl, bh))) => {
177                        let lo = al.max(bl);
178                        let hi = ah.min(bh);
179                        debug_assert!(lo <= hi, "contradictory ranges merged — unsound rule?");
180                        if lo <= hi { Some((lo, hi)) } else { Some((al, ah)) }
181                    }
182                    (r, None) | (None, r) => r,
183                },
184                // Kind facts are sound statements about the shared value:
185                // either side's proof carries over.
186                collection: x.collection || y.collection,
187                list: x.list || y.list,
188            },
189            (Some(x), None) | (None, Some(x)) => x,
190            (None, None) => ClassFact::default(),
191        };
192        self.facts.insert(root, merged);
193        self.dirty.push(root);
194        true
195    }
196
197    /// Restore the congruence invariant: parents of merged classes are
198    /// re-canonicalized; newly congruent pairs are unioned (worklist).
199    pub fn rebuild(&mut self) {
200        while let Some(r) = self.dirty.pop() {
201            let root = self.uf.find(r);
202            let parents = self.uses.remove(&root).unwrap_or_default();
203            let mut kept: Vec<NodeId> = Vec::with_capacity(parents.len());
204            for p in parents {
205                let canon = self.nodes[p].map_children(|c| self.uf.find(c));
206                if let Some(&q) = self.memo.get(&canon) {
207                    if self.uf.find(q) != self.uf.find(p) {
208                        self.union(p, q);
209                    }
210                }
211                let rep = self.uf.find(p);
212                self.memo.insert(canon, rep);
213                kept.push(p);
214            }
215            let now = self.uf.find(root);
216            self.uses.entry(now).or_default().extend(kept);
217        }
218    }
219
220    // ----- facts ---------------------------------------------------------
221
222    pub fn set_scalar(&mut self, id: NodeId, kind: ScalarKind) {
223        let root = self.find(id);
224        self.facts.entry(root).or_default().scalar = Some(kind);
225    }
226
227    pub fn set_int_range(&mut self, id: NodeId, lo: i64, hi: i64) {
228        let root = self.find(id);
229        let fact = self.facts.entry(root).or_default();
230        fact.scalar = Some(ScalarKind::Int);
231        fact.range = Some((lo, hi));
232    }
233
234    pub fn scalar_of(&mut self, id: NodeId) -> Option<ScalarKind> {
235        let root = self.find(id);
236        self.facts.get(&root).and_then(|f| f.scalar)
237    }
238
239    pub fn int_range(&mut self, id: NodeId) -> Option<(i64, i64)> {
240        let root = self.find(id);
241        self.facts.get(&root).and_then(|f| f.range)
242    }
243
244    /// A PROVEN single value (point interval) — sound grounds for folding.
245    pub fn int_value(&mut self, id: NodeId) -> Option<i64> {
246        match self.int_range(id) {
247            Some((lo, hi)) if lo == hi => Some(lo),
248            _ => None,
249        }
250    }
251
252    pub fn proven_nonneg(&mut self, id: NodeId) -> bool {
253        matches!(self.int_range(id), Some((lo, _)) if lo >= 0)
254    }
255
256    /// Mark the class a proven collection (list/text/map/set).
257    pub fn set_collection(&mut self, id: NodeId) {
258        let root = self.find(id);
259        self.facts.entry(root).or_default().collection = true;
260    }
261
262    /// Mark the class a proven LIST (implies collection).
263    pub fn set_list(&mut self, id: NodeId) {
264        let root = self.find(id);
265        let fact = self.facts.entry(root).or_default();
266        fact.collection = true;
267        fact.list = true;
268    }
269
270    pub fn proven_collection(&mut self, id: NodeId) -> bool {
271        let root = self.find(id);
272        self.facts.get(&root).is_some_and(|f| f.collection)
273    }
274
275    pub fn proven_list(&mut self, id: NodeId) -> bool {
276        let root = self.find(id);
277        self.facts.get(&root).is_some_and(|f| f.list)
278    }
279
280    /// Does the class contain the literal `Bool(b)`?
281    pub fn class_has_bool(&mut self, id: NodeId, b: bool) -> bool {
282        self.class_members(id)
283            .into_iter()
284            .any(|m| self.nodes[m] == CompilerENode::Bool(b))
285    }
286
287    /// Whole-tree totality: the class has at least one member whose op is
288    /// total and whose children are recursively total. A rewrite may only
289    /// DELETE (stop evaluating) a provably total subterm — anything else
290    /// could erase a runtime error or an effect. Bottom-up fixpoint,
291    /// cycle-tolerant (cycles stay "not yet proven" until a leaf grounds
292    /// them).
293    ///
294    /// Totality is FACT-CONDITIONAL where the op's only failure mode is a
295    /// kind error: `Len`/`Copy`/`Contains` never raise over a PROVEN
296    /// collection, and raise a type error over anything unproven — so
297    /// without the fact they stay non-total, exactly like fold.rs's
298    /// fail-closed `expr_is_total`.
299    pub fn provably_total(&mut self, id: NodeId) -> bool {
300        let mut total: HashMap<NodeId, bool> = HashMap::new();
301        loop {
302            let mut changed = false;
303            for n in 0..self.nodes.len() {
304                let node = self.nodes[n];
305                let root = self.uf.find(n);
306                if *total.get(&root).unwrap_or(&false) {
307                    continue;
308                }
309                let op_total = match node {
310                    CompilerENode::Len(c)
311                    | CompilerENode::Copy(c)
312                    | CompilerENode::Contains(c, _) => self.proven_collection(c),
313                    _ => node.op_is_total(),
314                };
315                if op_total
316                    && node
317                        .children()
318                        .iter()
319                        .all(|&c| *total.get(&self.uf.find(c)).unwrap_or(&false))
320                {
321                    total.insert(root, true);
322                    changed = true;
323                }
324            }
325            if !changed {
326                break;
327            }
328        }
329        let root = self.find(id);
330        *total.get(&root).unwrap_or(&false)
331    }
332
333    // ----- saturation ----------------------------------------------------
334
335    /// Budgeted equality saturation: deterministic node order × rule order,
336    /// unions applied immediately, congruence repaired per iteration. Stops
337    /// at fixpoint, iteration cap, or the node-bank cap.
338    pub fn saturate(&mut self, rules: &[rules::Rewrite]) {
339        self.rebuild();
340        'outer: for _ in 0..SATURATION_ITERS {
341            let snapshot = self.nodes.len();
342            let mut changed = false;
343            for id in 0..snapshot {
344                for rule in rules {
345                    if self.nodes.len() + 4 > MAX_NODES {
346                        self.rebuild();
347                        break 'outer;
348                    }
349                    if let Some(replacement) = (rule.apply)(self, id) {
350                        changed |= self.union(id, replacement);
351                    }
352                }
353            }
354            self.rebuild();
355            if !changed && self.nodes.len() == snapshot {
356                break;
357            }
358        }
359    }
360}
361
362impl Default for CompilerEGraph {
363    fn default() -> Self {
364        Self::new()
365    }
366}
logicaffeine_compile/optimize/egraph/mod.rs

logicaffeine_compile/optimize/egraph/
mod.rs
