logicaffeine_kernel/
cc.rs

1//! Congruence Closure Tactic
2//!
3//! Proves equalities over uninterpreted functions using Union-Find with
4//! congruence propagation. This implements a simplified Downey-Sethi-Tarjan
5//! algorithm.
6//!
7//! # Algorithm
8//!
9//! The congruence closure tactic works in four steps:
10//! 1. **Build E-graph**: Add all subterms from goal and hypotheses
11//! 2. **Merge hypotheses**: For each hypothesis `x = y`, merge x and y
12//! 3. **Propagate**: When `x = y` and `f(x)`, `f(y)` exist, merge `f(x)` and `f(y)`
13//! 4. **Check**: Goal `a = b` holds iff a and b are in the same equivalence class
14//!
15//! # Supported Goals
16//!
17//! - Direct equalities: `Eq a a` (reflexivity)
18//! - Implications: `(Eq x y) -> (Eq (f x) (f y))` (congruence)
19//! - Nested implications with multiple hypotheses
20//!
21//! # E-Graph Structure
22//!
23//! The E-graph maintains:
24//! - A Union-Find for equivalence classes
25//! - Hash-consing for structural sharing
26//! - Use lists for efficient congruence propagation
27
28use std::collections::HashMap;
29
30use crate::term::{Literal, Term};
31
32type TermId = usize;
33
34/// The shared Union-Find — one equivalence engine underneath the kernel's
35/// congruence closure and the compiler's e-graph.
36pub use logicaffeine_base::union_find::UnionFind;
37
38// =============================================================================
39// E-GRAPH DATA STRUCTURE
40// =============================================================================
41
42/// Node in the E-graph representing a term.
43///
44/// Terms are represented in a curried style where function application
45/// is a binary node. For example, `f(x, y)` is `App(App(f, x), y)`.
46#[derive(Debug, Clone, PartialEq, Eq, Hash)]
47pub enum ENode {
48    /// Integer literal from `SLit`.
49    Lit(i64),
50    /// De Bruijn variable from `SVar`.
51    Var(i64),
52    /// Named constant or function symbol from `SName`.
53    Name(String),
54    /// Function application (curried): `func` applied to `arg`.
55    App {
56        /// The function being applied.
57        func: TermId,
58        /// The argument.
59        arg: TermId,
60    },
61}
62
63/// E-graph with congruence closure.
64///
65/// Combines a Union-Find for equivalence classes with hash-consing for
66/// structural sharing and use lists for efficient congruence propagation.
67pub struct EGraph {
68    /// The nodes stored in the graph.
69    nodes: Vec<ENode>,
70    /// Union-Find tracking equivalence classes.
71    uf: UnionFind,
72    /// Hash-consing map: node content to its canonical ID.
73    node_map: HashMap<ENode, TermId>,
74    /// Pending merges to propagate (worklist algorithm).
75    pending: Vec<(TermId, TermId)>,
76    /// Use lists: for each term, the App nodes that use it as func or arg.
77    /// Used to find potential congruences when terms are merged.
78    use_list: Vec<Vec<TermId>>,
79}
80
81impl EGraph {
82    pub fn new() -> Self {
83        EGraph {
84            nodes: Vec::new(),
85            uf: UnionFind::new(),
86            node_map: HashMap::new(),
87            pending: Vec::new(),
88            use_list: Vec::new(),
89        }
90    }
91
92    /// Add a node, return its ID (hash-consed)
93    pub fn add(&mut self, node: ENode) -> TermId {
94        // Hash-consing: return existing ID if node already exists
95        if let Some(&id) = self.node_map.get(&node) {
96            return id;
97        }
98
99        let id = self.nodes.len();
100        self.nodes.push(node.clone());
101        self.node_map.insert(node.clone(), id);
102        self.uf.make_set();
103        self.use_list.push(Vec::new());
104
105        // Register in use lists for congruence detection
106        if let ENode::App { func, arg } = &node {
107            self.use_list[*func].push(id);
108            self.use_list[*arg].push(id);
109        }
110
111        id
112    }
113
114    /// Merge two terms and propagate congruences
115    pub fn merge(&mut self, a: TermId, b: TermId) {
116        self.pending.push((a, b));
117        self.propagate();
118    }
119
120    /// Propagate congruences until fixed point
121    fn propagate(&mut self) {
122        while let Some((a, b)) = self.pending.pop() {
123            let ra = self.uf.find(a);
124            let rb = self.uf.find(b);
125            if ra == rb {
126                continue;
127            }
128
129            // Before merging, collect uses for congruence checking
130            let uses_a: Vec<TermId> = self.use_list[ra].clone();
131            let uses_b: Vec<TermId> = self.use_list[rb].clone();
132
133            // Merge equivalence classes
134            self.uf.union(ra, rb);
135            let new_root = self.uf.find(ra);
136
137            // Check for new congruences first (before modifying use lists)
138            // If f(a) and f(b) exist, and a=b now, then f(a)=f(b)
139            for &ua in &uses_a {
140                for &ub in &uses_b {
141                    if self.congruent(ua, ub) {
142                        self.pending.push((ua, ub));
143                    }
144                }
145            }
146
147            // Merge use lists (now safe to consume uses_a/uses_b)
148            if new_root == ra {
149                for u in uses_b {
150                    self.use_list[ra].push(u);
151                }
152            } else {
153                for u in uses_a {
154                    self.use_list[rb].push(u);
155                }
156            }
157        }
158    }
159
160    /// Check if two application nodes are congruent
161    fn congruent(&mut self, a: TermId, b: TermId) -> bool {
162        match (&self.nodes[a].clone(), &self.nodes[b].clone()) {
163            (ENode::App { func: f1, arg: a1 }, ENode::App { func: f2, arg: a2 }) => {
164                self.uf.find(*f1) == self.uf.find(*f2) && self.uf.find(*a1) == self.uf.find(*a2)
165            }
166            _ => false,
167        }
168    }
169
170    /// Check if two terms are in the same equivalence class
171    pub fn equivalent(&mut self, a: TermId, b: TermId) -> bool {
172        self.uf.find(a) == self.uf.find(b)
173    }
174}
175
176// =============================================================================
177// SYNTAX TERM REIFICATION
178// =============================================================================
179
180/// Reify a Syntax term into an E-graph node.
181///
182/// Converts the deep embedding (Syntax) into E-graph nodes, returning
183/// the ID of the root node. Subterms are recursively reified and
184/// hash-consed (duplicate structures share the same ID).
185///
186/// # Returns
187///
188/// `Some(id)` on successful reification, `None` if the term cannot be reified.
189pub fn reify(egraph: &mut EGraph, term: &Term) -> Option<TermId> {
190    // SLit n -> Lit(n)
191    if let Some(n) = extract_slit(term) {
192        return Some(egraph.add(ENode::Lit(n)));
193    }
194
195    // SVar i -> Var(i)
196    if let Some(i) = extract_svar(term) {
197        return Some(egraph.add(ENode::Var(i)));
198    }
199
200    // SName s -> Name(s)
201    if let Some(name) = extract_sname(term) {
202        return Some(egraph.add(ENode::Name(name)));
203    }
204
205    // SApp f a -> App { func, arg }
206    if let Some((func_term, arg_term)) = extract_sapp(term) {
207        let func = reify(egraph, &func_term)?;
208        let arg = reify(egraph, &arg_term)?;
209        return Some(egraph.add(ENode::App { func, arg }));
210    }
211
212    None
213}
214
215// =============================================================================
216// GOAL DECOMPOSITION
217// =============================================================================
218
219/// Decompose a goal into hypotheses and conclusion.
220///
221/// Peels off nested implications to extract equality hypotheses.
222/// For example, `(h1 -> h2 -> conclusion)` becomes `([h1, h2], conclusion)`.
223///
224/// # Returns
225///
226/// A tuple of:
227/// - Vector of equality hypothesis pairs (LHS, RHS)
228/// - The final conclusion term
229///
230/// Only equalities in hypothesis position are extracted; other hypotheses
231/// are ignored.
232pub fn decompose_goal(goal: &Term) -> (Vec<(Term, Term)>, Term) {
233    let mut hypotheses = Vec::new();
234    let mut current = goal.clone();
235
236    // Peel off nested implications
237    while let Some((hyp, rest)) = extract_implication(&current) {
238        if let Some((lhs, rhs)) = extract_equality(&hyp) {
239            hypotheses.push((lhs, rhs));
240        }
241        current = rest;
242    }
243
244    (hypotheses, current)
245}
246
247/// Check if a goal is provable by congruence closure.
248///
249/// This is the main entry point for the CC tactic. It builds an E-graph,
250/// adds hypothesis equalities, propagates congruences, and checks if the
251/// conclusion follows.
252///
253/// # Supported Goals
254///
255/// - Direct equalities: `(Eq a a)` succeeds by reflexivity
256/// - Implications: `(implies (Eq x y) (Eq (f x) (f y)))` succeeds by congruence
257/// - Multiple hypotheses: `(implies (Eq a b) (implies (Eq b c) (Eq a c)))`
258///
259/// # Returns
260///
261/// `true` if the goal is provable by congruence closure, `false` otherwise.
262pub fn check_goal(goal: &Term) -> bool {
263    let (hypotheses, conclusion) = decompose_goal(goal);
264
265    // Conclusion must be an equality
266    let (lhs, rhs) = match extract_equality(&conclusion) {
267        Some(eq) => eq,
268        None => return false,
269    };
270
271    let mut egraph = EGraph::new();
272
273    // IMPORTANT: Reify conclusion FIRST so that fx and fy exist in the graph
274    // with their use lists populated. This way when we merge x=y, congruence
275    // will propagate to fx=fy.
276    let lhs_id = match reify(&mut egraph, &lhs) {
277        Some(id) => id,
278        None => return false,
279    };
280
281    let rhs_id = match reify(&mut egraph, &rhs) {
282        Some(id) => id,
283        None => return false,
284    };
285
286    // Now reify and merge hypothesis equalities
287    // The subterms (x, y) will be hash-consed with the ones in fx, fy
288    for (h_lhs, h_rhs) in &hypotheses {
289        let h_lhs_id = match reify(&mut egraph, h_lhs) {
290            Some(id) => id,
291            None => return false,
292        };
293        let h_rhs_id = match reify(&mut egraph, h_rhs) {
294            Some(id) => id,
295            None => return false,
296        };
297        egraph.merge(h_lhs_id, h_rhs_id);
298    }
299
300    // Check if conclusion follows by congruence
301    egraph.equivalent(lhs_id, rhs_id)
302}
303
304// =============================================================================
305// HELPER EXTRACTORS
306// =============================================================================
307
308/// Extract integer from SLit n
309fn extract_slit(term: &Term) -> Option<i64> {
310    if let Term::App(ctor, arg) = term {
311        if let Term::Global(name) = ctor.as_ref() {
312            if name == "SLit" {
313                if let Term::Lit(Literal::Int(n)) = arg.as_ref() {
314                    return Some(*n);
315                }
316            }
317        }
318    }
319    None
320}
321
322/// Extract variable index from SVar i
323fn extract_svar(term: &Term) -> Option<i64> {
324    if let Term::App(ctor, arg) = term {
325        if let Term::Global(name) = ctor.as_ref() {
326            if name == "SVar" {
327                if let Term::Lit(Literal::Int(i)) = arg.as_ref() {
328                    return Some(*i);
329                }
330            }
331        }
332    }
333    None
334}
335
336/// Extract name from SName "x"
337fn extract_sname(term: &Term) -> Option<String> {
338    if let Term::App(ctor, arg) = term {
339        if let Term::Global(name) = ctor.as_ref() {
340            if name == "SName" {
341                if let Term::Lit(Literal::Text(s)) = arg.as_ref() {
342                    return Some(s.clone());
343                }
344            }
345        }
346    }
347    None
348}
349
350/// Extract unary application: SApp f a
351fn extract_sapp(term: &Term) -> Option<(Term, Term)> {
352    if let Term::App(outer, arg) = term {
353        if let Term::App(sapp, func) = outer.as_ref() {
354            if let Term::Global(ctor) = sapp.as_ref() {
355                if ctor == "SApp" {
356                    return Some((func.as_ref().clone(), arg.as_ref().clone()));
357                }
358            }
359        }
360    }
361    None
362}
363
364/// Extract implication: SApp (SApp (SName "implies") hyp) concl
365fn extract_implication(term: &Term) -> Option<(Term, Term)> {
366    if let Some((op, hyp, concl)) = extract_binary_app(term) {
367        if op == "implies" {
368            return Some((hyp, concl));
369        }
370    }
371    None
372}
373
374/// Extract equality: SApp (SApp (SName "Eq") lhs) rhs
375fn extract_equality(term: &Term) -> Option<(Term, Term)> {
376    if let Some((op, lhs, rhs)) = extract_binary_app(term) {
377        if op == "Eq" || op == "eq" {
378            return Some((lhs, rhs));
379        }
380    }
381    None
382}
383
384/// Extract binary application: SApp (SApp (SName "op") a) b
385fn extract_binary_app(term: &Term) -> Option<(String, Term, Term)> {
386    if let Term::App(outer, b) = term {
387        if let Term::App(sapp_outer, inner) = outer.as_ref() {
388            if let Term::Global(ctor) = sapp_outer.as_ref() {
389                if ctor == "SApp" {
390                    if let Term::App(partial, a) = inner.as_ref() {
391                        if let Term::App(sapp_inner, op_term) = partial.as_ref() {
392                            if let Term::Global(ctor2) = sapp_inner.as_ref() {
393                                if ctor2 == "SApp" {
394                                    if let Some(op) = extract_sname(op_term) {
395                                        return Some((
396                                            op,
397                                            a.as_ref().clone(),
398                                            b.as_ref().clone(),
399                                        ));
400                                    }
401                                }
402                            }
403                        }
404                    }
405                }
406            }
407        }
408    }
409    None
410}
411
412// =============================================================================
413// UNIT TESTS
414// =============================================================================
415
416#[cfg(test)]
417mod tests {
418    use super::*;
419
420    #[test]
421    fn test_union_find_basic() {
422        let mut uf = UnionFind::new();
423        let a = uf.make_set();
424        let b = uf.make_set();
425        assert_ne!(uf.find(a), uf.find(b));
426        uf.union(a, b);
427        assert_eq!(uf.find(a), uf.find(b));
428    }
429
430    #[test]
431    fn test_union_find_transitivity() {
432        let mut uf = UnionFind::new();
433        let a = uf.make_set();
434        let b = uf.make_set();
435        let c = uf.make_set();
436        uf.union(a, b);
437        uf.union(b, c);
438        assert_eq!(uf.find(a), uf.find(c));
439    }
440
441    #[test]
442    fn test_egraph_reflexive() {
443        let mut eg = EGraph::new();
444        let x = eg.add(ENode::Var(0));
445        assert!(eg.equivalent(x, x));
446    }
447
448    #[test]
449    fn test_egraph_congruence() {
450        let mut eg = EGraph::new();
451        let x = eg.add(ENode::Var(0));
452        let y = eg.add(ENode::Var(1));
453        let f = eg.add(ENode::Name("f".to_string()));
454        let fx = eg.add(ENode::App { func: f, arg: x });
455        let fy = eg.add(ENode::App { func: f, arg: y });
456
457        // Before merging x=y, f(x) != f(y)
458        assert!(!eg.equivalent(fx, fy));
459
460        // After merging x=y, congruence gives f(x) = f(y)
461        eg.merge(x, y);
462        assert!(eg.equivalent(fx, fy));
463    }
464
465    #[test]
466    fn test_egraph_nested_congruence() {
467        let mut eg = EGraph::new();
468        let a = eg.add(ENode::Var(0));
469        let b = eg.add(ENode::Var(1));
470        let c = eg.add(ENode::Var(2));
471        let f = eg.add(ENode::Name("f".to_string()));
472
473        let fa = eg.add(ENode::App { func: f, arg: a });
474        let fc = eg.add(ENode::App { func: f, arg: c });
475        let ffa = eg.add(ENode::App { func: f, arg: fa });
476        let ffc = eg.add(ENode::App { func: f, arg: fc });
477
478        // a = b, b = c should give f(f(a)) = f(f(c))
479        eg.merge(a, b);
480        eg.merge(b, c);
481        assert!(eg.equivalent(ffa, ffc));
482    }
483
484    #[test]
485    fn test_egraph_binary_congruence() {
486        let mut eg = EGraph::new();
487        let a = eg.add(ENode::Var(0));
488        let b = eg.add(ENode::Var(1));
489        let c = eg.add(ENode::Var(2));
490        let add = eg.add(ENode::Name("add".to_string()));
491
492        // add(a, c) and add(b, c) as curried: add a c = (add a) c
493        let add_a = eg.add(ENode::App { func: add, arg: a });
494        let add_b = eg.add(ENode::App { func: add, arg: b });
495        let add_a_c = eg.add(ENode::App { func: add_a, arg: c });
496        let add_b_c = eg.add(ENode::App { func: add_b, arg: c });
497
498        assert!(!eg.equivalent(add_a_c, add_b_c));
499        eg.merge(a, b);
500        assert!(eg.equivalent(add_a_c, add_b_c));
501    }
502
503    // =========================================================================
504    // EXTRACTION TESTS
505    // =========================================================================
506
507    /// Helper to build SName "s"
508    fn make_sname(s: &str) -> Term {
509        Term::App(
510            Box::new(Term::Global("SName".to_string())),
511            Box::new(Term::Lit(Literal::Text(s.to_string()))),
512        )
513    }
514
515    /// Helper to build SVar i
516    fn make_svar(i: i64) -> Term {
517        Term::App(
518            Box::new(Term::Global("SVar".to_string())),
519            Box::new(Term::Lit(Literal::Int(i))),
520        )
521    }
522
523    /// Helper to build SApp f a
524    fn make_sapp(f: Term, a: Term) -> Term {
525        Term::App(
526            Box::new(Term::App(
527                Box::new(Term::Global("SApp".to_string())),
528                Box::new(f),
529            )),
530            Box::new(a),
531        )
532    }
533
534    #[test]
535    fn test_extract_sname() {
536        let term = make_sname("f");
537        assert_eq!(extract_sname(&term), Some("f".to_string()));
538    }
539
540    #[test]
541    fn test_extract_svar() {
542        let term = make_svar(0);
543        assert_eq!(extract_svar(&term), Some(0));
544    }
545
546    #[test]
547    fn test_extract_sapp() {
548        // SApp (SName "f") (SVar 0)
549        let term = make_sapp(make_sname("f"), make_svar(0));
550        let result = extract_sapp(&term);
551        assert!(result.is_some());
552        let (func, arg) = result.unwrap();
553        assert_eq!(extract_sname(&func), Some("f".to_string()));
554        assert_eq!(extract_svar(&arg), Some(0));
555    }
556
557    #[test]
558    fn test_extract_binary_app() {
559        // SApp (SApp (SName "Eq") (SVar 0)) (SVar 1)
560        let term = make_sapp(make_sapp(make_sname("Eq"), make_svar(0)), make_svar(1));
561        let result = extract_binary_app(&term);
562        assert!(result.is_some(), "Should extract binary app");
563        let (op, a, b) = result.unwrap();
564        assert_eq!(op, "Eq");
565        assert_eq!(extract_svar(&a), Some(0));
566        assert_eq!(extract_svar(&b), Some(1));
567    }
568
569    #[test]
570    fn test_extract_equality() {
571        // SApp (SApp (SName "Eq") (SVar 0)) (SVar 1)
572        let term = make_sapp(make_sapp(make_sname("Eq"), make_svar(0)), make_svar(1));
573        let result = extract_equality(&term);
574        assert!(result.is_some(), "Should extract equality");
575        let (lhs, rhs) = result.unwrap();
576        assert_eq!(extract_svar(&lhs), Some(0));
577        assert_eq!(extract_svar(&rhs), Some(1));
578    }
579
580    #[test]
581    fn test_extract_implication() {
582        // Build: SApp (SApp (SName "implies") hyp) concl
583        // hyp = SApp (SApp (SName "Eq") x) y
584        // concl = SApp (SApp (SName "Eq") fx) fy
585        let x = make_svar(0);
586        let y = make_svar(1);
587        let hyp = make_sapp(make_sapp(make_sname("Eq"), x.clone()), y.clone());
588
589        let f = make_sname("f");
590        let fx = make_sapp(f.clone(), x);
591        let fy = make_sapp(f, y);
592        let concl = make_sapp(make_sapp(make_sname("Eq"), fx), fy);
593
594        let goal = make_sapp(make_sapp(make_sname("implies"), hyp.clone()), concl.clone());
595
596        let result = extract_implication(&goal);
597        assert!(result.is_some(), "Should extract implication");
598        let (hyp_extracted, concl_extracted) = result.unwrap();
599
600        // Verify hypothesis is the equality x = y
601        let hyp_eq = extract_equality(&hyp_extracted);
602        assert!(hyp_eq.is_some(), "Hypothesis should be equality");
603        let (h_lhs, h_rhs) = hyp_eq.unwrap();
604        assert_eq!(extract_svar(&h_lhs), Some(0));
605        assert_eq!(extract_svar(&h_rhs), Some(1));
606
607        // Verify conclusion is an equality
608        let concl_eq = extract_equality(&concl_extracted);
609        assert!(concl_eq.is_some(), "Conclusion should be equality");
610    }
611
612    #[test]
613    fn test_decompose_goal_with_hypothesis() {
614        // Build: (implies (Eq x y) (Eq (f x) (f y)))
615        let x = make_svar(0);
616        let y = make_svar(1);
617        let hyp = make_sapp(make_sapp(make_sname("Eq"), x.clone()), y.clone());
618
619        let f = make_sname("f");
620        let fx = make_sapp(f.clone(), x);
621        let fy = make_sapp(f, y);
622        let concl = make_sapp(make_sapp(make_sname("Eq"), fx), fy);
623
624        let goal = make_sapp(make_sapp(make_sname("implies"), hyp), concl);
625
626        let (hypotheses, conclusion) = decompose_goal(&goal);
627        assert_eq!(hypotheses.len(), 1, "Should have 1 hypothesis");
628
629        // Verify conclusion is an equality
630        let concl_eq = extract_equality(&conclusion);
631        assert!(concl_eq.is_some(), "Conclusion should be equality");
632    }
633
634    #[test]
635    fn test_check_goal_with_hypothesis() {
636        // Build: (implies (Eq x y) (Eq (f x) (f y)))
637        // This should be provable by CC
638        let x = make_svar(0);
639        let y = make_svar(1);
640        let hyp = make_sapp(make_sapp(make_sname("Eq"), x.clone()), y.clone());
641
642        let f = make_sname("f");
643        let fx = make_sapp(f.clone(), x.clone());
644        let fy = make_sapp(f.clone(), y.clone());
645        let concl = make_sapp(make_sapp(make_sname("Eq"), fx), fy);
646
647        let goal = make_sapp(make_sapp(make_sname("implies"), hyp), concl);
648
649        assert!(check_goal(&goal), "CC should prove x=y → f(x)=f(y)");
650    }
651}
logicaffeine_kernel/cc.rs

logicaffeine_kernel/
cc.rs
