Skip to main content

kermit_algos/
const_rewrite.rs

1//! Const-view rewrite implementing Veldhuizen 2014 §3.4 point 4.
2//!
3//! Transforms body atoms (e.g. `p(X, c42)`) into fresh variables
4//! filtered by synthetic unary `Const_c42` predicates, so the existing
5//! LFTJ engine can handle them without modification. Intended to run
6//! immediately before [`crate::JoinAlgo::join_iter`].
7
8use {
9    kermit_parser::{JoinQuery, Predicate, Term},
10    std::fmt,
11};
12
13/// Error returned by [`rewrite_atoms`] when an atom does not match the
14/// expected `c<digits>` shape.
15#[derive(Debug, Clone, PartialEq, Eq)]
16pub enum RewriteError {
17    /// An atom was not of the form `c<digits>`. kermit currently only
18    /// supports constants encoded as dictionary IDs using this
19    /// convention.
20    BadAtom(String),
21}
22
23impl fmt::Display for RewriteError {
24    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
25        match self {
26            | RewriteError::BadAtom(s) => write!(
27                f,
28                "atom {s:?} does not match the expected c<digits> shape — kermit currently only \
29                 supports constants encoded as dictionary IDs",
30            ),
31        }
32    }
33}
34
35impl std::error::Error for RewriteError {}
36
37/// Pairs a synthetic predicate name (e.g. `"Const_c42"`) with its
38/// dictionary ID. One entry is produced per rewritten atom occurrence.
39pub type ConstSpec = (String, usize);
40
41/// Rewrites `query.body`: each `Term::Atom("c<id>")` becomes a fresh
42/// variable `K<i>`, with a new unary predicate `Const_c<id>(K<i>)`
43/// appended to the body.
44///
45/// Each atom occurrence gets its own fresh variable, even if the same
46/// dictionary ID appears multiple times. This avoids forcing equality
47/// between unrelated body positions.
48///
49/// # Head asymmetry
50///
51/// **Only body atoms are rewritten.** Head atoms (e.g.
52/// `Q(c5) :- p(X).`) are left unchanged. The head list describes the
53/// output shape and does not flow through the LFTJ engine the way
54/// body predicates do, so filtering there is the parser / caller's
55/// responsibility. The preprocessor emits queries of the form
56/// `Head(V0, …, Vn) :- body.` where every head term is a variable, so
57/// in practice head atoms never reach this function from the WatDiv
58/// pipeline. Keep this asymmetry in mind if authoring queries by
59/// hand: a `Term::Atom` in the head position will not be filtered.
60///
61/// # Errors
62///
63/// Returns [`RewriteError::BadAtom`] if any atom doesn't match `c\d+`.
64pub fn rewrite_atoms(mut query: JoinQuery) -> Result<(JoinQuery, Vec<ConstSpec>), RewriteError> {
65    // Fresh variables use the `K<n>` shape: `K` is just an unlikely letter
66    // (no special meaning) and `<n>` is a counter. To avoid colliding with
67    // user-supplied variables that already happen to be named `K0`, `K1`,
68    // …, we scan both the body and the head for the highest existing
69    // `K<n>` index and start the counter past it. The
70    // `fresh_var_allocation_avoids_existing_k_names` test pins this
71    // behaviour.
72    let mut next_k = highest_k_index(&query).map_or(0, |n| n + 1);
73    let mut specs: Vec<ConstSpec> = Vec::new();
74    let mut new_preds: Vec<Predicate> = Vec::new();
75
76    // Body only — head atoms are intentionally not rewritten; see the
77    // "Head asymmetry" section of this function's doc-comment.
78    for pred in &mut query.body {
79        for term in &mut pred.terms {
80            let atom = match term {
81                | Term::Atom(s) => s.clone(),
82                | _ => continue,
83            };
84            let id = parse_const_atom(&atom)?;
85            let fresh = format!("K{next_k}");
86            next_k += 1;
87            *term = Term::Var(fresh.clone());
88            let const_name = format!("Const_{atom}");
89            new_preds.push(Predicate {
90                name: const_name.clone(),
91                terms: vec![Term::Var(fresh)],
92            });
93            specs.push((const_name, id));
94        }
95    }
96    query.body.extend(new_preds);
97    Ok((query, specs))
98}
99
100fn parse_const_atom(s: &str) -> Result<usize, RewriteError> {
101    let rest = s
102        .strip_prefix('c')
103        .ok_or_else(|| RewriteError::BadAtom(s.to_string()))?;
104    if rest.is_empty() || !rest.chars().all(|c| c.is_ascii_digit()) {
105        return Err(RewriteError::BadAtom(s.to_string()));
106    }
107    rest.parse::<usize>()
108        .map_err(|_| RewriteError::BadAtom(s.to_string()))
109}
110
111fn highest_k_index(query: &JoinQuery) -> Option<usize> {
112    let scan = |p: &Predicate| -> Option<usize> {
113        p.terms
114            .iter()
115            .filter_map(|t| match t {
116                | Term::Var(name) => name.strip_prefix('K').and_then(|r| r.parse::<usize>().ok()),
117                | _ => None,
118            })
119            .max()
120    };
121    query
122        .body
123        .iter()
124        .chain(std::iter::once(&query.head))
125        .filter_map(scan)
126        .max()
127}
128
129#[cfg(test)]
130mod tests {
131    use super::*;
132
133    fn parse(q: &str) -> JoinQuery { q.parse().unwrap() }
134
135    #[test]
136    fn zero_atoms_is_identity() {
137        let q = parse("Q(X) :- p(X), r(X, Y).");
138        let (out, specs) = rewrite_atoms(q.clone()).unwrap();
139        assert_eq!(out, q);
140        assert!(specs.is_empty());
141    }
142
143    #[test]
144    fn single_atom_produces_one_fresh_var_and_one_const_pred() {
145        let q = parse("Q(X) :- p(X, c42).");
146        let (out, specs) = rewrite_atoms(q).unwrap();
147        assert_eq!(out.body.len(), 2);
148        assert_eq!(out.body[0].name, "p");
149        assert!(matches!(out.body[0].terms[1], Term::Var(ref n) if n == "K0"));
150        assert_eq!(out.body[1].name, "Const_c42");
151        assert!(matches!(out.body[1].terms[0], Term::Var(ref n) if n == "K0"));
152        assert_eq!(specs, vec![("Const_c42".into(), 42)]);
153    }
154
155    #[test]
156    fn multiple_atoms_get_distinct_fresh_vars() {
157        let q = parse("Q(X) :- p(X, c42), r(Y, c99).");
158        let (out, specs) = rewrite_atoms(q).unwrap();
159        assert_eq!(out.body.len(), 4);
160        assert_eq!(specs, vec![
161            ("Const_c42".into(), 42),
162            ("Const_c99".into(), 99),
163        ]);
164    }
165
166    #[test]
167    fn repeated_atom_value_gets_distinct_vars_but_same_const_pred() {
168        let q = parse("Q(X) :- p(X, c5), r(Y, c5).");
169        let (out, specs) = rewrite_atoms(q).unwrap();
170        assert_eq!(out.body.len(), 4);
171        assert_eq!(specs.len(), 2);
172        assert_eq!(specs[0].0, "Const_c5");
173        assert_eq!(specs[1].0, "Const_c5");
174        let k0 = match &out.body[0].terms[1] {
175            | Term::Var(n) => n.clone(),
176            | _ => panic!(),
177        };
178        let k1 = match &out.body[1].terms[1] {
179            | Term::Var(n) => n.clone(),
180            | _ => panic!(),
181        };
182        assert_ne!(k0, k1);
183    }
184
185    #[test]
186    fn fresh_var_allocation_avoids_existing_k_names() {
187        let q = parse("Q(K5) :- p(K5, c7).");
188        let (out, _) = rewrite_atoms(q).unwrap();
189        let fresh = match &out.body[0].terms[1] {
190            | Term::Var(n) => n.clone(),
191            | _ => panic!(),
192        };
193        let n: usize = fresh.strip_prefix('K').unwrap().parse().unwrap();
194        assert!(n > 5, "got {fresh}, expected > K5");
195    }
196
197    #[test]
198    fn malformed_atom_errors() {
199        for bad in ["foo", "c", "c1x", "cc5", "x42"] {
200            let q = JoinQuery {
201                head: Predicate {
202                    name: "Q".into(),
203                    terms: vec![Term::Var("X".into())],
204                },
205                body: vec![Predicate {
206                    name: "p".into(),
207                    terms: vec![Term::Var("X".into()), Term::Atom(bad.into())],
208                }],
209            };
210            assert!(
211                matches!(rewrite_atoms(q), Err(RewriteError::BadAtom(_))),
212                "expected error for {bad}"
213            );
214        }
215    }
216
217    #[test]
218    fn placeholders_left_alone() {
219        let q = parse("Q(X) :- p(X, _), r(_, c7).");
220        let (out, specs) = rewrite_atoms(q).unwrap();
221        assert_eq!(out.body.len(), 3);
222        assert!(matches!(out.body[0].terms[1], Term::Placeholder));
223        assert_eq!(specs, vec![("Const_c7".into(), 7)]);
224    }
225
226    #[test]
227    fn head_atoms_are_not_rewritten() {
228        // Head atoms are outside this function's contract; see the
229        // module docstring. This test pins the asymmetry so a future
230        // refactor can't accidentally start rewriting head terms.
231        let q = JoinQuery {
232            head: Predicate {
233                name: "Q".into(),
234                terms: vec![Term::Atom("c5".into()), Term::Var("X".into())],
235            },
236            body: vec![Predicate {
237                name: "p".into(),
238                terms: vec![Term::Var("X".into()), Term::Atom("c7".into())],
239            }],
240        };
241        let (out, specs) = rewrite_atoms(q).unwrap();
242        assert!(matches!(out.head.terms[0], Term::Atom(ref s) if s == "c5"));
243        assert!(matches!(out.head.terms[1], Term::Var(ref n) if n == "X"));
244        assert_eq!(specs, vec![("Const_c7".into(), 7)]);
245    }
246}