1use {
9 kermit_parser::{JoinQuery, Predicate, Term},
10 std::fmt,
11};
12
13#[derive(Debug, Clone, PartialEq, Eq)]
16pub enum RewriteError {
17 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
37pub type ConstSpec = (String, usize);
40
41pub fn rewrite_atoms(mut query: JoinQuery) -> Result<(JoinQuery, Vec<ConstSpec>), RewriteError> {
65 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 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 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}