Skip to main content

chess_startpos_rs/
problem.rs

1//! [`Problem`] — the constraint-satisfaction specification and its
2//! enumerate / count entry points.
3
4use std::collections::{HashMap, HashSet};
5use std::fmt;
6
7use crate::{ColorKind, Constraint, PieceKind, SquareColor};
8
9/// Reasons a [`Problem`] can fail [`validate`](Problem::validate).
10///
11/// Returned by [`Problem::validate`] and
12/// [`ProblemBuilder::try_build`]. Variants only carry the
13/// problem-relative indices needed to locate the offending item; they
14/// don't carry `P` or `C` values so that the error type can be
15/// `Eq + Hash + Clone` without dragging those bounds onto the user's
16/// piece / colour kinds.
17#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
18#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
19#[non_exhaustive]
20pub enum ValidationError {
21    /// `square_colors` is non-empty but its length doesn't match
22    /// `num_squares`. An empty `square_colors` is treated as
23    /// "no colour partition declared" and is **not** an error —
24    /// colour-keyed constraints will see zero matches.
25    ColorLengthMismatch {
26        /// `num_squares`.
27        expected: usize,
28        /// `square_colors.len()`.
29        actual: usize,
30    },
31    /// A constraint references a piece kind that is not in the
32    /// problem's alphabet (`pieces`).
33    UnknownPiece,
34    /// A `CountOnColor` constraint references a colour that doesn't
35    /// appear in `square_colors`.
36    UnknownColor,
37    /// An `At` / `NotAt` constraint references a square index
38    /// `>= num_squares`.
39    SquareOutOfRange {
40        /// The offending square index.
41        square: usize,
42        /// `num_squares`.
43        num_squares: usize,
44    },
45    /// An `Order` or `Relative` constraint references a piece
46    /// instance index that exceeds the count declared for that piece
47    /// via `Constraint::Count { Eq, n }`. Only checked when the
48    /// piece has a declared `Eq`-count.
49    InstanceOutOfRange {
50        /// The offending instance index.
51        instance: usize,
52        /// The declared count of that piece.
53        declared: usize,
54    },
55}
56
57impl fmt::Display for ValidationError {
58    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
59        match self {
60            Self::ColorLengthMismatch { expected, actual } => write!(
61                f,
62                "square_colors has length {actual}, expected {expected} (= num_squares)",
63            ),
64            Self::UnknownPiece => f.write_str("constraint references a piece not in the alphabet"),
65            Self::UnknownColor => {
66                f.write_str("CountOnColor references a colour not in square_colors")
67            }
68            Self::SquareOutOfRange {
69                square,
70                num_squares,
71            } => write!(
72                f,
73                "constraint references square {square} but num_squares is {num_squares}",
74            ),
75            Self::InstanceOutOfRange {
76                instance,
77                declared,
78            } => write!(
79                f,
80                "Order / Relative references instance {instance} of a piece declared with count {declared}",
81            ),
82        }
83    }
84}
85
86impl std::error::Error for ValidationError {}
87
88/// A constraint-satisfaction problem: a fixed board (size + per-square
89/// colour from a user-defined colour set), an alphabet of available
90/// piece kinds, and a constraint to satisfy.
91///
92/// `pieces` is the **alphabet** — the set of distinct kinds available.
93/// Duplicate entries are silently deduplicated; first-appearance order
94/// is preserved.
95///
96/// `C` is the colour kind. The default is [`SquareColor`] — the
97/// binary light/dark partition used by chess. For N-way partitions
98/// define your own enum and use it as the `C` type parameter.
99///
100/// `square_colors` must either be empty (no colour partition
101/// declared; colour-keyed constraints always see zero matches) or
102/// have length exactly `num_squares`. Call [`validate`](Self::validate)
103/// to check this and that every constraint references only declared
104/// pieces, colours, and squares.
105///
106/// Solve by calling [`Problem::count`] for the population size or
107/// [`Problem::iter`] to stream all satisfying arrangements.
108///
109/// If `pieces` is empty or `num_squares == 0`, the problem has no
110/// arrangements to enumerate and `count()` returns `0`.
111#[derive(Clone, Debug)]
112#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
113#[cfg_attr(
114    feature = "serde",
115    serde(bound(
116        serialize = "P: serde::Serialize, C: serde::Serialize",
117        deserialize = "P: serde::Deserialize<'de>, C: serde::Deserialize<'de>"
118    ))
119)]
120pub struct Problem<P: PieceKind, C: ColorKind = SquareColor> {
121    /// Number of squares (e.g. `8` for a chess back rank).
122    pub num_squares: usize,
123    /// Colour of each square. `square_colors.len() == num_squares`.
124    pub square_colors: Vec<C>,
125    /// Alphabet of available piece kinds. Duplicates are silently
126    /// deduped by the solver; size doesn't have to equal
127    /// `num_squares`.
128    pub pieces: Vec<P>,
129    /// Root constraint. Use [`Constraint::And`] for conjunctions.
130    pub constraint: Constraint<P, C>,
131}
132
133impl<P: PieceKind, C: ColorKind> Problem<P, C> {
134    /// Checks `self` is internally consistent and returns a
135    /// `ValidationError` if not.
136    ///
137    /// Specifically:
138    /// - `square_colors` is either empty or `num_squares` long.
139    /// - Every piece referenced by a constraint is in `self.pieces`.
140    /// - Every colour referenced by a `CountOnColor` is in
141    ///   `self.square_colors`. An empty `square_colors` is valid
142    ///   only for problems that have **no** `CountOnColor`
143    ///   constraints; otherwise the colour reference is treated as
144    ///   unknown.
145    /// - Every `At` / `NotAt` square index is `< num_squares`.
146    /// - Every `Order` / `Relative` instance index is `< declared
147    ///   count` for that piece (only checked when the piece has a
148    ///   `Constraint::Count { Eq, n }` constraint declaring its
149    ///   count).
150    ///
151    /// `count()` / `iter()` / `sample()` do **not** auto-validate; if
152    /// you need correctness up front, call this first or use
153    /// [`ProblemBuilder::try_build`].
154    pub fn validate(&self) -> Result<(), ValidationError> {
155        if !self.square_colors.is_empty() && self.square_colors.len() != self.num_squares {
156            return Err(ValidationError::ColorLengthMismatch {
157                expected: self.num_squares,
158                actual: self.square_colors.len(),
159            });
160        }
161
162        let alphabet: HashSet<P> = self.dedup_alphabet().into_iter().collect();
163        let counts: HashMap<P, usize> = self.constraint.collect_eq_counts().into_iter().collect();
164        let check_instance = |p: &P, idx: usize| -> Option<ValidationError> {
165            counts
166                .get(p)
167                .filter(|&&declared| idx >= declared)
168                .map(|&declared| ValidationError::InstanceOutOfRange {
169                    instance: idx,
170                    declared,
171                })
172        };
173
174        let mut error: Option<ValidationError> = None;
175        self.constraint.walk(&mut |c| {
176            if error.is_some() {
177                return;
178            }
179            match c {
180                Constraint::Count { piece, .. } if !alphabet.contains(piece) => {
181                    error = Some(ValidationError::UnknownPiece);
182                }
183                Constraint::CountOnColor { piece, color, .. } => {
184                    if !alphabet.contains(piece) {
185                        error = Some(ValidationError::UnknownPiece);
186                    } else if !self.square_colors.contains(color) {
187                        // Includes the case where `square_colors` is
188                        // empty: declaring "no colour partition" and
189                        // then referencing a colour is a contradiction.
190                        error = Some(ValidationError::UnknownColor);
191                    }
192                }
193                Constraint::At { piece, square } | Constraint::NotAt { piece, square } => {
194                    if !alphabet.contains(piece) {
195                        error = Some(ValidationError::UnknownPiece);
196                    } else if *square >= self.num_squares {
197                        error = Some(ValidationError::SquareOutOfRange {
198                            square: *square,
199                            num_squares: self.num_squares,
200                        });
201                    }
202                }
203                Constraint::Order(chain) => {
204                    for (p, idx) in chain {
205                        if !alphabet.contains(p) {
206                            error = Some(ValidationError::UnknownPiece);
207                            return;
208                        }
209                        if let Some(e) = check_instance(p, *idx) {
210                            error = Some(e);
211                            return;
212                        }
213                    }
214                }
215                Constraint::Relative { lhs, rhs, .. }
216                    if !alphabet.contains(&lhs.0) || !alphabet.contains(&rhs.0) =>
217                {
218                    error = Some(ValidationError::UnknownPiece);
219                }
220                Constraint::Relative { lhs, rhs, .. } => {
221                    if let Some(e) = check_instance(&lhs.0, lhs.1) {
222                        error = Some(e);
223                    } else if let Some(e) = check_instance(&rhs.0, rhs.1) {
224                        error = Some(e);
225                    }
226                }
227                _ => {}
228            }
229        });
230        match error {
231            Some(e) => Err(e),
232            None => Ok(()),
233        }
234    }
235
236    /// Returns a fresh empty [`ProblemBuilder`] for fluent
237    /// construction.
238    ///
239    /// ```
240    /// use chess_startpos_rs::{chess, Constraint, CountOp, Problem, SquareColor};
241    ///
242    /// let problem: Problem<chess::Piece> = Problem::builder()
243    ///     .squares(8)
244    ///     .alternating_colors(SquareColor::Dark, SquareColor::Light)
245    ///     .pieces([chess::Piece::King, chess::Piece::Queen])
246    ///     .constraint(Constraint::Count {
247    ///         piece: chess::Piece::King,
248    ///         op: CountOp::Eq,
249    ///         value: 1,
250    ///     })
251    ///     .constraint(Constraint::Count {
252    ///         piece: chess::Piece::Queen,
253    ///         op: CountOp::Eq,
254    ///         value: 7,
255    ///     })
256    ///     .build();
257    /// assert_eq!(problem.count(), 8); // 1 king × 8 placements over 7 queens
258    /// ```
259    ///
260    /// The struct-literal API ([`Problem`] is a regular public-fields
261    /// struct) remains a fully supported alternative.
262    #[must_use]
263    pub fn builder() -> ProblemBuilder<P, C> {
264        ProblemBuilder::new()
265    }
266
267    /// Number of distinct arrangements satisfying the constraint.
268    #[must_use]
269    pub fn count(&self) -> u64 {
270        self.iter().count() as u64
271    }
272
273    /// Streams all distinct arrangements satisfying the constraint
274    /// in canonical lexicographic order.
275    pub fn iter(&self) -> impl Iterator<Item = Vec<P>> + '_ {
276        ProblemIter::new(self)
277    }
278
279    /// Returns the `index`-th satisfying arrangement in canonical
280    /// lexicographic order, or `None` if `index >= self.count()`.
281    ///
282    /// Equivalent to `self.iter().nth(index)` with `u64` indexing.
283    /// O(index) — for repeated indexed lookup prefer iterating once.
284    #[must_use]
285    pub fn at(&self, index: u64) -> Option<Vec<P>> {
286        let idx = usize::try_from(index).ok()?;
287        self.iter().nth(idx)
288    }
289
290    /// Returns a uniformly-random arrangement satisfying the
291    /// constraint, deterministic in `seed`. `None` if the constraint
292    /// is unsatisfiable.
293    ///
294    /// Single pass over `self.iter()` using reservoir-of-size-1
295    /// sampling: each satisfying arrangement is selected with
296    /// probability `1 / k` where `k` is its 1-based index in the
297    /// iterator. The result is uniformly random over the satisfying
298    /// arrangements without materialising them all.
299    #[must_use]
300    pub fn sample(&self, seed: u64) -> Option<Vec<P>> {
301        let mut rng = fastrand::Rng::with_seed(seed);
302        let mut chosen: Option<Vec<P>> = None;
303        for (i, arrangement) in self.iter().enumerate() {
304            // 1-based position used as the reservoir denominator.
305            let seen = (i as u64).saturating_add(1);
306            if rng.u64(0..seen) == 0 {
307                chosen = Some(arrangement);
308            }
309        }
310        chosen
311    }
312
313    /// Returns a copy of `self` with `c` added via AND-composition.
314    #[must_use]
315    pub fn with_constraint(&self, c: Constraint<P, C>) -> Self {
316        let mut next = self.clone();
317        next.constraint = match next.constraint {
318            Constraint::And(mut cs) => {
319                cs.push(c);
320                Constraint::And(cs)
321            }
322            existing => Constraint::And(vec![existing, c]),
323        };
324        next
325    }
326
327    /// Returns the alphabet with duplicates removed, preserving
328    /// first-appearance order. O(n) — uses a `HashSet` to detect
329    /// duplicates.
330    fn dedup_alphabet(&self) -> Vec<P> {
331        let mut seen: HashSet<P> = HashSet::with_capacity(self.pieces.len());
332        let mut ordered: Vec<P> = Vec::with_capacity(self.pieces.len());
333        for p in &self.pieces {
334            if seen.insert(*p) {
335                ordered.push(*p);
336            }
337        }
338        ordered
339    }
340
341    /// Internal optimisation: when every alphabet member has a
342    /// top-level (root or And-nested) `Constraint::Count { Eq, n }`
343    /// fixing its count and the values sum to `num_squares`, the
344    /// declared counts define a single multiset. The enumerator can
345    /// then iterate that multiset's distinct permutations via
346    /// next-permutation instead of running the full Cartesian
347    /// product. Returns `None` (falls back to Cartesian enumeration)
348    /// whenever the fast path can't apply — including when
349    /// `Count{Eq}` constraints sit inside `Or` / `Not` (we don't
350    /// infer counts across disjunction).
351    fn fixed_count_arrangement(&self) -> Option<Vec<P>> {
352        let alphabet = self.dedup_alphabet();
353        if alphabet.is_empty() {
354            return None;
355        }
356
357        let counts = self.constraint.collect_eq_counts();
358        if counts.is_empty() {
359            return None;
360        }
361
362        let mut arrangement: Vec<P> = Vec::with_capacity(self.num_squares);
363        let mut total = 0usize;
364        for kind in &alphabet {
365            // First Eq-count wins if the user accidentally specified
366            // two conflicting counts; the iterator filter will then
367            // catch the contradiction by emitting zero arrangements.
368            let count = counts.iter().find(|(p, _)| p == kind).map(|(_, n)| *n)?;
369            for _ in 0..count {
370                arrangement.push(*kind);
371            }
372            total += count;
373        }
374
375        if total != self.num_squares {
376            return None;
377        }
378        arrangement.sort();
379        Some(arrangement)
380    }
381}
382
383/// Fluent builder for [`Problem`].
384///
385/// Construct via [`Problem::builder`]. Every method consumes `self`
386/// and returns the updated builder, so calls chain. Finalise with
387/// [`ProblemBuilder::build`].
388///
389/// The builder is an alternative to direct struct-literal
390/// construction of `Problem`; both produce the same value.
391///
392/// ```
393/// use chess_startpos_rs::{Constraint, CountOp, Problem, SquareColor};
394///
395/// #[derive(Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Hash, Debug)]
396/// enum Card { Ace, King, Queen }
397///
398/// let problem: Problem<Card> = Problem::builder()
399///     .squares(6)
400///     .alternating_colors(SquareColor::Light, SquareColor::Dark)
401///     .pieces([Card::Ace, Card::King, Card::Queen])
402///     .constraint(Constraint::Count { piece: Card::Ace,   op: CountOp::Eq, value: 2 })
403///     .constraint(Constraint::Count { piece: Card::King,  op: CountOp::Eq, value: 2 })
404///     .constraint(Constraint::Count { piece: Card::Queen, op: CountOp::Eq, value: 2 })
405///     .build();
406///
407/// assert_eq!(problem.count(), 90); // 6! / (2! · 2! · 2!)
408/// ```
409#[derive(Clone, Debug)]
410pub struct ProblemBuilder<P: PieceKind, C: ColorKind> {
411    num_squares: usize,
412    square_colors: Vec<C>,
413    pieces: Vec<P>,
414    constraints: Vec<Constraint<P, C>>,
415}
416
417impl<P: PieceKind, C: ColorKind> Default for ProblemBuilder<P, C> {
418    fn default() -> Self {
419        Self::new()
420    }
421}
422
423impl<P: PieceKind, C: ColorKind> ProblemBuilder<P, C> {
424    /// Returns a fresh empty builder.
425    #[must_use]
426    pub fn new() -> Self {
427        Self {
428            num_squares: 0,
429            square_colors: Vec::new(),
430            pieces: Vec::new(),
431            constraints: Vec::new(),
432        }
433    }
434
435    /// Sets the board size.
436    #[must_use]
437    pub fn squares(mut self, n: usize) -> Self {
438        self.num_squares = n;
439        self
440    }
441
442    /// Sets `square_colors` to `colors`. The slice's length should
443    /// match `num_squares`; mismatch isn't a build-time error but
444    /// will produce empty results from the solver.
445    #[must_use]
446    pub fn colors(mut self, colors: Vec<C>) -> Self {
447        self.square_colors = colors;
448        self
449    }
450
451    /// Sets `square_colors` to a length-`num_squares` sequence
452    /// alternating between `first` (even indices) and `second` (odd
453    /// indices). Call after `.squares(n)`; otherwise produces an
454    /// empty colour vector.
455    #[must_use]
456    pub fn alternating_colors(mut self, first: C, second: C) -> Self
457    where
458        C: Copy,
459    {
460        self.square_colors = crate::alternating(self.num_squares, first, second);
461        self
462    }
463
464    /// Sets `square_colors` to a length-`num_squares` sequence with
465    /// every square coloured `c`. Call after `.squares(n)`.
466    #[must_use]
467    pub fn uniform_colors(mut self, c: C) -> Self
468    where
469        C: Copy,
470    {
471        self.square_colors = crate::uniform(self.num_squares, c);
472        self
473    }
474
475    /// Sets `square_colors` to `(0..num_squares).map(f).collect()`.
476    /// The closure receives each 0-based square index and returns
477    /// its colour. Call after `.squares(n)`.
478    ///
479    /// ```
480    /// use chess_startpos_rs::{Problem, SquareColor};
481    ///
482    /// // 6-square board split into halves: first 3 Light, last 3 Dark.
483    /// let problem: Problem<u8> = Problem::builder()
484    ///     .squares(6)
485    ///     .colors_fn(|i| if i < 3 { SquareColor::Light } else { SquareColor::Dark })
486    ///     .build();
487    /// assert_eq!(problem.square_colors.len(), 6);
488    /// ```
489    #[must_use]
490    pub fn colors_fn<F>(mut self, f: F) -> Self
491    where
492        F: FnMut(usize) -> C,
493    {
494        self.square_colors = (0..self.num_squares).map(f).collect();
495        self
496    }
497
498    /// Replaces the alphabet with `alphabet`. Duplicates are silently
499    /// deduplicated by the solver.
500    #[must_use]
501    pub fn pieces<I: IntoIterator<Item = P>>(mut self, alphabet: I) -> Self {
502        self.pieces = alphabet.into_iter().collect();
503        self
504    }
505
506    /// Appends `p` to the alphabet.
507    #[must_use]
508    pub fn piece(mut self, p: P) -> Self {
509        self.pieces.push(p);
510        self
511    }
512
513    /// Appends a constraint. Multiple calls AND-compose at
514    /// [`build`](Self::build) time.
515    #[must_use]
516    pub fn constraint(mut self, c: Constraint<P, C>) -> Self {
517        self.constraints.push(c);
518        self
519    }
520
521    /// Finalises into a [`Problem`] without validating it. The
522    /// accumulated constraints are wrapped in [`Constraint::And`] (a
523    /// single constraint is stored bare; zero constraints become
524    /// `And(vec![])`, which is always satisfied).
525    ///
526    /// To check internal consistency before solving, use
527    /// [`try_build`](Self::try_build) instead, or call
528    /// [`Problem::validate`] on the returned value.
529    #[must_use]
530    pub fn build(self) -> Problem<P, C> {
531        let constraint = match self.constraints.len() {
532            0 => Constraint::And(Vec::new()),
533            1 => self.constraints.into_iter().next().expect("len == 1"),
534            _ => Constraint::And(self.constraints),
535        };
536        Problem {
537            num_squares: self.num_squares,
538            square_colors: self.square_colors,
539            pieces: self.pieces,
540            constraint,
541        }
542    }
543
544    /// Like [`build`](Self::build) but runs [`Problem::validate`] on
545    /// the result, returning the validation error instead of the
546    /// problem if it fails.
547    ///
548    /// ```
549    /// use chess_startpos_rs::{chess, Problem, SquareColor, ValidationError};
550    ///
551    /// // Builder forgot to call `.alternating_colors(...)` after
552    /// // `.squares(8)`; `square_colors` ends up length 0.
553    /// // That's allowed iff no constraint references colours — but
554    /// // here we omit colours *and* don't add any constraints, so
555    /// // it's fine. Now break it by mismatching the colour vec:
556    /// let bad: Result<Problem<chess::Piece>, _> = Problem::builder()
557    ///     .squares(8)
558    ///     .colors(vec![SquareColor::Light]) // length 1, not 8
559    ///     .pieces([chess::Piece::King])
560    ///     .try_build();
561    /// assert!(matches!(
562    ///     bad,
563    ///     Err(ValidationError::ColorLengthMismatch { expected: 8, actual: 1 }),
564    /// ));
565    /// ```
566    pub fn try_build(self) -> Result<Problem<P, C>, ValidationError> {
567        let problem = self.build();
568        problem.validate()?;
569        Ok(problem)
570    }
571}
572
573/// Iterator over satisfying arrangements; see [`Problem::iter`].
574struct ProblemIter<'a, P: PieceKind, C: ColorKind> {
575    problem: &'a Problem<P, C>,
576    state: IterState<P>,
577}
578
579enum IterState<P: PieceKind> {
580    /// Fully constrained — the declared counts give us a single
581    /// multiset; iterate its distinct permutations via the
582    /// next-permutation algorithm. Each emitted permutation is then
583    /// filtered against the rest of the constraint tree.
584    /// `current` is the next candidate to emit (before filtering).
585    Permutation { current: Option<Vec<P>> },
586    /// Partial / unconstrained — iterate Cartesian-product
587    /// length-N sequences from the alphabet. `current` is the next
588    /// candidate to emit. `index` maps each alphabet member to its
589    /// position so advancing is O(1) per slot.
590    Cartesian {
591        alphabet: Vec<P>,
592        index: HashMap<P, usize>,
593        current: Option<Vec<P>>,
594    },
595}
596
597impl<'a, P: PieceKind, C: ColorKind> ProblemIter<'a, P, C> {
598    fn new(problem: &'a Problem<P, C>) -> Self {
599        let state = match problem.fixed_count_arrangement() {
600            Some(start) => IterState::Permutation {
601                current: Some(start),
602            },
603            None => {
604                let alphabet = problem.dedup_alphabet();
605                let current = if alphabet.is_empty() || problem.num_squares == 0 {
606                    None
607                } else {
608                    Some(vec![alphabet[0]; problem.num_squares])
609                };
610                let index = alphabet.iter().enumerate().map(|(i, p)| (*p, i)).collect();
611                IterState::Cartesian {
612                    alphabet,
613                    index,
614                    current,
615                }
616            }
617        };
618        Self { problem, state }
619    }
620}
621
622impl<P: PieceKind, C: ColorKind> Iterator for ProblemIter<'_, P, C> {
623    type Item = Vec<P>;
624
625    fn next(&mut self) -> Option<Self::Item> {
626        loop {
627            let candidate = match &mut self.state {
628                IterState::Permutation { current } => {
629                    let candidate = current.take()?;
630                    let mut next = candidate.clone();
631                    if next_permutation(&mut next) {
632                        *current = Some(next);
633                    }
634                    candidate
635                }
636                IterState::Cartesian {
637                    alphabet,
638                    index,
639                    current,
640                } => {
641                    let candidate = current.take()?;
642                    let mut next = candidate.clone();
643                    if next_cartesian(&mut next, alphabet, index) {
644                        *current = Some(next);
645                    }
646                    candidate
647                }
648            };
649
650            if candidate.len() == self.problem.num_squares
651                && self
652                    .problem
653                    .constraint
654                    .evaluate(&candidate, &self.problem.square_colors)
655            {
656                return Some(candidate);
657            }
658        }
659    }
660}
661
662/// Standard next-permutation algorithm (lexicographic). When the
663/// input contains duplicate elements, equal-element swaps never
664/// produce a new ordering, so the iteration naturally yields each
665/// distinct permutation once. Returns `false` if `v` was already
666/// the last permutation.
667fn next_permutation<T: Ord>(v: &mut [T]) -> bool {
668    let n = v.len();
669    if n < 2 {
670        return false;
671    }
672    // Step 1: find the largest i such that v[i-1] < v[i].
673    let mut i = n - 1;
674    while i > 0 && v[i - 1] >= v[i] {
675        i -= 1;
676    }
677    if i == 0 {
678        return false;
679    }
680    // Step 2: find the largest j >= i such that v[j] > v[i-1].
681    let mut j = n - 1;
682    while v[j] <= v[i - 1] {
683        j -= 1;
684    }
685    // Step 3: swap v[i-1] and v[j].
686    v.swap(i - 1, j);
687    // Step 4: reverse v[i..].
688    v[i..].reverse();
689    true
690}
691
692/// Lexicographically advances `v` to the next length-`v.len()`
693/// sequence drawn from `alphabet` (Cartesian product, repetition
694/// allowed). `index` maps each alphabet member to its position
695/// (built once in [`ProblemIter::new`]). Returns `false` if `v`
696/// was already the last sequence.
697fn next_cartesian<P: PieceKind>(v: &mut [P], alphabet: &[P], index: &HashMap<P, usize>) -> bool {
698    if alphabet.is_empty() {
699        return false;
700    }
701    let last = alphabet[alphabet.len() - 1];
702    for i in (0..v.len()).rev() {
703        if v[i] != last {
704            let pos = *index.get(&v[i]).expect("in alphabet");
705            v[i] = alphabet[pos + 1];
706            for slot in v.iter_mut().skip(i + 1) {
707                *slot = alphabet[0];
708            }
709            return true;
710        }
711    }
712    false
713}
714
715#[cfg(test)]
716mod tests {
717    use super::*;
718    use crate::{alternating, CountOp};
719
720    /// A tiny piece kind for unit tests.
721    #[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
722    #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
723    enum Tile {
724        A,
725        B,
726        C,
727    }
728
729    fn light_dark(n: usize) -> Vec<SquareColor> {
730        alternating(n, SquareColor::Light, SquareColor::Dark)
731    }
732
733    fn fixed_counts(items: &[(Tile, usize)]) -> Constraint<Tile> {
734        Constraint::And(
735            items
736                .iter()
737                .map(|(p, n)| Constraint::Count {
738                    piece: *p,
739                    op: CountOp::Eq,
740                    value: *n,
741                })
742                .collect(),
743        )
744    }
745
746    #[test]
747    fn count_constraint_drives_piece_counts() {
748        // Alphabet {A, B}; counts {A: 2, B: 1}; expect 3 perms over
749        // the resulting [A, A, B] arrangement.
750        let problem = Problem {
751            num_squares: 3,
752            square_colors: light_dark(3),
753            pieces: vec![Tile::A, Tile::B],
754            constraint: fixed_counts(&[(Tile::A, 2), (Tile::B, 1)]),
755        };
756        assert_eq!(problem.count(), 3);
757    }
758
759    #[test]
760    fn duplicates_in_pieces_are_deduped() {
761        let problem = Problem {
762            num_squares: 3,
763            square_colors: light_dark(3),
764            pieces: vec![Tile::A, Tile::A, Tile::B], // duplicate Tile::A
765            constraint: fixed_counts(&[(Tile::A, 2), (Tile::B, 1)]),
766        };
767        assert_eq!(problem.count(), 3);
768    }
769
770    #[test]
771    fn at_constraint() {
772        let problem = Problem {
773            num_squares: 3,
774            square_colors: light_dark(3),
775            pieces: vec![Tile::A, Tile::B, Tile::C],
776            constraint: Constraint::And(vec![
777                Constraint::Count {
778                    piece: Tile::A,
779                    op: CountOp::Eq,
780                    value: 1,
781                },
782                Constraint::Count {
783                    piece: Tile::B,
784                    op: CountOp::Eq,
785                    value: 1,
786                },
787                Constraint::Count {
788                    piece: Tile::C,
789                    op: CountOp::Eq,
790                    value: 1,
791                },
792                Constraint::At {
793                    piece: Tile::A,
794                    square: 0,
795                },
796            ]),
797        };
798        assert_eq!(problem.count(), 2);
799    }
800
801    #[test]
802    fn unconstrained_cartesian_fallback() {
803        // No Count-Eq constraints; alphabet of 2 on 3 squares = 2^3 = 8.
804        let problem = Problem {
805            num_squares: 3,
806            square_colors: light_dark(3),
807            pieces: vec![Tile::A, Tile::B],
808            constraint: Constraint::And(vec![]),
809        };
810        assert_eq!(problem.count(), 8);
811    }
812
813    #[test]
814    fn unconstrained_with_at_filter() {
815        // Cartesian fallback + At filter.
816        let problem = Problem {
817            num_squares: 3,
818            square_colors: light_dark(3),
819            pieces: vec![Tile::A, Tile::B],
820            constraint: Constraint::At {
821                piece: Tile::A,
822                square: 0,
823            },
824        };
825        // Tile::A at square 0 + any of 2^2 sequences on squares 1, 2.
826        assert_eq!(problem.count(), 4);
827    }
828
829    #[test]
830    fn order_constraint_king_between_rooks() {
831        #[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
832        enum Piece {
833            R,
834            K,
835        }
836        let problem = Problem {
837            num_squares: 3,
838            square_colors: light_dark(3),
839            pieces: vec![Piece::R, Piece::K],
840            constraint: Constraint::And(vec![
841                Constraint::Count {
842                    piece: Piece::R,
843                    op: CountOp::Eq,
844                    value: 2,
845                },
846                Constraint::Count {
847                    piece: Piece::K,
848                    op: CountOp::Eq,
849                    value: 1,
850                },
851                Constraint::Order(vec![(Piece::R, 0), (Piece::K, 0), (Piece::R, 1)]),
852            ]),
853        };
854        assert_eq!(problem.count(), 1);
855    }
856
857    #[test]
858    fn order_with_out_of_range_instance_is_unsatisfied() {
859        let problem = Problem {
860            num_squares: 4,
861            square_colors: light_dark(4),
862            pieces: vec![Tile::A, Tile::B, Tile::C],
863            constraint: Constraint::And(vec![
864                Constraint::Count {
865                    piece: Tile::A,
866                    op: CountOp::Eq,
867                    value: 1,
868                },
869                Constraint::Count {
870                    piece: Tile::B,
871                    op: CountOp::Eq,
872                    value: 2,
873                },
874                Constraint::Count {
875                    piece: Tile::C,
876                    op: CountOp::Eq,
877                    value: 1,
878                },
879                Constraint::Order(vec![
880                    (Tile::B, 0),
881                    (Tile::B, 1),
882                    (Tile::B, 2), // doesn't exist
883                ]),
884            ]),
885        };
886        assert_eq!(problem.count(), 0);
887    }
888
889    #[test]
890    fn relative_constraint_exact_offset() {
891        let problem = Problem {
892            num_squares: 3,
893            square_colors: light_dark(3),
894            pieces: vec![Tile::A, Tile::B, Tile::C],
895            constraint: Constraint::And(vec![
896                Constraint::Count {
897                    piece: Tile::A,
898                    op: CountOp::Eq,
899                    value: 1,
900                },
901                Constraint::Count {
902                    piece: Tile::B,
903                    op: CountOp::Eq,
904                    value: 1,
905                },
906                Constraint::Count {
907                    piece: Tile::C,
908                    op: CountOp::Eq,
909                    value: 1,
910                },
911                Constraint::Relative {
912                    lhs: (Tile::A, 0),
913                    rhs: (Tile::B, 0),
914                    op: CountOp::Eq,
915                    offset: 1,
916                },
917            ]),
918        };
919        let arrangements: Vec<Vec<Tile>> = problem.iter().collect();
920        assert_eq!(
921            arrangements,
922            vec![
923                vec![Tile::B, Tile::A, Tile::C],
924                vec![Tile::C, Tile::B, Tile::A],
925            ],
926        );
927    }
928
929    #[test]
930    fn relative_constraint_absolute_distance_via_and() {
931        let problem = Problem {
932            num_squares: 3,
933            square_colors: light_dark(3),
934            pieces: vec![Tile::A, Tile::B, Tile::C],
935            constraint: Constraint::And(vec![
936                Constraint::Count {
937                    piece: Tile::A,
938                    op: CountOp::Eq,
939                    value: 1,
940                },
941                Constraint::Count {
942                    piece: Tile::B,
943                    op: CountOp::Eq,
944                    value: 1,
945                },
946                Constraint::Count {
947                    piece: Tile::C,
948                    op: CountOp::Eq,
949                    value: 1,
950                },
951                Constraint::Relative {
952                    lhs: (Tile::A, 0),
953                    rhs: (Tile::B, 0),
954                    op: CountOp::Le,
955                    offset: 1,
956                },
957                Constraint::Relative {
958                    lhs: (Tile::A, 0),
959                    rhs: (Tile::B, 0),
960                    op: CountOp::Ge,
961                    offset: -1,
962                },
963            ]),
964        };
965        assert_eq!(problem.count(), 4);
966    }
967
968    #[test]
969    fn and_or_not_combinators() {
970        let problem = Problem {
971            num_squares: 3,
972            square_colors: light_dark(3),
973            pieces: vec![Tile::A, Tile::B, Tile::C],
974            constraint: Constraint::And(vec![
975                Constraint::Count {
976                    piece: Tile::A,
977                    op: CountOp::Eq,
978                    value: 1,
979                },
980                Constraint::Count {
981                    piece: Tile::B,
982                    op: CountOp::Eq,
983                    value: 1,
984                },
985                Constraint::Count {
986                    piece: Tile::C,
987                    op: CountOp::Eq,
988                    value: 1,
989                },
990                Constraint::Not(Box::new(Constraint::At {
991                    piece: Tile::A,
992                    square: 0,
993                })),
994                Constraint::Or(vec![
995                    Constraint::At {
996                        piece: Tile::B,
997                        square: 0,
998                    },
999                    Constraint::At {
1000                        piece: Tile::C,
1001                        square: 0,
1002                    },
1003                ]),
1004            ]),
1005        };
1006        assert_eq!(problem.count(), 4);
1007    }
1008
1009    #[test]
1010    fn count_on_color() {
1011        // Two A's that must both be on light squares (positions 0, 2).
1012        let problem = Problem {
1013            num_squares: 3,
1014            square_colors: light_dark(3), // light, dark, light
1015            pieces: vec![Tile::A, Tile::B],
1016            constraint: Constraint::And(vec![
1017                Constraint::Count {
1018                    piece: Tile::A,
1019                    op: CountOp::Eq,
1020                    value: 2,
1021                },
1022                Constraint::Count {
1023                    piece: Tile::B,
1024                    op: CountOp::Eq,
1025                    value: 1,
1026                },
1027                Constraint::CountOnColor {
1028                    piece: Tile::A,
1029                    color: SquareColor::Light,
1030                    op: CountOp::Eq,
1031                    value: 2,
1032                },
1033            ]),
1034        };
1035        assert_eq!(problem.count(), 1);
1036    }
1037
1038    #[test]
1039    fn at_returns_lexicographic_arrangements() {
1040        let problem = Problem {
1041            num_squares: 3,
1042            square_colors: light_dark(3),
1043            pieces: vec![Tile::A, Tile::B],
1044            constraint: fixed_counts(&[(Tile::A, 2), (Tile::B, 1)]),
1045        };
1046        assert_eq!(problem.at(0), Some(vec![Tile::A, Tile::A, Tile::B]));
1047        assert_eq!(problem.at(1), Some(vec![Tile::A, Tile::B, Tile::A]));
1048        assert_eq!(problem.at(2), Some(vec![Tile::B, Tile::A, Tile::A]));
1049        assert_eq!(problem.at(3), None);
1050    }
1051
1052    #[test]
1053    fn sample_is_deterministic_and_in_range() {
1054        let problem = Problem {
1055            num_squares: 3,
1056            square_colors: light_dark(3),
1057            pieces: vec![Tile::A, Tile::B, Tile::C],
1058            constraint: Constraint::And(vec![
1059                Constraint::Count {
1060                    piece: Tile::A,
1061                    op: CountOp::Eq,
1062                    value: 1,
1063                },
1064                Constraint::Count {
1065                    piece: Tile::B,
1066                    op: CountOp::Eq,
1067                    value: 1,
1068                },
1069                Constraint::Count {
1070                    piece: Tile::C,
1071                    op: CountOp::Eq,
1072                    value: 1,
1073                },
1074            ]),
1075        };
1076        let arrangements: Vec<Vec<Tile>> = problem.iter().collect();
1077        let first = problem.sample(42).expect("non-empty");
1078        let again = problem.sample(42).expect("non-empty");
1079        assert_eq!(first, again);
1080        assert!(arrangements.contains(&first));
1081    }
1082
1083    #[test]
1084    fn sample_returns_none_for_unsatisfiable() {
1085        let problem = Problem {
1086            num_squares: 3,
1087            square_colors: light_dark(3),
1088            pieces: vec![Tile::A, Tile::B, Tile::C],
1089            constraint: Constraint::And(vec![
1090                Constraint::Count {
1091                    piece: Tile::A,
1092                    op: CountOp::Eq,
1093                    value: 1,
1094                },
1095                Constraint::Count {
1096                    piece: Tile::B,
1097                    op: CountOp::Eq,
1098                    value: 1,
1099                },
1100                Constraint::Count {
1101                    piece: Tile::C,
1102                    op: CountOp::Eq,
1103                    value: 1,
1104                },
1105                Constraint::At {
1106                    piece: Tile::A,
1107                    square: 0,
1108                },
1109                Constraint::NotAt {
1110                    piece: Tile::A,
1111                    square: 0,
1112                },
1113            ]),
1114        };
1115        assert_eq!(problem.count(), 0);
1116        assert_eq!(problem.sample(0), None);
1117    }
1118
1119    #[test]
1120    fn count_op_all_six_variants_filter_correctly() {
1121        // 3-square board, alphabet {A, B}. 2^3 = 8 sequences without
1122        // any Count-Eq constraint. Pre-compute the count of As in each:
1123        // 0:[B,B,B]=0 1:[A,B,B]=1 2:[B,A,B]=1 3:[A,A,B]=2
1124        // 4:[B,B,A]=1 5:[A,B,A]=2 6:[B,A,A]=2 7:[A,A,A]=3
1125        // → A-counts {0,1,1,2,1,2,2,3} → distribution 1×0, 3×1, 3×2, 1×3.
1126        let make = |op: CountOp, value: usize| Problem {
1127            num_squares: 3,
1128            square_colors: light_dark(3),
1129            pieces: vec![Tile::A, Tile::B],
1130            constraint: Constraint::Count {
1131                piece: Tile::A,
1132                op,
1133                value,
1134            },
1135        };
1136        assert_eq!(make(CountOp::Eq, 2).count(), 3);
1137        assert_eq!(make(CountOp::NotEq, 2).count(), 5); // 8 − 3
1138        assert_eq!(make(CountOp::Lt, 2).count(), 4); // counts 0, 1 → 1 + 3
1139        assert_eq!(make(CountOp::Le, 2).count(), 7); // 0, 1, 2 → 1 + 3 + 3
1140        assert_eq!(make(CountOp::Gt, 1).count(), 4); // 2, 3 → 3 + 1
1141        assert_eq!(make(CountOp::Ge, 1).count(), 7); // 1, 2, 3 → 3 + 3 + 1
1142    }
1143
1144    #[test]
1145    fn count_constraint_with_inequality_filters() {
1146        // Alphabet {A, B} on 3 squares (Cartesian regime — no Count-Eq).
1147        // Filter to arrangements with at least 1 A and at most 2 A.
1148        let problem = Problem {
1149            num_squares: 3,
1150            square_colors: light_dark(3),
1151            pieces: vec![Tile::A, Tile::B],
1152            constraint: Constraint::And(vec![
1153                Constraint::Count {
1154                    piece: Tile::A,
1155                    op: CountOp::Ge,
1156                    value: 1,
1157                },
1158                Constraint::Count {
1159                    piece: Tile::A,
1160                    op: CountOp::Le,
1161                    value: 2,
1162                },
1163            ]),
1164        };
1165        // 2^3 = 8 total, minus the all-B (0 A's) and all-A (3 A's) extremes.
1166        assert_eq!(problem.count(), 6);
1167    }
1168
1169    #[test]
1170    fn count_constraint_inside_or_does_not_activate_fast_path() {
1171        // Or-wrapped Count-Eq must NOT be picked up as a piece-count
1172        // declaration — the solver should fall back to Cartesian
1173        // enumeration and filter via the Or.
1174        let problem = Problem {
1175            num_squares: 2,
1176            square_colors: light_dark(2),
1177            pieces: vec![Tile::A, Tile::B],
1178            constraint: Constraint::Or(vec![
1179                Constraint::Count {
1180                    piece: Tile::A,
1181                    op: CountOp::Eq,
1182                    value: 2,
1183                },
1184                Constraint::Count {
1185                    piece: Tile::B,
1186                    op: CountOp::Eq,
1187                    value: 2,
1188                },
1189            ]),
1190        };
1191        // Cartesian regime: 2^2 = 4 candidates (AA, AB, BA, BB).
1192        // AA satisfies first arm, BB satisfies second. Count = 2.
1193        assert_eq!(problem.count(), 2);
1194    }
1195
1196    #[test]
1197    fn empty_alphabet_yields_zero_arrangements() {
1198        let problem: Problem<Tile> = Problem {
1199            num_squares: 3,
1200            square_colors: light_dark(3),
1201            pieces: vec![],
1202            constraint: Constraint::And(vec![]),
1203        };
1204        assert_eq!(problem.count(), 0);
1205        assert_eq!(problem.at(0), None);
1206        assert_eq!(problem.sample(0), None);
1207    }
1208
1209    #[test]
1210    fn zero_squares_yields_zero_arrangements() {
1211        let problem: Problem<Tile> = Problem {
1212            num_squares: 0,
1213            square_colors: vec![],
1214            pieces: vec![Tile::A, Tile::B],
1215            constraint: Constraint::And(vec![]),
1216        };
1217        assert_eq!(problem.count(), 0);
1218        assert_eq!(problem.at(0), None);
1219        assert_eq!(problem.sample(0), None);
1220    }
1221
1222    #[test]
1223    fn validate_accepts_consistent_problem() {
1224        let problem = Problem {
1225            num_squares: 3,
1226            square_colors: light_dark(3),
1227            pieces: vec![Tile::A, Tile::B],
1228            constraint: fixed_counts(&[(Tile::A, 1), (Tile::B, 2)]),
1229        };
1230        assert!(problem.validate().is_ok());
1231    }
1232
1233    #[test]
1234    fn validate_accepts_empty_colors() {
1235        // Empty `square_colors` is "no colour partition declared".
1236        let problem: Problem<Tile> = Problem {
1237            num_squares: 3,
1238            square_colors: vec![],
1239            pieces: vec![Tile::A, Tile::B],
1240            constraint: Constraint::And(vec![]),
1241        };
1242        assert!(problem.validate().is_ok());
1243    }
1244
1245    #[test]
1246    fn validate_rejects_mismatched_color_length() {
1247        let problem: Problem<Tile> = Problem {
1248            num_squares: 3,
1249            square_colors: vec![SquareColor::Light], // only 1 of 3
1250            pieces: vec![Tile::A],
1251            constraint: Constraint::And(vec![]),
1252        };
1253        assert_eq!(
1254            problem.validate(),
1255            Err(ValidationError::ColorLengthMismatch {
1256                expected: 3,
1257                actual: 1,
1258            }),
1259        );
1260    }
1261
1262    #[test]
1263    fn validate_rejects_unknown_piece() {
1264        let problem: Problem<Tile> = Problem {
1265            num_squares: 3,
1266            square_colors: light_dark(3),
1267            pieces: vec![Tile::A], // Tile::B not declared
1268            constraint: Constraint::At {
1269                piece: Tile::B,
1270                square: 0,
1271            },
1272        };
1273        assert_eq!(problem.validate(), Err(ValidationError::UnknownPiece));
1274    }
1275
1276    #[test]
1277    fn validate_rejects_unknown_color() {
1278        // square_colors only has Light; constraint references Dark.
1279        let problem: Problem<Tile> = Problem {
1280            num_squares: 2,
1281            square_colors: vec![SquareColor::Light, SquareColor::Light],
1282            pieces: vec![Tile::A],
1283            constraint: Constraint::CountOnColor {
1284                piece: Tile::A,
1285                color: SquareColor::Dark,
1286                op: CountOp::Eq,
1287                value: 0,
1288            },
1289        };
1290        assert_eq!(problem.validate(), Err(ValidationError::UnknownColor));
1291    }
1292
1293    #[test]
1294    fn validate_rejects_instance_out_of_range_in_order() {
1295        // Declare exactly 2 Bs but reference (B, 2).
1296        let problem: Problem<Tile> = Problem {
1297            num_squares: 3,
1298            square_colors: light_dark(3),
1299            pieces: vec![Tile::A, Tile::B],
1300            constraint: Constraint::And(vec![
1301                Constraint::Count {
1302                    piece: Tile::A,
1303                    op: CountOp::Eq,
1304                    value: 1,
1305                },
1306                Constraint::Count {
1307                    piece: Tile::B,
1308                    op: CountOp::Eq,
1309                    value: 2,
1310                },
1311                Constraint::Order(vec![(Tile::B, 0), (Tile::B, 1), (Tile::B, 2)]),
1312            ]),
1313        };
1314        assert_eq!(
1315            problem.validate(),
1316            Err(ValidationError::InstanceOutOfRange {
1317                instance: 2,
1318                declared: 2,
1319            }),
1320        );
1321    }
1322
1323    #[test]
1324    fn validate_rejects_instance_out_of_range_in_relative() {
1325        let problem: Problem<Tile> = Problem {
1326            num_squares: 2,
1327            square_colors: light_dark(2),
1328            pieces: vec![Tile::A, Tile::B],
1329            constraint: Constraint::And(vec![
1330                Constraint::Count {
1331                    piece: Tile::A,
1332                    op: CountOp::Eq,
1333                    value: 1,
1334                },
1335                Constraint::Count {
1336                    piece: Tile::B,
1337                    op: CountOp::Eq,
1338                    value: 1,
1339                },
1340                Constraint::Relative {
1341                    lhs: (Tile::A, 1), // only 1 A declared
1342                    rhs: (Tile::B, 0),
1343                    op: CountOp::Eq,
1344                    offset: 0,
1345                },
1346            ]),
1347        };
1348        assert_eq!(
1349            problem.validate(),
1350            Err(ValidationError::InstanceOutOfRange {
1351                instance: 1,
1352                declared: 1,
1353            }),
1354        );
1355    }
1356
1357    #[test]
1358    fn validate_rejects_count_on_color_when_no_colors_declared() {
1359        // Empty square_colors + a CountOnColor → contradiction.
1360        let problem: Problem<Tile> = Problem {
1361            num_squares: 2,
1362            square_colors: vec![],
1363            pieces: vec![Tile::A],
1364            constraint: Constraint::CountOnColor {
1365                piece: Tile::A,
1366                color: SquareColor::Light,
1367                op: CountOp::Eq,
1368                value: 1,
1369            },
1370        };
1371        assert_eq!(problem.validate(), Err(ValidationError::UnknownColor));
1372    }
1373
1374    #[test]
1375    fn validate_rejects_square_out_of_range() {
1376        let problem: Problem<Tile> = Problem {
1377            num_squares: 3,
1378            square_colors: light_dark(3),
1379            pieces: vec![Tile::A],
1380            constraint: Constraint::At {
1381                piece: Tile::A,
1382                square: 5, // out of range
1383            },
1384        };
1385        assert_eq!(
1386            problem.validate(),
1387            Err(ValidationError::SquareOutOfRange {
1388                square: 5,
1389                num_squares: 3,
1390            }),
1391        );
1392    }
1393
1394    #[test]
1395    fn validate_walks_into_combinators() {
1396        // UnknownPiece inside a nested Or should still be caught.
1397        let problem: Problem<Tile> = Problem {
1398            num_squares: 2,
1399            square_colors: light_dark(2),
1400            pieces: vec![Tile::A],
1401            constraint: Constraint::Or(vec![
1402                Constraint::At {
1403                    piece: Tile::A,
1404                    square: 0,
1405                },
1406                Constraint::Not(Box::new(Constraint::At {
1407                    piece: Tile::B, // unknown
1408                    square: 1,
1409                })),
1410            ]),
1411        };
1412        assert_eq!(problem.validate(), Err(ValidationError::UnknownPiece));
1413    }
1414
1415    #[test]
1416    fn try_build_returns_error_for_invalid_problem() {
1417        let result: Result<Problem<Tile>, _> = Problem::builder()
1418            .squares(3)
1419            .colors(vec![SquareColor::Light]) // mismatched
1420            .pieces([Tile::A])
1421            .try_build();
1422        assert!(matches!(
1423            result,
1424            Err(ValidationError::ColorLengthMismatch { .. })
1425        ));
1426    }
1427
1428    #[test]
1429    fn builder_colors_fn_assigns_per_index() {
1430        let problem: Problem<Tile> = Problem::builder()
1431            .squares(6)
1432            .colors_fn(|i| {
1433                if i < 3 {
1434                    SquareColor::Light
1435                } else {
1436                    SquareColor::Dark
1437                }
1438            })
1439            .pieces([Tile::A])
1440            .build();
1441        assert_eq!(problem.square_colors.len(), 6);
1442        assert_eq!(problem.square_colors[0], SquareColor::Light);
1443        assert_eq!(problem.square_colors[5], SquareColor::Dark);
1444    }
1445
1446    #[test]
1447    fn builder_matches_struct_literal_on_chess_960() {
1448        // Build the same problem two ways and assert they have the
1449        // same population.
1450        let from_preset = crate::chess::chess_960().into_problem();
1451
1452        use crate::chess::{back_rank_colors, Piece};
1453        let from_builder: Problem<Piece> = Problem::builder()
1454            .squares(8)
1455            .colors(back_rank_colors())
1456            .pieces([
1457                Piece::King,
1458                Piece::Queen,
1459                Piece::Rook,
1460                Piece::Bishop,
1461                Piece::Knight,
1462            ])
1463            .constraint(Constraint::Count {
1464                piece: Piece::King,
1465                op: CountOp::Eq,
1466                value: 1,
1467            })
1468            .constraint(Constraint::Count {
1469                piece: Piece::Queen,
1470                op: CountOp::Eq,
1471                value: 1,
1472            })
1473            .constraint(Constraint::Count {
1474                piece: Piece::Rook,
1475                op: CountOp::Eq,
1476                value: 2,
1477            })
1478            .constraint(Constraint::Count {
1479                piece: Piece::Bishop,
1480                op: CountOp::Eq,
1481                value: 2,
1482            })
1483            .constraint(Constraint::Count {
1484                piece: Piece::Knight,
1485                op: CountOp::Eq,
1486                value: 2,
1487            })
1488            .constraint(Constraint::CountOnColor {
1489                piece: Piece::Bishop,
1490                color: SquareColor::Light,
1491                op: CountOp::Eq,
1492                value: 1,
1493            })
1494            .constraint(Constraint::CountOnColor {
1495                piece: Piece::Bishop,
1496                color: SquareColor::Dark,
1497                op: CountOp::Eq,
1498                value: 1,
1499            })
1500            .constraint(Constraint::Order(vec![
1501                (Piece::Rook, 0),
1502                (Piece::King, 0),
1503                (Piece::Rook, 1),
1504            ]))
1505            .build();
1506
1507        assert_eq!(from_builder.count(), from_preset.count());
1508        assert_eq!(from_builder.count(), 960);
1509    }
1510
1511    #[test]
1512    fn builder_with_zero_constraints_is_and_empty() {
1513        let problem: Problem<Tile> = Problem::builder()
1514            .squares(3)
1515            .alternating_colors(SquareColor::Light, SquareColor::Dark)
1516            .pieces([Tile::A, Tile::B])
1517            .build();
1518        // Unconstrained Cartesian regime: 2^3 = 8.
1519        assert_eq!(problem.count(), 8);
1520        assert!(matches!(problem.constraint, Constraint::And(ref v) if v.is_empty()));
1521    }
1522
1523    #[test]
1524    fn builder_with_single_constraint_unwraps_and() {
1525        let problem: Problem<Tile> = Problem::builder()
1526            .squares(3)
1527            .uniform_colors(SquareColor::Light)
1528            .pieces([Tile::A, Tile::B])
1529            .constraint(Constraint::At {
1530                piece: Tile::A,
1531                square: 0,
1532            })
1533            .build();
1534        // A pinned at square 0, 2 choices on each remaining square → 4.
1535        assert_eq!(problem.count(), 4);
1536        // Single constraint is stored bare, not wrapped in And.
1537        assert!(matches!(problem.constraint, Constraint::At { .. }));
1538    }
1539
1540    #[cfg(feature = "serde")]
1541    #[test]
1542    fn problem_serde_roundtrip() {
1543        let problem = Problem {
1544            num_squares: 3,
1545            square_colors: light_dark(3),
1546            pieces: vec![Tile::A, Tile::B],
1547            constraint: fixed_counts(&[(Tile::A, 1), (Tile::B, 2)]),
1548        };
1549        let json = serde_json::to_string(&problem).expect("serialise");
1550        let back: Problem<Tile> = serde_json::from_str(&json).expect("deserialise");
1551        assert_eq!(back.num_squares, problem.num_squares);
1552        assert_eq!(back.pieces, problem.pieces);
1553        assert_eq!(back.count(), problem.count());
1554    }
1555
1556    #[test]
1557    fn builder_piece_method_appends() {
1558        let problem: Problem<Tile> = Problem::builder()
1559            .squares(3)
1560            .uniform_colors(SquareColor::Light)
1561            .piece(Tile::A)
1562            .piece(Tile::B)
1563            .constraint(Constraint::Count {
1564                piece: Tile::A,
1565                op: CountOp::Eq,
1566                value: 2,
1567            })
1568            .constraint(Constraint::Count {
1569                piece: Tile::B,
1570                op: CountOp::Eq,
1571                value: 1,
1572            })
1573            .build();
1574        assert_eq!(problem.count(), 3);
1575    }
1576
1577    #[test]
1578    fn with_constraint_adds_via_and() {
1579        let problem = Problem {
1580            num_squares: 3,
1581            square_colors: light_dark(3),
1582            pieces: vec![Tile::A, Tile::B, Tile::C],
1583            constraint: fixed_counts(&[(Tile::A, 1), (Tile::B, 1), (Tile::C, 1)]),
1584        };
1585        assert_eq!(problem.count(), 6);
1586        let narrowed = problem.with_constraint(Constraint::At {
1587            piece: Tile::A,
1588            square: 0,
1589        });
1590        assert_eq!(narrowed.count(), 2);
1591    }
1592}