logicaffeine_proof/
discrimination.rs

1//! First-order discrimination tree: the shared pattern index under `simp`
2//! rule lookup, `exact?`/`apply?` library search, and `crush`'s E-matching.
3//!
4//! # How it works
5//!
6//! A pattern is flattened to its preorder key sequence — head symbols carrying
7//! their arity, with `*` at each pattern-variable position — and inserted into
8//! a trie. Retrieval walks the trie in step with the query's own flattening:
9//! a symbol edge must match the query symbol exactly; a `*` edge skips one
10//! complete query subterm (the arity annotations make subterm extents
11//! computable without re-walking the term).
12//!
13//! # The contract
14//!
15//! Retrieval NEVER MISSES: every inserted pattern that one-sided matches the
16//! query (per [`crate::unify::match_term_pattern`]) is among the candidates.
17//! It may over-approximate — variable positions are treated independently, so
18//! the non-linear pattern `f(x, x)` is retrieved for `f(a, b)` — and the
19//! matcher is the arbiter. The value of the tree is pruning: a query touches
20//! only patterns sharing its head structure, not the whole rule set.
21//!
22//! Quantified subexpressions flatten to `*` on both sides (opaque), which
23//! preserves never-miss at the cost of extra candidates — the conservative
24//! choice for shapes the matcher itself treats conservatively.
25
26use std::collections::BTreeMap;
27
28use crate::{ProofExpr, ProofTerm};
29
30/// One symbol in a flattened pattern/query.
31#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
32enum Key {
33    /// A head symbol with its arity. The `String` is namespaced by kind
34    /// (predicate/function/constant/connective) so same-named symbols of
35    /// different kinds never collide.
36    Sym(String, usize),
37    /// A variable position — spans one complete subterm.
38    Star,
39}
40
41struct Node<T> {
42    children: BTreeMap<Key, Node<T>>,
43    /// Values whose pattern's key sequence ends exactly here.
44    values: Vec<T>,
45}
46
47impl<T> Default for Node<T> {
48    fn default() -> Self {
49        Node { children: BTreeMap::new(), values: Vec::new() }
50    }
51}
52
53/// A discrimination tree mapping term/expression patterns to payloads.
54///
55/// Term and expression patterns share one tree — their key alphabets are
56/// disjoint by construction, so they never cross-match.
57pub struct DiscTree<T> {
58    root: Node<T>,
59}
60
61impl<T> Default for DiscTree<T> {
62    fn default() -> Self {
63        Self::new()
64    }
65}
66
67impl<T> DiscTree<T> {
68    pub fn new() -> Self {
69        DiscTree { root: Node::default() }
70    }
71
72    /// Index a term pattern (its `Variable`s are the wildcard positions).
73    pub fn insert_term(&mut self, pattern: &ProofTerm, value: T) {
74        let mut keys = Vec::new();
75        flatten_term(pattern, &mut keys);
76        self.insert_keys(&keys, value);
77    }
78
79    /// Index an expression pattern (its `Variable`s are the wildcard positions).
80    pub fn insert_expr(&mut self, pattern: &ProofExpr, value: T) {
81        let mut keys = Vec::new();
82        flatten_expr(pattern, &mut keys);
83        self.insert_keys(&keys, value);
84    }
85
86    /// Every pattern that could one-sided match the query term (superset —
87    /// confirm with [`crate::unify::match_term_pattern`]).
88    pub fn candidates_term(&self, query: &ProofTerm) -> Vec<&T> {
89        let mut keys = Vec::new();
90        flatten_term(query, &mut keys);
91        self.candidates_keys(&keys)
92    }
93
94    /// Every pattern that could one-sided match the query expression
95    /// (superset — confirm with [`crate::unify::match_expr_pattern`]).
96    pub fn candidates_expr(&self, query: &ProofExpr) -> Vec<&T> {
97        let mut keys = Vec::new();
98        flatten_expr(query, &mut keys);
99        self.candidates_keys(&keys)
100    }
101
102    fn insert_keys(&mut self, keys: &[Key], value: T) {
103        let mut node = &mut self.root;
104        for key in keys {
105            node = node.children.entry(key.clone()).or_default();
106        }
107        node.values.push(value);
108    }
109
110    fn candidates_keys(&self, keys: &[Key]) -> Vec<&T> {
111        let jump = subterm_jumps(keys);
112        let mut out = Vec::new();
113        retrieve(&self.root, keys, &jump, 0, &mut out);
114        out
115    }
116}
117
118/// `jump[i]` = the index one past the subterm that starts at position `i`.
119fn subterm_jumps(keys: &[Key]) -> Vec<usize> {
120    let mut jump = vec![0usize; keys.len()];
121    fill_jumps(keys, 0, &mut jump);
122    jump
123}
124
125fn fill_jumps(keys: &[Key], i: usize, jump: &mut [usize]) -> usize {
126    let end = match &keys[i] {
127        Key::Star => i + 1,
128        Key::Sym(_, arity) => {
129            let mut j = i + 1;
130            for _ in 0..*arity {
131                j = fill_jumps(keys, j, jump);
132            }
133            j
134        }
135    };
136    jump[i] = end;
137    end
138}
139
140fn retrieve<'a, T>(
141    node: &'a Node<T>,
142    keys: &[Key],
143    jump: &[usize],
144    i: usize,
145    out: &mut Vec<&'a T>,
146) {
147    if i == keys.len() {
148        out.extend(node.values.iter());
149        return;
150    }
151    match &keys[i] {
152        sym @ Key::Sym(..) => {
153            // A pattern with the same head continues in lockstep.
154            if let Some(child) = node.children.get(sym) {
155                retrieve(child, keys, jump, i + 1, out);
156            }
157            // A pattern variable here swallows this whole query subterm.
158            if let Some(star) = node.children.get(&Key::Star) {
159                retrieve(star, keys, jump, jump[i], out);
160            }
161        }
162        Key::Star => {
163            // The query is opaque here (a query variable or a shape flattened
164            // conservatively): any complete pattern subterm may align with it.
165            let mut landings = Vec::new();
166            skip_pattern_subterm(node, 1, &mut landings);
167            for landing in landings {
168                retrieve(landing, keys, jump, i + 1, out);
169            }
170        }
171    }
172}
173
174/// All trie nodes exactly one complete pattern subterm below `node`.
175fn skip_pattern_subterm<'a, T>(node: &'a Node<T>, pending: usize, out: &mut Vec<&'a Node<T>>) {
176    if pending == 0 {
177        out.push(node);
178        return;
179    }
180    for (key, child) in &node.children {
181        let need = pending - 1
182            + match key {
183                Key::Sym(_, arity) => *arity,
184                Key::Star => 0,
185            };
186        skip_pattern_subterm(child, need, out);
187    }
188}
189
190fn flatten_term(t: &ProofTerm, out: &mut Vec<Key>) {
191    match t {
192        ProofTerm::Variable(_) => out.push(Key::Star),
193        ProofTerm::Constant(s) => out.push(Key::Sym(format!("c:{s}"), 0)),
194        ProofTerm::BoundVarRef(s) => out.push(Key::Sym(format!("b:{s}"), 0)),
195        ProofTerm::Function(name, args) => {
196            out.push(Key::Sym(format!("f:{name}"), args.len()));
197            for a in args {
198                flatten_term(a, out);
199            }
200        }
201        ProofTerm::Group(args) => {
202            out.push(Key::Sym("(,)".to_string(), args.len()));
203            for a in args {
204                flatten_term(a, out);
205            }
206        }
207    }
208}
209
210fn flatten_expr(e: &ProofExpr, out: &mut Vec<Key>) {
211    match e {
212        ProofExpr::Predicate { name, args, .. } => {
213            out.push(Key::Sym(format!("p:{name}"), args.len()));
214            for a in args {
215                flatten_term(a, out);
216            }
217        }
218        ProofExpr::Identity(l, r) => {
219            out.push(Key::Sym("=".to_string(), 2));
220            flatten_term(l, out);
221            flatten_term(r, out);
222        }
223        ProofExpr::Atom(s) => out.push(Key::Sym(format!("a:{s}"), 0)),
224        ProofExpr::And(l, r) => {
225            out.push(Key::Sym("∧".to_string(), 2));
226            flatten_expr(l, out);
227            flatten_expr(r, out);
228        }
229        ProofExpr::Or(l, r) => {
230            out.push(Key::Sym("∨".to_string(), 2));
231            flatten_expr(l, out);
232            flatten_expr(r, out);
233        }
234        ProofExpr::Implies(l, r) => {
235            out.push(Key::Sym("→".to_string(), 2));
236            flatten_expr(l, out);
237            flatten_expr(r, out);
238        }
239        ProofExpr::Iff(l, r) => {
240            out.push(Key::Sym("↔".to_string(), 2));
241            flatten_expr(l, out);
242            flatten_expr(r, out);
243        }
244        ProofExpr::Not(p) => {
245            out.push(Key::Sym("¬".to_string(), 1));
246            flatten_expr(p, out);
247        }
248        // Binders and every other shape are opaque: `*` preserves never-miss
249        // (the matcher treats them conservatively too).
250        _ => out.push(Key::Star),
251    }
252}
logicaffeine_proof/discrimination.rs

logicaffeine_proof/
discrimination.rs
