1use std::collections::HashMap;
2
3use itertools::Itertools;
4
5use crate::{
6 BinRel, Types,
7 constraint::{Bind, Constant, Constraint, Expr, Pred, Qualifier},
8 constraint_fragments::ConstraintFragments,
9 graph::topological_sort_sccs,
10 is_constraint_satisfiable,
11 parser::ParsingTypes,
12};
13
14impl<T: Types> Constraint<T> {
15 pub fn depth_first_fragments(&self) -> ConstraintFragments<'_, T> {
24 ConstraintFragments::new(self)
25 }
26
27 pub fn kvar_deps(&self) -> Vec<T::KVar> {
28 match self {
29 Constraint::Conj(_) => panic!("Conjunctions should not occur in fragments"),
30 Constraint::ForAll(bind, inner) => {
31 let mut dependencies = bind.pred.kvars();
32 dependencies.extend_from_slice(&inner.kvar_deps());
33 dependencies
34 }
35 Constraint::Pred(_, _) => vec![],
36 }
37 }
38
39 pub fn kvar_mappings(
40 &self,
41 ) -> (HashMap<T::KVar, Vec<Constraint<T>>>, HashMap<T::KVar, Vec<T::KVar>>) {
42 let mut kvar_to_fragments: HashMap<T::KVar, Vec<Constraint<T>>> = HashMap::new();
43 let mut kvar_to_dependencies: HashMap<T::KVar, Vec<T::KVar>> = HashMap::new();
44 for frag in self.depth_first_fragments() {
45 if let Some(kvar) = frag.fragment_kvar_head() {
46 kvar_to_dependencies
47 .entry(kvar.clone())
48 .or_insert_with(Vec::new)
49 .extend_from_slice(&frag.kvar_deps().into_iter().unique().collect::<Vec<_>>());
50 kvar_to_fragments
51 .entry(kvar.clone())
52 .or_insert_with(Vec::new)
53 .push(frag);
54 }
55 }
56 (kvar_to_fragments, kvar_to_dependencies)
57 }
58
59 pub fn topo_order_fragments(&self) -> Vec<Self> {
60 let (mut kvar_to_fragments, kvar_to_dependencies) = self.kvar_mappings();
61 let topologically_ordered_kvids = topological_sort_sccs::<T>(&kvar_to_dependencies);
62 topologically_ordered_kvids
63 .into_iter()
64 .rev()
65 .flatten()
66 .map(|kvid| kvar_to_fragments.remove(&kvid).unwrap())
67 .flatten()
68 .collect()
69 }
70
71 pub fn fragment_kvar_head(&self) -> Option<T::KVar> {
72 match self {
73 Constraint::ForAll(_bind, inner) => inner.fragment_kvar_head(),
74 Constraint::Pred(Pred::Expr(_expr), _tag) => None,
75 Constraint::Pred(Pred::KVar(name, _args), _tag) => Some(name.clone()),
76 _ => panic!("Conjunctions should not occur in fragments"),
77 }
78 }
79
80 pub fn sub_all_kvars(
81 &self,
82 assignments: &HashMap<T::KVar, Vec<(&Qualifier<T>, Vec<usize>)>>,
83 ) -> Self {
84 match self {
85 Constraint::ForAll(bind, inner) => {
86 Constraint::ForAll(
87 Bind {
88 name: bind.name.clone(),
89 sort: bind.sort.clone(),
90 pred: bind.pred.sub_kvars(assignments),
91 },
92 Box::new(inner.sub_all_kvars(assignments)),
93 )
94 }
95 Constraint::Pred(pred, _tag) => Constraint::Pred(pred.sub_kvars(assignments), None),
96 Constraint::Conj(conjuncts) => {
97 Constraint::Conj(
98 conjuncts
99 .iter()
100 .map(|cstr| cstr.sub_all_kvars(assignments))
101 .collect(),
102 )
103 }
104 }
105 }
106
107 pub fn sub_kvars_except_head(
108 &self,
109 assignments: &HashMap<T::KVar, Vec<(&Qualifier<T>, Vec<usize>)>>,
110 ) -> Self {
111 match self {
112 Constraint::ForAll(bind, inner) => {
113 Constraint::ForAll(
114 Bind {
115 name: bind.name.clone(),
116 sort: bind.sort.clone(),
117 pred: bind.pred.sub_kvars(assignments),
118 },
119 Box::new(inner.sub_kvars_except_head(assignments)),
120 )
121 }
122 Constraint::Pred(pred, _tag) => Constraint::Pred(pred.clone(), None),
123 _ => panic!("Conjunctions should not occur in constraint fragments"),
124 }
125 }
126
127 pub fn sub_head(&self, assignment: &(&Qualifier<T>, Vec<usize>)) -> Self {
128 match self {
129 Constraint::ForAll(bind, inner) => {
130 Constraint::ForAll(bind.clone(), Box::new(inner.sub_head(assignment)))
131 }
132 Constraint::Pred(pred, _tag) => Constraint::Pred(pred.sub_head(assignment), None),
133 _ => panic!("Conjunctions should not occur in constraint fragments"),
134 }
135 }
136
137 pub fn scope(&self, var: &T::KVar) -> Self {
138 self.scope_help(var)
139 .unwrap_or(Constraint::Pred(Pred::Expr(Expr::Constant(Constant::Boolean(true))), None))
140 }
141
142 fn scope_help(&self, var: &T::KVar) -> Option<Constraint<T>> {
143 match self {
144 Constraint::ForAll(bind, inner) => {
145 if bind.pred.kvars().contains(var) {
146 Some(self.clone())
147 } else {
148 inner.scope_help(var)
149 }
150 }
151 Constraint::Pred(Pred::KVar(kvid, _args), _tag) if var.eq(kvid) => Some(self.clone()),
152 Constraint::Pred(_, _) => None,
153 Constraint::Conj(conjuncts) => {
154 match conjuncts
155 .iter()
156 .filter_map(|inner| inner.scope_help(var))
157 .collect::<Vec<Self>>()
158 .as_slice()
159 {
160 [] => Some(self.clone()),
161 [cstr] => Some(cstr.clone()),
162 _ => Some(self.clone()),
163 }
164 }
165 }
166 }
167}
168
169impl Constraint<ParsingTypes> {
170 pub fn uniquify(&mut self) {
171 let mut used_names = HashMap::new();
172 self.uniquify_help(&mut used_names);
173 }
174
175 fn uniquify_help(&mut self, used_names: &mut HashMap<String, u32>) {
176 match self {
177 Constraint::ForAll(bind, inner) => {
178 let bound_name = &bind.name;
179 if let Some(count) = used_names.get_mut(bound_name) {
180 *count += 1;
181 let new_name = format!("{}#{}", bind.name, count);
182 used_names.insert(new_name.clone(), 0);
183 inner.rename(bound_name, &new_name);
184 bind.pred.rename(bound_name, &new_name);
185 bind.name = new_name
186 } else {
187 used_names.insert(bind.name.clone(), 0);
188 }
189 inner.uniquify_help(used_names);
190 }
191 Constraint::Conj(conjuncts) => {
192 conjuncts.iter_mut().for_each(|cstr| {
193 cstr.uniquify_help(used_names);
194 });
195 }
196 Constraint::Pred(_, _) => {}
197 }
198 }
199
200 fn rename(&mut self, from: &String, to: &String) {
201 match self {
202 Constraint::ForAll(bind, inner) => {
203 if !bind.name.eq(from) {
204 bind.name = to.clone();
205 bind.pred.rename(from, to);
206 inner.rename(from, to);
207 }
208 }
209 Constraint::Conj(conjuncts) => {
210 conjuncts.iter_mut().for_each(|cstr| {
211 cstr.rename(from, to);
212 });
213 }
214 Constraint::Pred(pred, _tag) => pred.rename(from, to),
215 }
216 }
217
218 pub fn is_satisfiable(&self) -> bool {
219 is_constraint_satisfiable(self)
220 }
221
222 pub fn sol1(&self, var: &String) -> Vec<(Vec<Bind<ParsingTypes>>, Vec<Expr<ParsingTypes>>)> {
223 match self {
224 Constraint::ForAll(bind, inner) => {
225 inner
226 .sol1(var)
227 .into_iter()
228 .map(|(mut binders, exprs)| {
229 binders.push(bind.clone());
230 (binders, exprs)
231 })
232 .collect()
233 }
234 Constraint::Conj(conjuncts) => {
235 conjuncts
236 .iter()
237 .map(|cstr| cstr.sol1(var))
238 .flatten()
239 .collect()
240 }
241 Constraint::Pred(Pred::KVar(kvid, args), _tag) if var.eq(kvid) => {
242 let args_eq: Vec<Expr<ParsingTypes>> = (0..)
243 .zip(args.iter())
244 .map(|(idx, arg)| {
245 Expr::Atom(
246 BinRel::Eq,
247 Box::new([
248 Expr::Var(format!("karg${}#{}", kvid, idx).to_string()),
249 Expr::Var(arg.clone()),
250 ]),
251 )
252 })
253 .collect();
254 vec![(vec![], args_eq)]
255 }
256 Constraint::Pred(_, _) => vec![],
257 }
258 }
259
260 pub fn elim(&self, vars: &Vec<String>) -> Self {
261 vars.iter().fold(self.clone(), |acc, var| acc.elim1(var))
262 }
263
264 pub fn elim1(&self, var: &String) -> Self {
265 let solution = self.scope(var).sol1(var);
266 self.do_elim(var, &solution)
267 }
268
269 fn do_elim(
270 &self,
271 var: &String,
272 solution: &Vec<(Vec<Bind<ParsingTypes>>, Vec<Expr<ParsingTypes>>)>,
273 ) -> Self {
274 match self {
275 Constraint::Conj(conjuncts) => {
276 Constraint::Conj(
277 conjuncts
278 .iter()
279 .map(|cstr| cstr.do_elim(var, solution))
280 .collect(),
281 )
282 }
283 Constraint::ForAll(Bind { name, sort, pred }, inner) => {
284 let inner_elimmed = inner.do_elim(var, solution);
285 if pred.kvars().contains(var) {
286 let cstrs: Vec<Constraint<ParsingTypes>> = solution
287 .iter()
288 .map(|(binders, eqs)| {
289 let (kvar_instances, other_preds) = pred.partition_pred(var);
290 let kvar_instances_subbed: Vec<Pred<ParsingTypes>> = {
291 kvar_instances
292 .into_iter()
293 .map(|(kvid, args)| {
294 eqs.iter()
295 .enumerate()
296 .zip(args.iter())
297 .map(|((i, eq), arg)| {
298 Pred::Expr(eq.clone().substitute(
299 &format!("karg${}#{}", &kvid, i),
300 arg,
301 ))
302 })
303 .collect::<Vec<_>>()
304 })
305 .flatten()
306 .collect()
307 };
308 let mut preds = kvar_instances_subbed;
309 preds.extend(other_preds.into_iter());
310 let init = Constraint::ForAll(
311 Bind {
312 name: name.clone(),
313 sort: sort.clone(),
314 pred: Pred::And(preds),
315 },
316 Box::new(inner_elimmed.clone()),
317 );
318 binders.iter().fold(init, |acc, binder| {
319 Constraint::ForAll(binder.clone(), Box::new(acc))
320 })
321 })
322 .collect();
323 if cstrs.len() == 1 { cstrs[0].clone() } else { Constraint::Conj(cstrs) }
324 } else {
325 Constraint::ForAll(
326 Bind { name: name.clone(), sort: sort.clone(), pred: pred.clone() },
327 Box::new(inner_elimmed),
328 )
329 }
330 }
331 Constraint::Pred(Pred::KVar(kvid, _args), _tag) if var.eq(kvid) => {
332 Constraint::Pred(Pred::TRUE, None)
333 }
334 cpred => cpred.clone(),
335 }
336 }
337}
338
339impl<T: Types> Pred<T> {
340 pub fn contains_kvars(&self) -> bool {
341 match self {
342 Pred::And(ps) => ps.iter().any(Pred::contains_kvars),
343 Pred::KVar(_, _) => true,
344 Pred::Expr(_) => false,
345 }
346 }
347
348 pub fn kvars(&self) -> Vec<T::KVar> {
349 match self {
350 Pred::KVar(kvid, _args) => vec![kvid.clone()],
351 Pred::Expr(_expr) => vec![],
352 Pred::And(conjuncts) => {
353 conjuncts
354 .iter()
355 .map(Pred::kvars)
356 .flatten()
357 .unique()
358 .collect()
359 }
360 }
361 }
362
363 pub fn sub_kvars(
364 &self,
365 assignment: &HashMap<T::KVar, Vec<(&Qualifier<T>, Vec<usize>)>>,
366 ) -> Self {
367 match self {
368 Pred::KVar(kvid, args) => {
369 let qualifiers = assignment.get(kvid).unwrap();
370 if qualifiers.len() == 0 {
371 return Pred::Expr(Expr::Constant(Constant::Boolean(false)));
372 }
373 if qualifiers.len() == 1 {
374 let qualifier = qualifiers[0].0;
375 return Pred::Expr(
376 qualifier
377 .args
378 .iter()
379 .map(|arg| &arg.0)
380 .zip(qualifiers[0].1.iter().map(|arg_idx| &args[*arg_idx]))
381 .fold(qualifier.body.clone(), |acc, e| acc.substitute(e.0, e.1)),
382 );
383 } else {
384 return Pred::And(
385 qualifiers
386 .iter()
387 .map(|qualifier| {
388 Pred::Expr(
389 qualifier
390 .0
391 .args
392 .iter()
393 .map(|arg| &arg.0)
394 .zip(qualifier.1.iter().map(|arg_idx| &args[*arg_idx]))
395 .fold(qualifier.0.body.clone(), |acc, e| {
396 acc.substitute(e.0, e.1)
397 }),
398 )
399 })
400 .collect(),
401 );
402 }
403 }
404 Pred::Expr(expr) => Pred::Expr(expr.clone()),
405 Pred::And(conjuncts) => {
406 Pred::And(
407 conjuncts
408 .clone()
409 .into_iter()
410 .map(|pred| pred.sub_kvars(assignment))
411 .collect(),
412 )
413 }
414 }
415 }
416
417 pub fn sub_head(&self, assignment: &(&Qualifier<T>, Vec<usize>)) -> Self {
418 match self {
419 Pred::Expr(expr) => Pred::Expr(expr.clone()),
420 Pred::KVar(_kvid, args) => {
421 Pred::Expr(
422 assignment
423 .0
424 .args
425 .iter()
426 .map(|arg| &arg.0)
427 .zip(assignment.1.iter().map(|arg_idx| &args[*arg_idx]))
428 .fold(assignment.0.body.clone(), |acc, e| acc.substitute(e.0, e.1)),
429 )
430 }
431 _ => panic!("Conjunctions should not occur here"),
432 }
433 }
434
435 pub fn partition_pred(&self, var: &T::KVar) -> (Vec<(T::KVar, Vec<T::Var>)>, Vec<Pred<T>>) {
436 let mut kvar_instances = vec![];
437 let mut other_preds = vec![];
438 self.partition_pred_help(var, &mut kvar_instances, &mut other_preds);
439 (kvar_instances, other_preds)
440 }
441
442 fn partition_pred_help(
443 &self,
444 var: &T::KVar,
445 kvar_instances: &mut Vec<(T::KVar, Vec<T::Var>)>,
446 other_preds: &mut Vec<Pred<T>>,
447 ) {
448 match self {
449 Pred::And(conjuncts) => {
450 conjuncts
451 .iter()
452 .for_each(|pred| pred.partition_pred_help(var, kvar_instances, other_preds));
453 }
454 Pred::KVar(kvid, args) if var.eq(kvid) => {
455 kvar_instances.push((kvid.clone(), args.to_vec()));
456 }
457 _ => {
458 other_preds.push(self.clone());
459 }
460 }
461 }
462
463 pub fn rename(&mut self, from: &T::Var, to: &T::Var) {
464 match self {
465 Pred::Expr(expr) => {
466 expr.substitute_in_place(from, to);
467 }
468 Pred::KVar(_kvid, args) => {
469 args.iter_mut().for_each(|arg| {
470 if from.eq(arg) {
471 *arg = to.clone();
472 }
473 });
474 }
475 Pred::And(conjuncts) => {
476 conjuncts.iter_mut().for_each(|pred| pred.rename(from, to));
477 }
478 }
479 }
480}
481
482impl<T: Types> Expr<T> {
483 pub fn substitute_in_place(&mut self, v_from: &T::Var, v_to: &T::Var) {
484 match self {
485 Expr::Var(v) => {
486 if v == v_from {
487 *v = v_to.clone();
488 }
489 }
490 Expr::Constant(_) => {}
491 Expr::BinaryOp(_, operands) => {
492 operands[0].substitute_in_place(v_from, v_to);
493 operands[1].substitute_in_place(v_from, v_to);
494 }
495 Expr::Atom(_, operands) => {
496 operands[0].substitute_in_place(v_from, v_to);
497 operands[1].substitute_in_place(v_from, v_to);
498 }
499 Expr::And(conjuncts) => {
500 conjuncts
501 .iter_mut()
502 .for_each(|expr| expr.substitute_in_place(v_from, v_to));
503 }
504 Expr::App(func, args) => {
505 func.substitute_in_place(v_from, v_to);
506 args.iter_mut()
507 .for_each(|expr| expr.substitute_in_place(v_from, v_to));
508 }
509 _ => panic!("Not supported yet; implement as needed"),
510 }
511 }
512
513 pub fn substitute(&self, v_from: &T::Var, v_to: &T::Var) -> Self {
514 let mut new_expr = self.clone();
515 new_expr.substitute_in_place(v_from, v_to);
516 new_expr
517 }
518}