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(..)
174 | ArrayNode::Call { .. }
175 | ArrayNode::Pick { .. }
176 | ArrayNode::RowScope { .. } => {}
177 }
178 }
179 false
180 });
181 for d in dists {
182 for &v in &d.vars {
183 add(SoftKey::Var(v));
184 }
185 }
186 keys
187}
188
189pub fn wrap_soft(id: u32, body: Expr) -> Expr {
191 Expr::Binary(
192 BinOp::LogOr,
193 Box::new(Expr::Var(SOFT_BASE + id, Ty::BIT)),
194 Box::new(body),
195 Ty::BIT,
196 )
197}
198
199pub fn resolve_softs(e: &Expr, discard: &[bool]) -> Result<Expr, u32> {
202 let mapped = e.try_map_children(&mut |c| resolve_softs(c, discard))?;
203 if let Expr::Binary(BinOp::LogOr, a, b, _) = &mapped
204 && let Expr::Var(v, _) = **a
205 && v >= SOFT_BASE
206 {
207 let k = v - SOFT_BASE;
208 return match discard.get(k as usize) {
209 Some(true) => Ok(Expr::Const(Ty::BIT, BigUint::one())),
210 Some(false) => Ok((**b).clone()),
211 None => Err(k),
212 };
213 }
214 Ok(mapped)
215}
216
217pub fn has_soft(e: &Expr) -> bool {
219 e.any_node(&mut |n| matches!(n, Expr::Var(v, _) if *v >= SOFT_BASE))
220}
221
222pub fn solve_relaxing<T, E>(
235 softs: &[Soft],
236 seed: &mut u64,
237 mut attempt: impl FnMut(&[bool], &mut u64) -> Result<T, E>,
238 relaxable: impl Fn(&E) -> bool,
239) -> Result<T, E> {
240 let seed0 = *seed;
241 let initial: Vec<bool> = softs.iter().map(|s| s.disabled).collect();
242 let mut s = seed0;
243 match attempt(&initial, &mut s) {
244 Ok(t) => {
245 *seed = s;
246 return Ok(t);
247 }
248 Err(e) if !relaxable(&e) => return Err(e),
249 Err(_) => {}
250 }
251 let mut best_seed = seed0;
253 let mut best = attempt(&vec![true; softs.len()], &mut best_seed)?;
254 let mut order: Vec<usize> = (0..softs.len()).filter(|&k| !softs[k].disabled).collect();
255 order.sort_by(|&a, &b| softs[b].priority.cmp(&softs[a].priority).then(b.cmp(&a)));
256 let mut keep = vec![false; softs.len()];
257 for k in order {
258 keep[k] = true;
259 let discard: Vec<bool> = keep.iter().map(|&b| !b).collect();
260 let mut s = seed0;
261 match attempt(&discard, &mut s) {
262 Ok(t) => {
263 best = t;
264 best_seed = s;
265 }
266 Err(e) if relaxable(&e) => keep[k] = false,
267 Err(e) => return Err(e),
268 }
269 }
270 *seed = best_seed;
271 Ok(best)
272}
273
274pub fn object_order<K: Clone + PartialEq>(root: &K, contained: &impl Fn(&K) -> Vec<K>) -> Vec<K> {
281 fn walk<K: Clone + PartialEq>(
282 o: &K,
283 path: &mut Vec<K>,
284 contained: &impl Fn(&K) -> Vec<K>,
285 out: &mut Vec<K>,
286 ) {
287 if path.contains(o) {
288 return;
289 }
290 path.push(o.clone());
291 for c in contained(o) {
292 walk(&c, path, contained, out);
293 }
294 path.pop();
295 out.push(o.clone());
296 }
297 let mut all = Vec::new();
298 walk(root, &mut Vec::new(), contained, &mut all);
299 let mut kept: Vec<K> = Vec::new();
300 for o in all.into_iter().rev() {
301 if !kept.contains(&o) {
302 kept.push(o);
303 }
304 }
305 kept.reverse();
306 kept
307}
308
309#[cfg(test)]
310mod tests {
311 use super::*;
312 use crate::constraints::ir::BinOp;
313
314 fn s(object: u32, level: u32, seq: u32) -> Priority {
315 Priority { object, level, seq }
316 }
317
318 fn soft(priority: Priority, keys: Vec<SoftKey>) -> Soft {
319 Soft {
320 priority,
321 dists: Vec::new(),
322 keys,
323 disabled: false,
324 disabled_by: Vec::new(),
325 }
326 }
327
328 struct Class {
333 base: Option<&'static Class>,
334 blocks: &'static [&'static [&'static str]],
338 handles: &'static [(&'static str, &'static Class)],
340 }
341
342 static B1: Class = Class {
343 base: None,
344 blocks: &[&["a1", "a2"]],
345 handles: &[],
346 };
347 static D1: Class = Class {
348 base: Some(&B1),
349 blocks: &[&["b1"]],
350 handles: &[],
351 };
352 static B2: Class = Class {
353 base: None,
354 blocks: &[&["c1"]],
355 handles: &[],
356 };
357 static D2: Class = Class {
358 base: Some(&B2),
359 blocks: &[&["d1"], &["e1"], &["f1"]],
361 handles: &[("p1", &D1), ("p2", &B1), ("p3", &D1)],
362 };
363
364 #[derive(Clone, Copy)]
366 struct Obj {
367 class: &'static Class,
368 id: u32,
369 }
370
371 impl PartialEq for Obj {
373 fn eq(&self, other: &Obj) -> bool {
374 self.id == other.id
375 }
376 }
377
378 fn object_softs(o: Obj, rank: u32, prefix: &str) -> Vec<(Priority, String)> {
381 let mut chain: Vec<&'static Class> = Vec::new();
382 let mut c = Some(o.class);
383 while let Some(k) = c {
384 chain.push(k);
385 c = k.base;
386 }
387 chain.reverse(); let mut seq = 0;
389 let mut out = Vec::new();
390 for (level, class) in chain.iter().enumerate() {
391 for block in class.blocks {
392 for label in *block {
393 out.push((s(rank, level as u32, seq), format!("{prefix}{label}")));
394 seq += 1;
395 }
396 }
397 }
398 out
399 }
400
401 fn contained(same: bool) -> impl Fn(&Obj) -> Vec<Obj> {
404 move |o: &Obj| {
405 o.class
406 .handles
407 .iter()
408 .enumerate()
409 .map(|(i, (_, class))| Obj {
410 class,
411 id: if same && i == 2 { 1 } else { i as u32 + 1 },
412 })
413 .collect()
414 }
415 }
416
417 fn chain(same: bool) -> String {
420 let root = Obj { class: &D2, id: 0 };
421 let objects = object_order(&root, &contained(same));
422 let mut all: Vec<(Priority, String)> = Vec::new();
423 for (rank, o) in objects.iter().enumerate() {
424 let prefix = match (o.class as *const Class, o.id) {
425 (_, 0) => String::new(),
426 (_, 1) if same => "p3.".to_string(),
427 (_, 1) => "p1.".to_string(),
428 (_, 2) => "p2.".to_string(),
429 _ => "p3.".to_string(),
430 };
431 all.extend(object_softs(*o, rank as u32, &prefix));
432 }
433 all.push((s(u32::MAX, u32::MAX, 0), "i1".to_string()));
435 all.push((s(u32::MAX, u32::MAX, 1), "i2".to_string()));
436 all.sort_by_key(|b| std::cmp::Reverse(b.0));
437 all.into_iter()
438 .map(|(_, l)| l)
439 .collect::<Vec<_>>()
440 .join(" ")
441 }
442
443 #[test]
444 fn the_priority_chain_is_the_one_18_5_13_1_gives() {
445 assert_eq!(
446 chain(false),
447 "i2 i1 f1 e1 d1 c1 p3.b1 p3.a2 p3.a1 p2.a2 p2.a1 p1.b1 p1.a2 p1.a1"
448 );
449 }
450
451 #[test]
452 fn the_same_object_through_two_handles_takes_the_higher_rank() {
453 assert_eq!(
454 chain(true),
455 "i2 i1 f1 e1 d1 c1 p3.b1 p3.a2 p3.a1 p2.a2 p2.a1"
456 );
457 }
458
459 #[test]
460 fn object_order_expands_nested_objects_and_survives_cycles() {
461 let contained = |o: &u32| match *o {
463 0 => vec![1, 2],
464 1 => vec![3],
465 3 => vec![0],
466 _ => Vec::new(),
467 };
468 assert_eq!(object_order(&0, &contained), vec![3, 1, 2, 0]);
470 let again = |o: &u32| match *o {
472 0 => vec![1, 2, 1],
473 _ => Vec::new(),
474 };
475 assert_eq!(object_order(&0, &again), vec![2, 1, 0]);
476 }
477
478 #[test]
481 fn a_disable_discards_only_lower_priority_softs_that_mention_the_variable() {
482 let x = SoftKey::Var(0);
483 let y = SoftKey::Var(1);
484 let mut softs = vec![
485 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]), ];
490 let disables = [Disable {
491 priority: s(0, 0, 1),
492 key: x,
493 }];
494 apply_disables(&mut softs, &disables);
495 let flags: Vec<bool> = softs.iter().map(|s| s.disabled).collect();
496 assert_eq!(flags, vec![true, false, false, true]);
497 }
498
499 #[test]
500 fn a_disable_on_a_whole_array_reaches_its_size_and_a_size_disable_only_the_size() {
501 let size = SoftKey::ArraySize(3);
502 let elem = SoftKey::Array(3);
503 assert!(elem.disables(size) && elem.disables(elem));
504 assert!(size.disables(size) && !size.disables(elem));
505 assert!(!SoftKey::Array(4).disables(size));
506 assert!(!SoftKey::Var(3).disables(elem));
507 }
508
509 fn var(v: u32) -> Expr {
512 Expr::Var(v, Ty::unsigned(8))
513 }
514
515 fn lt(a: Expr, b: Expr) -> Expr {
516 Expr::Binary(BinOp::Lt, Box::new(a), Box::new(b), Ty::BIT)
517 }
518
519 #[test]
520 fn a_discarded_soft_is_true_and_a_kept_one_is_its_body() {
521 let body = lt(var(0), Expr::Const(Ty::unsigned(8), BigUint::from(3u32)));
522 let e = Expr::Binary(
523 BinOp::LogAnd,
524 Box::new(wrap_soft(0, body.clone())),
525 Box::new(wrap_soft(1, body.clone())),
526 Ty::BIT,
527 );
528 assert!(has_soft(&e));
529 let one = Expr::Const(Ty::BIT, BigUint::one());
530 let r = resolve_softs(&e, &[false, true]).unwrap();
531 assert_eq!(
532 r,
533 Expr::Binary(BinOp::LogAnd, Box::new(body), Box::new(one), Ty::BIT)
534 );
535 assert!(!has_soft(&r));
536 assert_eq!(resolve_softs(&e, &[false]), Err(1));
537 }
538
539 #[test]
540 fn the_variables_a_soft_mentions_exclude_the_flags_and_selectors() {
541 let e = wrap_soft(0, lt(var(2), var(5)));
542 assert_eq!(
543 keys_of(&e, &[]),
544 vec![SoftKey::Var(2), SoftKey::Var(5)],
545 "the discard flag is not a mention"
546 );
547 let arr = Expr::Array(ArrayNode::Size(7));
548 assert_eq!(keys_of(&arr, &[]), vec![SoftKey::ArraySize(7)]);
549 }
550
551 fn toy(
556 softs: &[Soft],
557 conflict: (usize, usize),
558 hard_fails: bool,
559 ) -> Result<Vec<usize>, &'static str> {
560 let mut seed = 5u64;
561 let mut calls = 0;
562 let r = solve_relaxing(
563 softs,
564 &mut seed,
565 |discard, s| {
566 calls += 1;
567 *s += 1;
568 if hard_fails {
569 return Err("unsat");
570 }
571 let kept: Vec<usize> = (0..discard.len()).filter(|&k| !discard[k]).collect();
572 if kept.contains(&conflict.0) && kept.contains(&conflict.1) {
573 return Err("unsat");
574 }
575 Ok(kept)
576 },
577 |e| *e == "unsat",
578 );
579 let _ = calls;
580 r
581 }
582
583 #[test]
584 fn without_a_contradiction_every_soft_is_kept() {
585 let softs = vec![soft(s(0, 0, 0), vec![]), soft(s(0, 0, 1), vec![])];
586 assert_eq!(toy(&softs, (7, 8), false), Ok(vec![0, 1]));
587 }
588
589 #[test]
590 fn the_higher_priority_of_two_conflicting_softs_wins() {
591 let softs = vec![
592 soft(s(0, 0, 0), vec![]),
593 soft(s(0, 0, 1), vec![]),
594 soft(s(0, 0, 2), vec![]),
595 ];
596 assert_eq!(toy(&softs, (0, 2), false), Ok(vec![1, 2]));
598 assert_eq!(toy(&softs, (1, 2), false), Ok(vec![0, 2]));
600 }
601
602 #[test]
603 fn a_higher_derived_level_beats_a_later_base_declaration() {
604 let softs = vec![soft(s(0, 1, 0), vec![]), soft(s(0, 0, 1), vec![])];
606 assert_eq!(toy(&softs, (0, 1), false), Ok(vec![0]));
607 }
608
609 #[test]
610 fn a_hard_failure_is_reported_and_disabled_softs_never_return() {
611 let softs = vec![soft(s(0, 0, 0), vec![])];
612 assert_eq!(toy(&softs, (7, 8), true), Err("unsat"));
613 let mut d = soft(s(0, 0, 1), vec![]);
614 d.disabled = true;
615 let softs = vec![soft(s(0, 0, 0), vec![]), d];
616 assert_eq!(toy(&softs, (0, 1), false), Ok(vec![0]));
617 }
618
619 #[test]
620 fn the_seed_is_the_kept_attempts_not_the_probes() {
621 let softs = vec![soft(s(0, 0, 0), vec![]), soft(s(0, 0, 1), vec![])];
622 let mut seed = 100u64;
623 let r = solve_relaxing(
624 &softs,
625 &mut seed,
626 |discard, s| {
627 *s += 1;
628 if discard.iter().all(|d| !d) {
629 Err("unsat")
630 } else {
631 Ok(discard.to_vec())
632 }
633 },
634 |e: &&str| *e == "unsat",
635 );
636 assert_eq!(r, Ok(vec![true, false]));
637 assert_eq!(seed, 101, "one draw on top of the call's own seed");
638 }
639}