1use super::dist::{Dist, SELECTOR_BASE};
45use super::ir::{ArrayId, ArrayNode, BinOp, Expr, Ty, VarId};
46use num_bigint::BigUint;
47use num_traits::One;
48
49pub const SOFT_BASE: VarId = 1 << 31;
53
54#[derive(
57 Clone,
58 Copy,
59 Debug,
60 Default,
61 PartialEq,
62 Eq,
63 PartialOrd,
64 Ord,
65 Hash,
66 serde::Serialize,
67 serde::Deserialize,
68)]
69pub struct Priority {
70 pub object: u32,
73 pub level: u32,
76 pub seq: u32,
78}
79
80#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
83pub enum SoftKey {
84 Var(VarId),
86 Array(ArrayId),
89 ArraySize(ArrayId),
91}
92
93impl SoftKey {
94 pub fn disables(self, mention: SoftKey) -> bool {
99 match (self, mention) {
100 (SoftKey::Var(a), SoftKey::Var(b)) => a == b,
101 (SoftKey::Array(a), SoftKey::Array(b) | SoftKey::ArraySize(b)) => a == b,
102 (SoftKey::ArraySize(a), SoftKey::ArraySize(b)) => a == b,
103 _ => false,
104 }
105 }
106}
107
108#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
110pub struct Soft {
111 pub priority: Priority,
113 pub dists: Vec<u32>,
115 pub keys: Vec<SoftKey>,
117 pub disabled: bool,
120 pub disabled_by: Vec<u32>,
125}
126
127#[derive(Clone, Copy, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
129pub struct Disable {
130 pub priority: Priority,
132 pub key: SoftKey,
134}
135
136pub fn apply_disables(softs: &mut [Soft], disables: &[Disable]) {
138 for s in softs {
139 s.disabled_by = disables
140 .iter()
141 .enumerate()
142 .filter(|(_, d)| d.priority > s.priority && s.keys.iter().any(|&k| d.key.disables(k)))
143 .map(|(i, _)| i as u32)
144 .collect();
145 s.disabled = !s.disabled_by.is_empty();
146 }
147}
148
149pub fn keys_of(e: &Expr, dists: &[&Dist]) -> Vec<SoftKey> {
151 let mut keys: Vec<SoftKey> = Vec::new();
152 let mut add = |k: SoftKey| {
153 if !keys.contains(&k) {
154 keys.push(k);
155 }
156 };
157 e.for_each_var(&mut |v, _| {
158 if v < SELECTOR_BASE {
159 add(SoftKey::Var(v));
160 }
161 });
162 e.any_node(&mut |n| {
163 if let Expr::Array(a) = n {
164 match a {
165 ArrayNode::Size(id) => add(SoftKey::ArraySize(*id)),
166 ArrayNode::Elem { array, .. }
167 | ArrayNode::Foreach { array, .. }
168 | ArrayNode::Reduce { array, .. } => add(SoftKey::Array(*array)),
169 ArrayNode::Distinct { arrays, .. } => {
170 arrays.iter().for_each(|a| add(SoftKey::Array(*a)));
171 }
172 ArrayNode::Loop(..) | ArrayNode::Call { .. } | ArrayNode::Pick { .. } => {}
174 }
175 }
176 false
177 });
178 for d in dists {
179 for &v in &d.vars {
180 add(SoftKey::Var(v));
181 }
182 }
183 keys
184}
185
186pub fn wrap_soft(id: u32, body: Expr) -> Expr {
188 Expr::Binary(
189 BinOp::LogOr,
190 Box::new(Expr::Var(SOFT_BASE + id, Ty::BIT)),
191 Box::new(body),
192 Ty::BIT,
193 )
194}
195
196pub fn resolve_softs(e: &Expr, discard: &[bool]) -> Result<Expr, u32> {
199 let mapped = e.try_map_children(&mut |c| resolve_softs(c, discard))?;
200 if let Expr::Binary(BinOp::LogOr, a, b, _) = &mapped
201 && let Expr::Var(v, _) = **a
202 && v >= SOFT_BASE
203 {
204 let k = v - SOFT_BASE;
205 return match discard.get(k as usize) {
206 Some(true) => Ok(Expr::Const(Ty::BIT, BigUint::one())),
207 Some(false) => Ok((**b).clone()),
208 None => Err(k),
209 };
210 }
211 Ok(mapped)
212}
213
214pub fn has_soft(e: &Expr) -> bool {
216 e.any_node(&mut |n| matches!(n, Expr::Var(v, _) if *v >= SOFT_BASE))
217}
218
219pub fn solve_relaxing<T, E>(
232 softs: &[Soft],
233 seed: &mut u64,
234 mut attempt: impl FnMut(&[bool], &mut u64) -> Result<T, E>,
235 relaxable: impl Fn(&E) -> bool,
236) -> Result<T, E> {
237 let seed0 = *seed;
238 let initial: Vec<bool> = softs.iter().map(|s| s.disabled).collect();
239 let mut s = seed0;
240 match attempt(&initial, &mut s) {
241 Ok(t) => {
242 *seed = s;
243 return Ok(t);
244 }
245 Err(e) if !relaxable(&e) => return Err(e),
246 Err(_) => {}
247 }
248 let mut best_seed = seed0;
250 let mut best = attempt(&vec![true; softs.len()], &mut best_seed)?;
251 let mut order: Vec<usize> = (0..softs.len()).filter(|&k| !softs[k].disabled).collect();
252 order.sort_by(|&a, &b| softs[b].priority.cmp(&softs[a].priority).then(b.cmp(&a)));
253 let mut keep = vec![false; softs.len()];
254 for k in order {
255 keep[k] = true;
256 let discard: Vec<bool> = keep.iter().map(|&b| !b).collect();
257 let mut s = seed0;
258 match attempt(&discard, &mut s) {
259 Ok(t) => {
260 best = t;
261 best_seed = s;
262 }
263 Err(e) if relaxable(&e) => keep[k] = false,
264 Err(e) => return Err(e),
265 }
266 }
267 *seed = best_seed;
268 Ok(best)
269}
270
271pub fn object_order<K: Clone + PartialEq>(root: &K, contained: &impl Fn(&K) -> Vec<K>) -> Vec<K> {
278 fn walk<K: Clone + PartialEq>(
279 o: &K,
280 path: &mut Vec<K>,
281 contained: &impl Fn(&K) -> Vec<K>,
282 out: &mut Vec<K>,
283 ) {
284 if path.contains(o) {
285 return;
286 }
287 path.push(o.clone());
288 for c in contained(o) {
289 walk(&c, path, contained, out);
290 }
291 path.pop();
292 out.push(o.clone());
293 }
294 let mut all = Vec::new();
295 walk(root, &mut Vec::new(), contained, &mut all);
296 let mut kept: Vec<K> = Vec::new();
297 for o in all.into_iter().rev() {
298 if !kept.contains(&o) {
299 kept.push(o);
300 }
301 }
302 kept.reverse();
303 kept
304}
305
306#[cfg(test)]
307mod tests {
308 use super::*;
309 use crate::constraints::ir::BinOp;
310
311 fn s(object: u32, level: u32, seq: u32) -> Priority {
312 Priority { object, level, seq }
313 }
314
315 fn soft(priority: Priority, keys: Vec<SoftKey>) -> Soft {
316 Soft {
317 priority,
318 dists: Vec::new(),
319 keys,
320 disabled: false,
321 disabled_by: Vec::new(),
322 }
323 }
324
325 struct Class {
330 base: Option<&'static Class>,
331 blocks: &'static [&'static [&'static str]],
335 handles: &'static [(&'static str, &'static Class)],
337 }
338
339 static B1: Class = Class {
340 base: None,
341 blocks: &[&["a1", "a2"]],
342 handles: &[],
343 };
344 static D1: Class = Class {
345 base: Some(&B1),
346 blocks: &[&["b1"]],
347 handles: &[],
348 };
349 static B2: Class = Class {
350 base: None,
351 blocks: &[&["c1"]],
352 handles: &[],
353 };
354 static D2: Class = Class {
355 base: Some(&B2),
356 blocks: &[&["d1"], &["e1"], &["f1"]],
358 handles: &[("p1", &D1), ("p2", &B1), ("p3", &D1)],
359 };
360
361 #[derive(Clone, Copy)]
363 struct Obj {
364 class: &'static Class,
365 id: u32,
366 }
367
368 impl PartialEq for Obj {
370 fn eq(&self, other: &Obj) -> bool {
371 self.id == other.id
372 }
373 }
374
375 fn object_softs(o: Obj, rank: u32, prefix: &str) -> Vec<(Priority, String)> {
378 let mut chain: Vec<&'static Class> = Vec::new();
379 let mut c = Some(o.class);
380 while let Some(k) = c {
381 chain.push(k);
382 c = k.base;
383 }
384 chain.reverse(); let mut seq = 0;
386 let mut out = Vec::new();
387 for (level, class) in chain.iter().enumerate() {
388 for block in class.blocks {
389 for label in *block {
390 out.push((s(rank, level as u32, seq), format!("{prefix}{label}")));
391 seq += 1;
392 }
393 }
394 }
395 out
396 }
397
398 fn contained(same: bool) -> impl Fn(&Obj) -> Vec<Obj> {
401 move |o: &Obj| {
402 o.class
403 .handles
404 .iter()
405 .enumerate()
406 .map(|(i, (_, class))| Obj {
407 class,
408 id: if same && i == 2 { 1 } else { i as u32 + 1 },
409 })
410 .collect()
411 }
412 }
413
414 fn chain(same: bool) -> String {
417 let root = Obj { class: &D2, id: 0 };
418 let objects = object_order(&root, &contained(same));
419 let mut all: Vec<(Priority, String)> = Vec::new();
420 for (rank, o) in objects.iter().enumerate() {
421 let prefix = match (o.class as *const Class, o.id) {
422 (_, 0) => String::new(),
423 (_, 1) if same => "p3.".to_string(),
424 (_, 1) => "p1.".to_string(),
425 (_, 2) => "p2.".to_string(),
426 _ => "p3.".to_string(),
427 };
428 all.extend(object_softs(*o, rank as u32, &prefix));
429 }
430 all.push((s(u32::MAX, u32::MAX, 0), "i1".to_string()));
432 all.push((s(u32::MAX, u32::MAX, 1), "i2".to_string()));
433 all.sort_by_key(|b| std::cmp::Reverse(b.0));
434 all.into_iter()
435 .map(|(_, l)| l)
436 .collect::<Vec<_>>()
437 .join(" ")
438 }
439
440 #[test]
441 fn the_priority_chain_is_the_one_18_5_13_1_gives() {
442 assert_eq!(
443 chain(false),
444 "i2 i1 f1 e1 d1 c1 p3.b1 p3.a2 p3.a1 p2.a2 p2.a1 p1.b1 p1.a2 p1.a1"
445 );
446 }
447
448 #[test]
449 fn the_same_object_through_two_handles_takes_the_higher_rank() {
450 assert_eq!(
451 chain(true),
452 "i2 i1 f1 e1 d1 c1 p3.b1 p3.a2 p3.a1 p2.a2 p2.a1"
453 );
454 }
455
456 #[test]
457 fn object_order_expands_nested_objects_and_survives_cycles() {
458 let contained = |o: &u32| match *o {
460 0 => vec![1, 2],
461 1 => vec![3],
462 3 => vec![0],
463 _ => Vec::new(),
464 };
465 assert_eq!(object_order(&0, &contained), vec![3, 1, 2, 0]);
467 let again = |o: &u32| match *o {
469 0 => vec![1, 2, 1],
470 _ => Vec::new(),
471 };
472 assert_eq!(object_order(&0, &again), vec![2, 1, 0]);
473 }
474
475 #[test]
478 fn a_disable_discards_only_lower_priority_softs_that_mention_the_variable() {
479 let x = SoftKey::Var(0);
480 let y = SoftKey::Var(1);
481 let mut softs = vec![
482 soft(s(0, 0, 0), vec![x]), soft(s(0, 0, 2), vec![x]), soft(s(0, 0, 0), vec![y]), soft(s(0, 0, 0), vec![x, y]), ];
487 let disables = [Disable {
488 priority: s(0, 0, 1),
489 key: x,
490 }];
491 apply_disables(&mut softs, &disables);
492 let flags: Vec<bool> = softs.iter().map(|s| s.disabled).collect();
493 assert_eq!(flags, vec![true, false, false, true]);
494 }
495
496 #[test]
497 fn a_disable_on_a_whole_array_reaches_its_size_and_a_size_disable_only_the_size() {
498 let size = SoftKey::ArraySize(3);
499 let elem = SoftKey::Array(3);
500 assert!(elem.disables(size) && elem.disables(elem));
501 assert!(size.disables(size) && !size.disables(elem));
502 assert!(!SoftKey::Array(4).disables(size));
503 assert!(!SoftKey::Var(3).disables(elem));
504 }
505
506 fn var(v: u32) -> Expr {
509 Expr::Var(v, Ty::unsigned(8))
510 }
511
512 fn lt(a: Expr, b: Expr) -> Expr {
513 Expr::Binary(BinOp::Lt, Box::new(a), Box::new(b), Ty::BIT)
514 }
515
516 #[test]
517 fn a_discarded_soft_is_true_and_a_kept_one_is_its_body() {
518 let body = lt(var(0), Expr::Const(Ty::unsigned(8), BigUint::from(3u32)));
519 let e = Expr::Binary(
520 BinOp::LogAnd,
521 Box::new(wrap_soft(0, body.clone())),
522 Box::new(wrap_soft(1, body.clone())),
523 Ty::BIT,
524 );
525 assert!(has_soft(&e));
526 let one = Expr::Const(Ty::BIT, BigUint::one());
527 let r = resolve_softs(&e, &[false, true]).unwrap();
528 assert_eq!(
529 r,
530 Expr::Binary(BinOp::LogAnd, Box::new(body), Box::new(one), Ty::BIT)
531 );
532 assert!(!has_soft(&r));
533 assert_eq!(resolve_softs(&e, &[false]), Err(1));
534 }
535
536 #[test]
537 fn the_variables_a_soft_mentions_exclude_the_flags_and_selectors() {
538 let e = wrap_soft(0, lt(var(2), var(5)));
539 assert_eq!(
540 keys_of(&e, &[]),
541 vec![SoftKey::Var(2), SoftKey::Var(5)],
542 "the discard flag is not a mention"
543 );
544 let arr = Expr::Array(ArrayNode::Size(7));
545 assert_eq!(keys_of(&arr, &[]), vec![SoftKey::ArraySize(7)]);
546 }
547
548 fn toy(
553 softs: &[Soft],
554 conflict: (usize, usize),
555 hard_fails: bool,
556 ) -> Result<Vec<usize>, &'static str> {
557 let mut seed = 5u64;
558 let mut calls = 0;
559 let r = solve_relaxing(
560 softs,
561 &mut seed,
562 |discard, s| {
563 calls += 1;
564 *s += 1;
565 if hard_fails {
566 return Err("unsat");
567 }
568 let kept: Vec<usize> = (0..discard.len()).filter(|&k| !discard[k]).collect();
569 if kept.contains(&conflict.0) && kept.contains(&conflict.1) {
570 return Err("unsat");
571 }
572 Ok(kept)
573 },
574 |e| *e == "unsat",
575 );
576 let _ = calls;
577 r
578 }
579
580 #[test]
581 fn without_a_contradiction_every_soft_is_kept() {
582 let softs = vec![soft(s(0, 0, 0), vec![]), soft(s(0, 0, 1), vec![])];
583 assert_eq!(toy(&softs, (7, 8), false), Ok(vec![0, 1]));
584 }
585
586 #[test]
587 fn the_higher_priority_of_two_conflicting_softs_wins() {
588 let softs = vec![
589 soft(s(0, 0, 0), vec![]),
590 soft(s(0, 0, 1), vec![]),
591 soft(s(0, 0, 2), vec![]),
592 ];
593 assert_eq!(toy(&softs, (0, 2), false), Ok(vec![1, 2]));
595 assert_eq!(toy(&softs, (1, 2), false), Ok(vec![0, 2]));
597 }
598
599 #[test]
600 fn a_higher_derived_level_beats_a_later_base_declaration() {
601 let softs = vec![soft(s(0, 1, 0), vec![]), soft(s(0, 0, 1), vec![])];
603 assert_eq!(toy(&softs, (0, 1), false), Ok(vec![0]));
604 }
605
606 #[test]
607 fn a_hard_failure_is_reported_and_disabled_softs_never_return() {
608 let softs = vec![soft(s(0, 0, 0), vec![])];
609 assert_eq!(toy(&softs, (7, 8), true), Err("unsat"));
610 let mut d = soft(s(0, 0, 1), vec![]);
611 d.disabled = true;
612 let softs = vec![soft(s(0, 0, 0), vec![]), d];
613 assert_eq!(toy(&softs, (0, 1), false), Ok(vec![0]));
614 }
615
616 #[test]
617 fn the_seed_is_the_kept_attempts_not_the_probes() {
618 let softs = vec![soft(s(0, 0, 0), vec![]), soft(s(0, 0, 1), vec![])];
619 let mut seed = 100u64;
620 let r = solve_relaxing(
621 &softs,
622 &mut seed,
623 |discard, s| {
624 *s += 1;
625 if discard.iter().all(|d| !d) {
626 Err("unsat")
627 } else {
628 Ok(discard.to_vec())
629 }
630 },
631 |e: &&str| *e == "unsat",
632 );
633 assert_eq!(r, Ok(vec![true, false]));
634 assert_eq!(seed, 101, "one draw on top of the call's own seed");
635 }
636}