sat_solver/sat/
assignment.rs

1#![warn(clippy::all, clippy::pedantic, clippy::nursery, clippy::cargo)]
2#![allow(
3    unsafe_code,
4    clippy::cast_possible_truncation,
5    clippy::cast_possible_wrap
6)]
7
8//! This module defines the `Assignment` trait and its implementations for managing variable states
9//! in the SAT solvers.
10//!
11//! It provides a way to track whether variables are assigned true, false, or remain unassigned.
12//!
13//! It includes two main implementations:
14//! - `VecAssignment`: Uses a `Vec<VarState>` for dense variable sets.
15//! - `HashMapAssignment`: Uses an `FxHashMap<Variable, VarState>` for sparse or non-contiguous variable sets.
16//!
17//! The `Assignment` trait defines methods for setting, resetting, and querying variable states,
18//! as well as retrieving solutions in a format compatible with SAT solver outputs.
19
20use crate::sat::literal::{Literal, Variable};
21use crate::sat::solver::Solutions;
22use clap::ValueEnum;
23use core::ops::{Index, IndexMut};
24use itertools::Itertools;
25use rustc_hash::FxHashMap;
26use std::fmt::{Debug, Display};
27use std::hash::Hash;
28
29/// Represents the assignment state of a propositional variable.
30///
31/// A variable can be unassigned, or assigned to true or false.
32#[derive(Debug, Clone, PartialEq, Eq, Copy, Default, Hash, PartialOrd, Ord)]
33pub enum VarState {
34    /// The variable has not been assigned a truth value.
35    #[default]
36    Unassigned,
37    /// The variable has been assigned a specific truth value.
38    Assigned(bool),
39}
40
41impl VarState {
42    /// Checks if the variable state is `Assigned`.
43    ///
44    /// # Returns
45    ///
46    /// `true` if the variable is assigned (either true or false), `false` otherwise.
47    #[must_use]
48    pub const fn is_assigned(self) -> bool {
49        matches!(self, Self::Assigned(_))
50    }
51
52    /// Checks if the variable state is `Unassigned`.
53    ///
54    /// # Returns
55    ///
56    /// `true` if the variable is unassigned, `false` otherwise.
57    #[must_use]
58    pub const fn is_unassigned(self) -> bool {
59        !self.is_assigned()
60    }
61
62    /// Checks if the variable state is `Assigned(true)`.
63    ///
64    /// # Returns
65    ///
66    /// `true` if the variable is assigned true, `false` otherwise.
67    #[must_use]
68    pub const fn is_true(self) -> bool {
69        matches!(self, Self::Assigned(true))
70    }
71
72    /// Checks if the variable state is `Assigned(false)`.
73    ///
74    /// # Returns
75    ///
76    /// `true` if the variable is assigned false, `false` otherwise.
77    #[must_use]
78    pub const fn is_false(self) -> bool {
79        matches!(self, Self::Assigned(false))
80    }
81}
82
83impl From<VarState> for Option<bool> {
84    /// Converts a `VarState` into an `Option<bool>`.
85    ///
86    /// `VarState::Assigned(b)` converts to `Some(b)`.
87    /// `VarState::Unassigned` converts to `None`.
88    fn from(s: VarState) -> Self {
89        match s {
90            VarState::Assigned(b) => Some(b),
91            VarState::Unassigned => None,
92        }
93    }
94}
95
96impl From<Option<bool>> for VarState {
97    /// Converts an `Option<bool>` into a `VarState`.
98    ///
99    /// `Some(b)` converts to `VarState::Assigned(b)`.
100    /// `None` converts to `VarState::Unassigned`.
101    fn from(b: Option<bool>) -> Self {
102        b.map_or(Self::Unassigned, VarState::Assigned)
103    }
104}
105
106/// Trait defining the interface for managing variable assignments.
107///
108/// This trait allows for different underlying data structures (e.g. `Vec`, `HashMap`)
109/// to store and manage the states of variables in a SAT solver.
110/// Variables are typically represented by `usize` indices for direct access,
111/// or `Variable` (a `u32` alias) for semantic clarity.
112/// Assumes 0-indexed variables if `Variable` is a numeric type used as an index.
113pub trait Assignment:
114    Index<usize, Output = VarState> + IndexMut<usize, Output = VarState> + Debug + Clone
115{
116    /// Creates a new assignment manager for `n_vars` variables.
117    ///
118    /// All variables are initially `Unassigned`.
119    ///
120    /// # Arguments
121    ///
122    /// * `n_vars`: The total number of variables to manage.
123    fn new(n_vars: usize) -> Self;
124
125    /// Returns the total number of variables this assignment manager is configured for.
126    fn num_vars(&self) -> usize;
127
128    /// Sets the truth value of a variable.
129    ///
130    /// # Arguments
131    ///
132    /// * `var`: The variable to assign.
133    /// * `b`: The boolean value to assign (`true` or `false`).
134    fn set(&mut self, var: Variable, b: bool);
135
136    /// Resets all variables to the `Unassigned` state.
137    fn reset(&mut self);
138
139    /// Assigns a truth value to a variable based on a literal.
140    ///
141    /// The variable of the literal is assigned the literal's polarity.
142    /// For example, if `l` is `x_i`, `x_i` is set to `true`. If `l` is `!x_i`, `x_i` is set to `false`.
143    ///
144    /// # Arguments
145    ///
146    /// * `l`: The literal to assign.
147    fn assign(&mut self, l: impl Literal) {
148        self.set(l.variable(), l.polarity());
149    }
150
151    /// Sets a variable to the `Unassigned` state.
152    ///
153    /// # Arguments
154    ///
155    /// * `var`: The variable to unassign.
156    fn unassign(&mut self, var: Variable);
157
158    /// Retrieves the current set of assigned variables as a `Solutions` object.
159    ///
160    /// `Solutions` expects 1-indexed variable IDs for DIMACS compatibility.
161    fn get_solutions(&self) -> Solutions;
162
163    /// Checks if a specific variable is assigned a truth value.
164    ///
165    /// # Arguments
166    ///
167    /// * `var`: The variable to check.
168    ///
169    /// # Returns
170    ///
171    /// `true` if the variable is assigned, `false` otherwise.
172    fn is_assigned(&self, var: Variable) -> bool {
173        self[var as usize].is_assigned()
174    }
175
176    /// Gets the truth value of a variable, if assigned.
177    ///
178    /// # Arguments
179    ///
180    /// * `var`: The variable whose value is requested.
181    ///
182    /// # Returns
183    ///
184    /// `Some(true)` if the variable is assigned true, `Some(false)` if assigned false,
185    /// or `None` if unassigned.
186    fn var_value(&self, var: Variable) -> Option<bool> {
187        self[var as usize].into()
188    }
189
190    /// Checks if all variables managed by this assignment are currently assigned.
191    ///
192    /// # Returns
193    ///
194    /// `true` if all variables have a truth value, `false` otherwise.
195    fn all_assigned(&self) -> bool;
196
197    /// Gets the truth value of a literal under the current assignment.
198    ///
199    /// # Arguments
200    ///
201    /// * `l`: The literal to evaluate.
202    ///
203    /// # Returns
204    ///
205    /// `Some(true)` if the literal is true, `Some(false)` if false,
206    /// or `None` if the literal's variable is unassigned.
207    fn literal_value(&self, l: impl Literal) -> Option<bool> {
208        self.var_value(l.variable()).map(|b| b == l.polarity())
209    }
210
211    /// Returns an iterator over all currently unassigned variables.
212    ///
213    /// The iterator yields `Variable` identifiers.
214    fn unassigned(&self) -> impl Iterator<Item = Variable> + '_ {
215        (0..self.num_vars()).filter_map(move |i| {
216            #[allow(clippy::cast_possible_truncation)]
217            let var = i as Variable;
218            if self[var as usize].is_unassigned() {
219                Some(var)
220            } else {
221                None
222            }
223        })
224    }
225}
226
227/// An implementation of `Assignment` using a `Vec<VarState>`.
228///
229/// This implementation is efficient for dense sets of variables, where variable IDs
230/// are contiguous and start from 0 (i.e. `0, 1, ..., n-1`).
231/// Indexing with a `usize` value greater than or equal to the number of variables
232/// will result in a panic.
233#[derive(Debug, Clone, PartialEq, Eq, Default)]
234pub struct VecAssignment {
235    /// Stores the state of each variable. The index in the vector corresponds to the `Variable` ID.
236    states: Vec<VarState>,
237}
238
239impl Index<usize> for VecAssignment {
240    type Output = VarState;
241
242    /// Accesses the state of the variable at the given `index`.
243    ///
244    /// # Panics
245    ///
246    /// Panics if `index` is out of bounds.
247    /// The `unsafe get_unchecked` is used for performance, assuming valid indices
248    /// from internal logic.
249    fn index(&self, index: usize) -> &Self::Output {
250        // Safety: Caller must ensure index is within bounds [0, states.len()).
251        // This should never fault when used correctly
252        unsafe { self.states.get_unchecked(index) }
253    }
254}
255
256impl IndexMut<usize> for VecAssignment {
257    /// Mutably accesses the state of the variable at the given `index`.
258    ///
259    /// # Panics
260    ///
261    /// Panics if `index` is out of bounds.
262    /// The `unsafe get_unchecked_mut` is used for performance, similar to `index`.
263    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
264        // Safety: Caller must ensure index is within bounds [0, states.len()).
265        // This should never fault when used correctly
266        unsafe { self.states.get_unchecked_mut(index) }
267    }
268}
269
270impl Assignment for VecAssignment {
271    /// Creates a new `VecAssignment` for `n_vars` variables.
272    /// Variables are indexed from `0` to `n_vars - 1`.
273    fn new(n_vars: usize) -> Self {
274        Self {
275            states: vec![VarState::Unassigned; n_vars],
276        }
277    }
278
279    fn num_vars(&self) -> usize {
280        self.states.len()
281    }
282
283    fn set(&mut self, var: Variable, b: bool) {
284        self[var as usize] = VarState::Assigned(b);
285    }
286
287    fn reset(&mut self) {
288        self.states.fill(VarState::Unassigned);
289    }
290
291    fn unassign(&mut self, var: Variable) {
292        self[var as usize] = VarState::Unassigned;
293    }
294
295    fn get_solutions(&self) -> Solutions {
296        Solutions::new(
297            &self
298                .states
299                .iter()
300                .enumerate()
301                .filter_map(|(i, s)| {
302                    #[allow(clippy::cast_possible_wrap, clippy::cast_possible_truncation)]
303                    let var_id = i as i32;
304                    match s {
305                        VarState::Assigned(true) => Some(var_id),
306                        VarState::Assigned(false) => Some(-var_id),
307                        _ => None,
308                    }
309                })
310                .collect_vec(),
311        )
312    }
313
314    fn all_assigned(&self) -> bool {
315        self.states.iter().all(|v| v.is_assigned())
316    }
317}
318
319/// An implementation of `Assignment` using an `FxHashMap<Variable, VarState>`.
320///
321/// This implementation is suitable for sparse sets of variables or when variable IDs
322/// are not necessarily small or contiguous. It stores states only for variables
323/// that have been explicitly set or unassigned. Accessing a variable not in the map
324/// via `Index` will return `VarState::Unassigned`.
325#[derive(Debug, Clone, PartialEq, Eq, Default)]
326pub struct HashMapAssignment {
327    /// Stores variable states. Keys are `Variable` IDs.
328    map: FxHashMap<Variable, VarState>,
329    /// The total number of distinct variables expected or managed by this assignment instance.
330    /// Used for `all_assigned` and iterating `unassigned` variables up to this conceptual limit.
331    num_vars: usize,
332}
333
334impl Index<usize> for HashMapAssignment {
335    type Output = VarState;
336
337    /// Accesses the state of the variable corresponding to `index`.
338    ///
339    /// If the variable (converted from `index`) is not in the map,
340    /// it's treated as `Unassigned`.
341    fn index(&self, index: usize) -> &Self::Output {
342        #[allow(clippy::cast_possible_truncation)]
343        self.map
344            .get(&(index as Variable))
345            .unwrap_or(&VarState::Unassigned)
346    }
347}
348
349impl IndexMut<usize> for HashMapAssignment {
350    /// Mutably accesses the state of the variable corresponding to `index`.
351    ///
352    /// If the variable (converted from `index`) is not in the map, it's inserted
353    /// with `VarState::Unassigned` before returning a mutable reference.
354    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
355        #[allow(clippy::cast_possible_truncation)]
356        self.map
357            .entry(index as Variable)
358            .or_insert(VarState::Unassigned)
359    }
360}
361
362impl Assignment for HashMapAssignment {
363    /// Creates a new `HashMapAssignment`.
364    ///
365    /// # Arguments
366    ///
367    /// * `n_vars`: The conceptual number of variables. This is used by `all_assigned`
368    ///   to determine if all *expected* variables have assignments, and by
369    ///   `unassigned` to iterate up to this many variables.
370    fn new(n_vars: usize) -> Self {
371        Self {
372            map: FxHashMap::default(),
373            num_vars: n_vars,
374        }
375    }
376
377    fn num_vars(&self) -> usize {
378        self.num_vars
379    }
380
381    fn set(&mut self, var: Variable, b: bool) {
382        self.map.insert(var, VarState::Assigned(b));
383    }
384
385    fn reset(&mut self) {
386        self.map.clear();
387    }
388
389    /// Unassigns a variable by setting its state to `VarState::Unassigned` in the map.
390    /// If the variable was not in the map, it will be inserted as `Unassigned`.
391    fn unassign(&mut self, var: Variable) {
392        self.map.insert(var, VarState::Unassigned);
393    }
394
395    fn get_solutions(&self) -> Solutions {
396        Solutions::new(
397            &self
398                .map
399                .iter()
400                .filter_map(|(&var, s)| {
401                    #[allow(clippy::cast_possible_wrap)]
402                    let var_id = var as i32;
403                    match s {
404                        VarState::Assigned(true) => Some(var_id),
405                        VarState::Assigned(false) => Some(-var_id),
406                        _ => None,
407                    }
408                })
409                .collect_vec(),
410        )
411    }
412
413    fn all_assigned(&self) -> bool {
414        self.map.len() == self.num_vars && self.map.values().all(|v| v.is_assigned())
415    }
416}
417
418/// Defines an enumeration of the possible assignment types.
419/// Allows for a rough dynamic dispatch.
420#[derive(Clone, Debug)]
421pub enum AssignmentImpls {
422    /// `Vec` assignment wrapper
423    Vec(VecAssignment),
424
425    /// `HashMap` assignment wrapper
426    HashMap(HashMapAssignment),
427}
428
429impl Index<usize> for AssignmentImpls {
430    type Output = VarState;
431    fn index(&self, index: usize) -> &Self::Output {
432        match self {
433            Self::Vec(v) => v.index(index),
434            Self::HashMap(v) => v.index(index),
435        }
436    }
437}
438
439impl IndexMut<usize> for AssignmentImpls {
440    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
441        match self {
442            Self::Vec(v) => v.index_mut(index),
443            Self::HashMap(v) => v.index_mut(index),
444        }
445    }
446}
447
448impl Assignment for AssignmentImpls {
449    fn new(n_vars: usize) -> Self {
450        Self::Vec(VecAssignment::new(n_vars))
451    }
452
453    fn num_vars(&self) -> usize {
454        match self {
455            Self::Vec(v) => v.num_vars(),
456            Self::HashMap(v) => v.num_vars(),
457        }
458    }
459
460    fn set(&mut self, var: Variable, b: bool) {
461        match self {
462            Self::Vec(v) => v.set(var, b),
463            Self::HashMap(v) => v.set(var, b),
464        }
465    }
466
467    fn reset(&mut self) {
468        match self {
469            Self::Vec(v) => v.reset(),
470            Self::HashMap(v) => v.reset(),
471        }
472    }
473    fn unassign(&mut self, var: Variable) {
474        match self {
475            Self::Vec(v) => v.unassign(var),
476            Self::HashMap(v) => v.unassign(var),
477        }
478    }
479
480    fn get_solutions(&self) -> Solutions {
481        match self {
482            Self::Vec(v) => v.get_solutions(),
483            Self::HashMap(v) => v.get_solutions(),
484        }
485    }
486    fn all_assigned(&self) -> bool {
487        match self {
488            Self::Vec(v) => v.all_assigned(),
489            Self::HashMap(v) => v.all_assigned(),
490        }
491    }
492}
493
494/// An enumeration of the types of assignment implementations available.
495#[derive(Debug, Clone, PartialEq, Eq, Copy, Default, Hash, ValueEnum)]
496pub enum AssignmentType {
497    /// Use a `Vec<VarState>` for dense variable sets.
498    #[default]
499    Vec,
500    /// Use an `FxHashMap<Variable, VarState>` for sparse or non-contiguous variable sets.
501    HashMap,
502}
503
504impl Display for AssignmentType {
505    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
506        match self {
507            Self::Vec => write!(f, "vec"),
508            Self::HashMap => write!(f, "hash-map"),
509        }
510    }
511}
512
513impl AssignmentType {
514    /// Creates a new `Assignment` instance based on the specified type.
515    ///
516    /// # Arguments
517    ///
518    /// * `n_vars`: The number of variables to manage.
519    ///
520    /// # Returns
521    ///
522    /// An instance of `Assignment` based on the specified type.
523    #[must_use]
524    pub fn to_impl(self, n_vars: usize) -> AssignmentImpls {
525        match self {
526            Self::Vec => AssignmentImpls::Vec(VecAssignment::new(n_vars)),
527            Self::HashMap => AssignmentImpls::HashMap(HashMapAssignment::new(n_vars)),
528        }
529    }
530}
531
532#[cfg(test)]
533mod tests {
534    use super::*;
535    use crate::sat::literal::PackedLiteral;
536
537    #[test]
538    fn test_var_state() {
539        assert!(VarState::Unassigned.is_unassigned());
540        assert!(!VarState::Unassigned.is_assigned());
541        assert!(!VarState::Unassigned.is_true());
542        assert!(!VarState::Unassigned.is_false());
543
544        assert!(!VarState::Assigned(true).is_unassigned());
545        assert!(VarState::Assigned(true).is_assigned());
546        assert!(VarState::Assigned(true).is_true());
547        assert!(!VarState::Assigned(true).is_false());
548
549        assert!(!VarState::Assigned(false).is_unassigned());
550        assert!(VarState::Assigned(false).is_assigned());
551        assert!(!VarState::Assigned(false).is_true());
552        assert!(VarState::Assigned(false).is_false());
553    }
554
555    #[allow(clippy::cognitive_complexity)]
556    fn test_assignment<A: Assignment>(a: &mut A) {
557        a.set(1, true);
558        a.set(2, false);
559        a.set(3, true);
560
561        assert!(a.is_assigned(1));
562        assert!(a.is_assigned(2));
563        assert!(a.is_assigned(3));
564        assert!(!a.is_assigned(0));
565
566        assert_eq!(a.var_value(1), Some(true));
567        assert_eq!(a.var_value(2), Some(false));
568        assert_eq!(a.var_value(3), Some(true));
569        assert_eq!(a.var_value(0), None);
570
571        assert_eq!(a.literal_value(PackedLiteral::new(1, false)), Some(false));
572        assert_eq!(a.literal_value(PackedLiteral::new(1, true)), Some(true));
573
574        assert_eq!(a.literal_value(PackedLiteral::new(2, false)), Some(true));
575        assert_eq!(a.literal_value(PackedLiteral::new(2, true)), Some(false));
576
577        assert_eq!(a.literal_value(PackedLiteral::new(3, false)), Some(false));
578        assert_eq!(a.literal_value(PackedLiteral::new(3, true)), Some(true));
579
580        a.unassign(1);
581        assert!(!a.is_assigned(1));
582        assert_eq!(a.var_value(1), None);
583        assert_eq!(a.literal_value(PackedLiteral::new(1, false)), None);
584
585        let s = a.get_solutions();
586        assert_eq!(s, Solutions::new(&[-2, 3]));
587
588        let unassigned_vars = a.unassigned().collect_vec();
589        assert_eq!(unassigned_vars, vec![0, 1]);
590
591        assert!(!a.all_assigned());
592        a.set(0, true);
593        a.set(1, true);
594        assert!(a.all_assigned());
595
596        a.reset();
597        assert!(!a.is_assigned(0));
598        assert!(!a.is_assigned(1));
599        assert!(!a.is_assigned(2));
600        assert!(!a.is_assigned(3));
601        assert!(!a.all_assigned());
602    }
603
604    #[test]
605    fn test_assignment_vec() {
606        let mut a: VecAssignment = Assignment::new(4);
607        test_assignment(&mut a);
608    }
609
610    #[test]
611    fn test_hashmap_assignment() {
612        let mut a = HashMapAssignment::new(4);
613        test_assignment(&mut a);
614    }
615
616    #[test]
617    fn test_assignment_unassigned_default_impl() {
618        let a = VecAssignment::new(4);
619        assert!(!a.is_assigned(0));
620        assert!(!a.is_assigned(1));
621        assert!(!a.is_assigned(2));
622        assert!(!a.is_assigned(3));
623        assert_eq!(a.unassigned().count(), 4);
624        assert!(!a.all_assigned());
625
626        let b = HashMapAssignment::new(4);
627        assert!(!b.is_assigned(0));
628        assert!(!b.is_assigned(1));
629        assert_eq!(b.unassigned().count(), 4);
630        assert!(!b.all_assigned());
631
632        let c = VecAssignment::new(0);
633        assert!(c.all_assigned());
634
635        let d = HashMapAssignment::new(0);
636        assert!(d.all_assigned());
637    }
638}
sat_solver/sat/assignment.rs

sat_solver/sat/
assignment.rs
