1use std::cell::Cell;
4use std::collections::BTreeSet;
5use std::fmt;
6use std::iter;
7use std::mem;
8use std::ops;
9use std::rc::Rc;
10use bit_set::BitSet;
11
12use ::{Constraint,LinExpr,PsResult,Solution,Val,VarToken};
13use constraint;
14
15#[derive(Clone,Debug,Eq,PartialEq)]
17enum Candidates {
18 None, Value(Val), Set(Rc<BTreeSet<Val>>), }
22
23#[derive(Clone,Debug)]
25enum VarState {
26 Assigned(Val),
27 Unassigned(Candidates),
28 Unified(VarToken),
29}
30
31pub struct Puzzle {
33 num_vars: usize,
35
36 num_guesses: Cell<u32>,
38
39 candidates: Vec<Candidates>,
41
42 constraints: Vec<Rc<Constraint>>,
44}
45
46struct PuzzleConstraints {
48 constraints: Vec<Rc<Constraint>>,
51
52 wake: Vec<BitSet>,
55}
56
57#[derive(Clone)]
59pub struct PuzzleSearch<'a> {
60 puzzle: &'a Puzzle,
61 constraints: Rc<PuzzleConstraints>,
62 vars: Vec<VarState>,
63
64 wake: BitSet,
66}
67
68impl Candidates {
71 fn len(&self) -> usize {
73 match self {
74 &Candidates::None => 0,
75 &Candidates::Value(_) => 1,
76 &Candidates::Set(ref rc) => rc.len(),
77 }
78 }
79
80 fn iter<'a>(&'a self) -> Box<Iterator<Item=Val> + 'a> {
82 match self {
83 &Candidates::None => Box::new(iter::empty()),
84 &Candidates::Value(val) => Box::new(iter::once(val)),
85 &Candidates::Set(ref rc) => Box::new(rc.iter().cloned()),
86 }
87 }
88}
89
90impl Puzzle {
93 pub fn new() -> Self {
101 Puzzle {
102 num_vars: 0,
103 num_guesses: Cell::new(0),
104 candidates: Vec::new(),
105 constraints: Vec::new(),
106 }
107 }
108
109 pub fn new_var(&mut self) -> VarToken {
119 let var = VarToken(self.num_vars);
120 self.num_vars = self.num_vars + 1;
121 self.candidates.push(Candidates::None);
122 var
123 }
124
125 pub fn new_var_with_candidates(&mut self, candidates: &[Val]) -> VarToken {
135 let var = self.new_var();
136 self.insert_candidates(var, candidates);
137 var
138 }
139
140 pub fn new_vars_with_candidates_1d(&mut self, n: usize, candidates: &[Val])
150 -> Vec<VarToken> {
151 let mut vars = Vec::with_capacity(n);
152 for _ in 0..n {
153 vars.push(self.new_var_with_candidates(candidates));
154 }
155 vars
156 }
157
158 pub fn new_vars_with_candidates_2d(self: &mut Puzzle,
168 width: usize, height: usize, candidates: &[Val])
169 -> Vec<Vec<VarToken>> {
170 let mut vars = Vec::with_capacity(height);
171 for _ in 0..height {
172 vars.push(self.new_vars_with_candidates_1d(width, candidates));
173 }
174 vars
175 }
176
177 pub fn set_value(&mut self, var: VarToken, value: Val) {
193 let VarToken(idx) = var;
194
195 if let Candidates::Value(val) = self.candidates[idx] {
196 if val != value {
197 panic!("attempt to set fixed variable");
198 }
199 }
200
201 self.candidates[idx] = Candidates::Value(value);
202 }
203
204 pub fn insert_candidates(&mut self, var: VarToken, candidates: &[Val]) {
216 let VarToken(idx) = var;
217
218 match &self.candidates[idx] {
219 &Candidates::Value(_) =>
220 panic!("attempt to set fixed variable"),
221
222 &Candidates::None => {
223 self.candidates[idx] = Candidates::Set(Rc::new(BTreeSet::new()));
224 },
225
226 &Candidates::Set(_) => (),
227 }
228
229 if let Candidates::Set(ref mut rc) = self.candidates[idx] {
230 let cs = Rc::get_mut(rc).expect("unique");
231 cs.extend(candidates);
232 }
233 }
234
235 pub fn remove_candidates(&mut self, var: VarToken, candidates: &[Val]) {
250 let VarToken(idx) = var;
251
252 match self.candidates[idx] {
253 Candidates::Value(_) =>
254 panic!("attempt to set fixed variable"),
255
256 Candidates::None => (),
257
258 Candidates::Set(ref mut rc) => {
259 let cs = Rc::get_mut(rc).expect("unique");
260 for c in candidates.iter() {
261 cs.remove(c);
262 }
263 },
264 }
265 }
266
267 pub fn intersect_candidates(&mut self, var: VarToken, candidates: &[Val]) {
281 let VarToken(idx) = var;
282
283 match self.candidates[idx] {
284 Candidates::Value(_) =>
285 panic!("attempt to set fixed variable"),
286
287 Candidates::None => (),
288
289 Candidates::Set(ref mut rc) => {
290 let cs = Rc::get_mut(rc).expect("unique");
291 let mut set = BTreeSet::new();
292 set.extend(candidates);
293 *cs = cs.intersection(&set).cloned().collect();
294 },
295 }
296 }
297
298 pub fn add_constraint<T>(&mut self, constraint: T)
300 where T: Constraint + 'static {
301 self.constraints.push(Rc::new(constraint));
302 }
303
304 pub fn all_different<'a, I>(&mut self, vars: I)
316 where I: IntoIterator<Item=&'a VarToken> {
317 self.add_constraint(constraint::AllDifferent::new(vars));
318 }
319
320 pub fn equals<L,R>(&mut self, lhs: L, rhs: R)
332 where L: Into<LinExpr>, R: Into<LinExpr> {
333 self.add_constraint(constraint::Equality::new(lhs.into() - rhs.into()));
334 }
335
336 pub fn unify(&mut self, var1: VarToken, var2: VarToken) {
350 self.add_constraint(constraint::Unify::new(var1, var2));
351 }
352
353 pub fn solve_any(&mut self) -> Option<Solution> {
366 let mut solutions = Vec::with_capacity(1);
367
368 self.num_guesses.set(0);
369 if self.num_vars > 0 {
370 let mut search = PuzzleSearch::new(self);
371 search.solve(1, &mut solutions);
372 }
373
374 solutions.pop()
375 }
376
377 pub fn solve_unique(&mut self) -> Option<Solution> {
391 self.num_guesses.set(0);
392 if self.num_vars > 0 {
393 let mut search = PuzzleSearch::new(self);
394 let mut solutions = Vec::with_capacity(2);
395 search.solve(2, &mut solutions);
396 if solutions.len() == 1 {
397 return solutions.pop();
398 }
399 }
400
401 None
402 }
403
404 pub fn solve_all(&mut self) -> Vec<Solution> {
417 let mut solutions = Vec::new();
418
419 self.num_guesses.set(0);
420 if self.num_vars > 0 {
421 let mut search = PuzzleSearch::new(self);
422 search.solve(::std::usize::MAX, &mut solutions);
423 }
424
425 solutions
426 }
427
428 pub fn step(&mut self) -> Option<PuzzleSearch> {
435 if self.num_vars > 0 {
436 let mut search = PuzzleSearch::new(self);
437 if search.constrain().is_ok() {
438 return Some(search);
439 }
440 }
441
442 None
443 }
444
445 pub fn num_guesses(&self) -> u32 {
447 self.num_guesses.get()
448 }
449}
450
451impl PuzzleConstraints {
454 fn new(puzzle: &Puzzle) -> Self {
456 let wake = Self::init_wake(&puzzle.constraints, puzzle.num_vars);
457 PuzzleConstraints {
458 constraints: puzzle.constraints.clone(),
459 wake: wake,
460 }
461 }
462
463 fn substitute(&self, from: VarToken, to: VarToken) -> PsResult<Self> {
466 let VarToken(idx) = from;
467 let mut new_constraints = self.constraints.clone();
468
469 for cidx in self.wake[idx].iter() {
470 let rc = try!(self.constraints[cidx].substitute(from, to));
471 new_constraints[cidx] = rc;
472 }
473
474 let wake = Self::init_wake(&new_constraints, self.wake.len());
475 Ok(PuzzleConstraints {
476 constraints: new_constraints,
477 wake: wake,
478 })
479 }
480
481 fn init_wake(constraints: &Vec<Rc<Constraint>>, num_vars: usize)
483 -> Vec<BitSet> {
484 let mut wake = vec![BitSet::new(); num_vars];
485 for cidx in 0..constraints.len() {
486 for &VarToken(idx) in constraints[cidx].vars() {
487 wake[idx].insert(cidx);
488 }
489 }
490
491 wake
492 }
493}
494
495impl<'a> PuzzleSearch<'a> {
498 fn new(puzzle: &'a Puzzle) -> Self {
500 let constraints = PuzzleConstraints::new(puzzle);
501 let vars = puzzle.candidates.iter().map(|cs|
502 VarState::Unassigned(cs.clone())).collect();
503 let mut wake = BitSet::new();
504
505 for cidx in 0..constraints.constraints.len() {
506 wake.insert(cidx);
507 }
508
509 PuzzleSearch {
510 puzzle: puzzle,
511 constraints: Rc::new(constraints),
512 vars: vars,
513 wake: wake,
514 }
515 }
516
517 pub fn is_assigned(&self, var: VarToken) -> bool {
519 let VarToken(idx) = var;
520 match &self.vars[idx] {
521 &VarState::Assigned(_) => true,
522 &VarState::Unassigned(_) => false,
523 &VarState::Unified(other) => self.is_assigned(other),
524 }
525 }
526
527 pub fn get_assigned(&self, var: VarToken) -> Option<Val> {
532 let VarToken(idx) = var;
533 match &self.vars[idx] {
534 &VarState::Assigned(val) => Some(val),
535 &VarState::Unassigned(_) => None,
536 &VarState::Unified(other) => self.get_assigned(other),
537 }
538 }
539
540 pub fn get_unassigned(&'a self, var: VarToken)
542 -> Box<Iterator<Item=Val> + 'a> {
543 let VarToken(idx) = var;
544 match &self.vars[idx] {
545 &VarState::Assigned(_) => Box::new(iter::empty()),
546 &VarState::Unassigned(ref cs) => cs.iter(),
547 &VarState::Unified(other) => self.get_unassigned(other),
548 }
549 }
550
551 pub fn get_min_max(&self, var: VarToken) -> PsResult<(Val, Val)> {
553 let VarToken(idx) = var;
554 match &self.vars[idx] {
555 &VarState::Assigned(val) => Ok((val, val)),
556 &VarState::Unassigned(ref cs) => match cs {
557 &Candidates::None => Err(()),
558 &Candidates::Value(val) => Ok((val, val)),
559 &Candidates::Set(ref rc) => {
560 rc.iter().cloned().min().into_iter()
561 .zip(rc.iter().cloned().max()).next()
562 .ok_or(())
563 }
564 },
565 &VarState::Unified(other) => self.get_min_max(other),
566 }
567 }
568
569 pub fn set_candidate(&mut self, var: VarToken, val: Val)
571 -> PsResult<()> {
572 let VarToken(idx) = var;
573
574 match &self.vars[idx] {
575 &VarState::Assigned(v) => return bool_to_result(v == val),
576 &VarState::Unassigned(ref cs) => match cs {
577 &Candidates::None => return Err(()),
578 &Candidates::Value(v) => return bool_to_result(v == val),
579 &Candidates::Set(_) => (),
580 },
581 &VarState::Unified(_) => (),
582 }
583
584 if let &VarState::Unified(other) = &self.vars[idx] {
585 self.set_candidate(other, val)
586 } else if let &mut VarState::Unassigned(Candidates::Set(ref mut rc))
587 = &mut self.vars[idx] {
588 if rc.contains(&val) {
589 let mut set = Rc::make_mut(rc);
590 set.clear();
591 set.insert(val);
592 self.wake.union_with(&self.constraints.wake[idx]);
593 Ok(())
594 } else {
595 Err(())
596 }
597 } else {
598 unreachable!();
599 }
600 }
601
602 pub fn remove_candidate(&mut self, var: VarToken, val: Val)
604 -> PsResult<()> {
605 let VarToken(idx) = var;
606
607 match &self.vars[idx] {
608 &VarState::Assigned(v) => return bool_to_result(v != val),
609 &VarState::Unassigned(ref cs) => match cs {
610 &Candidates::None => return Err(()),
611 &Candidates::Value(v) => return bool_to_result(v != val),
612 &Candidates::Set(_) => (),
613 },
614 &VarState::Unified(_) => (),
615 }
616
617 if let &VarState::Unified(other) = &self.vars[idx] {
618 self.remove_candidate(other, val)
619 } else if let &mut VarState::Unassigned(Candidates::Set(ref mut rc))
620 = &mut self.vars[idx] {
621 if rc.contains(&val) {
622 let mut set = Rc::make_mut(rc);
623 set.remove(&val);
624 self.wake.union_with(&self.constraints.wake[idx]);
625 }
626 bool_to_result(!rc.is_empty())
627 } else {
628 unreachable!();
629 }
630 }
631
632 pub fn bound_candidate_range(&mut self, var: VarToken, min: Val, max: Val)
634 -> PsResult<(Val, Val)> {
635 let VarToken(idx) = var;
636
637 match &self.vars[idx] {
638 &VarState::Assigned(v) =>
639 if min <= v && v <= max {
640 return Ok((v, v))
641 } else {
642 return Err(())
643 },
644 &VarState::Unassigned(ref cs) => match cs {
645 &Candidates::None => return Err(()),
646 &Candidates::Value(v) =>
647 if min <= v && v <= max {
648 return Ok((v, v))
649 } else {
650 return Err(())
651 },
652 &Candidates::Set(_) => (),
653 },
654 &VarState::Unified(_) => (),
655 }
656
657 if let &VarState::Unified(other) = &self.vars[idx] {
658 self.bound_candidate_range(other, min, max)
659 } else if let &mut VarState::Unassigned(Candidates::Set(ref mut rc))
660 = &mut self.vars[idx] {
661 let &curr_min = rc.iter().min().expect("candidates");
662 let &curr_max = rc.iter().max().expect("candidates");
663
664 if curr_min < min || max < curr_max {
665 {
666 let mut set = Rc::make_mut(rc);
667 *set = set.iter()
668 .filter(|&val| min <= *val && *val <= max)
669 .cloned()
670 .collect();
671 self.wake.union_with(&self.constraints.wake[idx]);
672 }
673 rc.iter().cloned().min().into_iter()
674 .zip(rc.iter().cloned().max()).next()
675 .ok_or(())
676 } else {
677 Ok((curr_min, curr_max))
678 }
679 } else {
680 unreachable!();
681 }
682 }
683
684 fn solve(&mut self, count: usize, solutions: &mut Vec<Solution>) {
686 if self.constrain().is_err() {
687 return;
688 }
689
690 let next_unassigned = self.vars.iter().enumerate().min_by_key(
691 |&(_, vs)| match vs {
692 &VarState::Unassigned(ref cs) => cs.len(),
693 _ => ::std::usize::MAX,
694 });
695
696 if let Some((idx, &VarState::Unassigned(ref cs))) = next_unassigned {
697 if cs.len() == 0 {
698 return;
700 }
701
702 for val in cs.iter() {
703 let num_guesses = self.puzzle.num_guesses.get() + 1;
704 self.puzzle.num_guesses.set(num_guesses);
705
706 let mut new = self.clone();
707 if new.assign(idx, val).is_err() {
708 continue;
709 }
710
711 new.solve(count, solutions);
712 if solutions.len() >= count {
713 return;
715 }
716 }
717 } else {
718 let vars = (0..self.puzzle.num_vars).map(|idx|
720 self[VarToken(idx)]).collect();
721 solutions.push(Solution{ vars: vars });
722 }
723 }
724
725 fn assign(&mut self, idx: usize, val: Val) -> PsResult<()> {
727 let var = VarToken(idx);
728 self.vars[idx] = VarState::Assigned(val);
729 self.wake.union_with(&self.constraints.wake[idx]);
730
731 for cidx in 0..self.constraints.constraints.len() {
732 if self.constraints.wake[idx].contains(cidx) {
733 let constraint = self.constraints.constraints[cidx].clone();
734 try!(constraint.on_assigned(self, var, val));
735 }
736 }
737
738 Ok(())
739 }
740
741 fn constrain(&mut self) -> PsResult<()> {
744 while !self.wake.is_empty() {
745 let cycle = self.vars.len();
750 let mut idx = 0;
751 let mut last_gimme = cycle - 1;
752 loop {
753 let gimme = match &self.vars[idx] {
754 &VarState::Assigned(_) => None,
755 &VarState::Unassigned(ref cs) => match cs.len() {
756 0 => return Err(()),
757 1 => cs.iter().next(),
758 _ => None,
759 },
760 &VarState::Unified(_) => None,
761 };
762
763 if let Some(val) = gimme {
764 try!(self.assign(idx, val));
765 last_gimme = idx;
766 } else if idx == last_gimme {
767 break;
768 }
769
770 idx = if idx + 1 >= cycle { 0 } else { idx + 1 };
771 }
772
773 if !self.wake.is_empty() {
775 let wake = mem::replace(&mut self.wake, BitSet::new());
776 for cidx in wake.iter() {
777 let constraint = self.constraints.constraints[cidx].clone();
778 try!(constraint.on_updated(self));
779 }
780 }
781 }
782
783 Ok(())
784 }
785
786 pub fn unify(&mut self, from: VarToken, to: VarToken)
788 -> PsResult<()> {
789 if from == to {
790 return Ok(());
791 }
792
793 let (from, to) = {
796 let VarToken(search) = from;
797 let VarToken(replace) = to;
798 match (&self.vars[search], &self.vars[replace]) {
799 (&VarState::Assigned(val1), &VarState::Assigned(val2)) =>
800 return bool_to_result(val1 == val2),
801
802 (&VarState::Unified(_), _) | (_, &VarState::Unified(_)) =>
804 panic!("internal representation corrupt"),
805
806 (&VarState::Unassigned(_), &VarState::Assigned(_)) =>
807 (to, from),
808
809 (&VarState::Assigned(_), &VarState::Unassigned(_))
810 | (&VarState::Unassigned(_), &VarState::Unassigned(_)) =>
811 (from, to),
812 }
813 };
814
815 let VarToken(search) = from;
816 let VarToken(replace) = to;
817
818 let new_constraints = try!(self.constraints.substitute(from, to));
820 self.constraints = Rc::new(new_constraints);
821 self.wake.union_with(&self.constraints.wake[replace]);
822 assert!(self.constraints.wake[search].is_empty());
823
824 if let &VarState::Assigned(val) = &self.vars[search] {
826 try!(self.set_candidate(to, val));
827 } else {
828 if let (&mut VarState::Unassigned(Candidates::Set(ref mut rc1)),
829 &mut VarState::Unassigned(Candidates::Set(ref mut rc2)))
830 = get_two_mut(&mut self.vars, search, replace) {
831 *rc2 = Rc::new(rc2.intersection(rc1).cloned().collect());
832 if rc2.is_empty() {
833 return Err(());
834 }
835 }
836 }
837
838 self.vars[search] = VarState::Unified(to);
839 Ok(())
840 }
841}
842
843impl<'a> fmt::Debug for PuzzleSearch<'a> {
844 fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
845 write!(f, "PuzzleSearch={{")?;
846 for (idx, var) in self.vars.iter().enumerate() {
847 writeln!(f, "")?;
848 match var {
849 &VarState::Assigned(val) => {
850 write!(f, " var {}: {}", idx, val)?;
851 },
852 &VarState::Unassigned(ref cs) => {
853 write!(f, " var {}:", idx)?;
854 for val in cs.iter() {
855 write!(f, " {}", val)?;
856 }
857 },
858 &VarState::Unified(VarToken(other)) => {
859 write!(f, " var {}: var {}", idx, other)?;
860 },
861 }
862 }
863 write!(f, "}}")?;
864 Ok(())
865 }
866}
867
868impl<'a> ops::Index<VarToken> for PuzzleSearch<'a> {
869 type Output = Val;
870
871 fn index(&self, var: VarToken) -> &Val {
877 let VarToken(idx) = var;
878 match &self.vars[idx] {
879 &VarState::Assigned(ref val) => val,
880 &VarState::Unassigned(_) => panic!("unassigned"),
881 &VarState::Unified(other) => &self[other],
882 }
883 }
884}
885
886fn bool_to_result(cond: bool) -> PsResult<()> {
887 if cond {
888 Ok(())
889 } else {
890 Err(())
891 }
892}
893
894fn get_two_mut<'a, T>(slice: &'a mut [T], a: usize, b: usize)
895 -> (&'a mut T, &'a mut T) {
896 assert!(a != b);
897 if a < b {
898 let (mut l, mut r) = slice.split_at_mut(b);
899 (&mut l[a], &mut r[0])
900 } else {
901 let (l, r) = slice.split_at_mut(a);
902 (&mut r[0], &mut l[b])
903 }
904}
905
906#[cfg(test)]
907mod tests {
908 use ::Puzzle;
909
910 #[test]
911 fn test_no_vars() {
912 let mut sys = Puzzle::new();
913 sys.solve_any();
914 sys.solve_unique();
915 sys.solve_all();
916 sys.step();
917 }
918}