stalmarck_sat/core/
formula.rs

1use std::collections::HashMap;
2
3/// Represents a formula in propositional logic
4#[derive(Debug, Clone, Default)]
5pub struct Formula {
6    clauses: Vec<Vec<i32>>,
7
8    /// Formula in implication form
9    implication_form: Option<ImplicationFormula>,
10
11    /// Triplet Representation of formula
12    triplets: Option<TripletFormula>,
13    num_vars: usize,
14}
15
16/// Represents a formula in implication form
17#[derive(Debug, Clone, PartialEq)]
18pub enum ImplicationFormula {
19    /// A variable (positive or negative literal)
20    Var(i32),
21
22    /// Negation of an expression (NOT)
23    Not(Box<ImplicationFormula>),
24
25    /// Implication relation (p → q)
26    Implies(Box<ImplicationFormula>, Box<ImplicationFormula>),
27
28    /// Boolean constants
29    Const(bool),
30}
31
32/// Represents a variable or constant in triplet form
33#[derive(Debug, Clone, PartialEq, Eq, Hash)]
34pub enum TripletVar {
35    /// A variable (positive or negative literal)
36    Var(i32),
37
38    /// Boolean constants
39    Const(bool),
40}
41
42/// Represents a formula in triplet form for Stalmarck's algorithm
43#[derive(Debug, Clone, PartialEq)]
44pub struct TripletFormula {
45    /// Collection of triplets (a, b, c) where each represents a logical relation
46    pub triplets: Vec<(TripletVar, TripletVar, TripletVar)>,
47
48    /// Mapping from bridge variables (b1, b2, etc.) to their corresponding triplet index
49    pub bridge_vars: std::collections::HashMap<TripletVar, usize>,
50
51    /// The next available bridge variable ID
52    pub next_bridge_var: i32,
53
54    /// The root variable representing the entire formula
55    pub root_var: Option<TripletVar>,
56
57    /// Maximum original variable ID (not including bridge variables)
58    pub max_original_var: i32,
59}
60
61impl Formula {
62    /// Create a new empty formula
63    pub fn new() -> Self {
64        Self::default()
65    }
66
67    /// Add a clause to the formula
68    pub fn add_clause(&mut self, literals: Vec<i32>) {
69        // Update the number of variables based on the literals in the clause
70        for &lit in &literals {
71            let var = lit.abs() as usize;
72            if var > self.num_vars {
73                self.num_vars = var;
74            }
75        }
76        self.clauses.push(literals);
77    }
78
79    /// Set the number of variables directly
80    pub fn set_num_variables(&mut self, num_vars: usize) {
81        self.num_vars = num_vars;
82    }
83
84    /// Pre-allocate space for the expected number of clauses
85    pub fn reserve_clauses(&mut self, num_clauses: usize) {
86        self.clauses.reserve(num_clauses);
87    }
88
89    /// Translate to implication form
90    pub fn translate_to_implication_form(&mut self) {
91        if self.clauses.is_empty() {
92            // Empty clause is unsatisfiable (FALSE)
93            self.implication_form = Some(ImplicationFormula::Not(Box::new(
94                ImplicationFormula::Const(true),
95            )));
96            return;
97        }
98
99        let clause_exprs: Vec<ImplicationFormula> = self
100            .clauses
101            .iter()
102            .map(|clause| self.clause_to_implication(clause))
103            .collect();
104
105        // Combine clauses using the rule: A AND B = NOT(A implies NOT B)
106        // Start with the first clause
107        let mut result = clause_exprs[0].clone();
108
109        // Process remaining clauses left to right
110        for clause_expr in clause_exprs.iter().skip(1) {
111            // Apply the transformation: result AND clause_expr = NOT(result implies NOT clause_expr)
112            result = ImplicationFormula::Not(Box::new(ImplicationFormula::Implies(
113                Box::new(result),
114                Box::new(ImplicationFormula::Not(Box::new(clause_expr.clone()))),
115            )));
116        }
117
118        // Remove all NOTs from result
119        self.remove_nots(&mut result);
120
121        self.implication_form = Some(result);
122    }
123
124    /// Helper method to convert a single clause to implication form
125    fn clause_to_implication(&self, clause: &[i32]) -> ImplicationFormula {
126        if clause.is_empty() {
127            // Empty clause is unsatisfiable (FALSE)
128            return ImplicationFormula::Not(Box::new(ImplicationFormula::Const(true)));
129        }
130
131        if clause.len() == 1 {
132            // Single literal clause
133            return ImplicationFormula::Var(clause[0]);
134        }
135
136        // Convert OR clause to implication form
137        // (p ∨ q ∨ r) = (¬p → (q ∨ r)) = (¬p → (¬q → r))
138        // Start with the first literal
139        let mut expr = ImplicationFormula::Var(clause[0]);
140
141        // Process literals left to right
142        for &lit in &clause[1..] {
143            // Create implication structure
144            expr = ImplicationFormula::Implies(
145                Box::new(ImplicationFormula::Not(Box::new(expr))),
146                Box::new(ImplicationFormula::Var(lit)),
147            );
148        }
149
150        expr
151    }
152
153    /// Helper method to remove NOTs from the implication formula
154    fn remove_nots(&self, formula: &mut ImplicationFormula) {
155        // Perform a "take and replace" to modify the enum variant in place
156        let original_node = std::mem::replace(formula, ImplicationFormula::Const(true));
157
158        match original_node {
159            ImplicationFormula::Not(mut sub_expr_box) => {
160                // Recursively call remove_nots on the inner expression
161                self.remove_nots(&mut *sub_expr_box);
162
163                // Replace Not(A) with Implies(A, Const(false))
164                *formula = ImplicationFormula::Implies(
165                    sub_expr_box,
166                    Box::new(ImplicationFormula::Const(false)),
167                );
168            }
169            ImplicationFormula::Implies(mut left_box, mut right_box) => {
170                // Recursively call remove_nots on both children
171                self.remove_nots(&mut *left_box);
172                self.remove_nots(&mut *right_box);
173
174                // Reconstruct the Implies node with its children
175                *formula = ImplicationFormula::Implies(left_box, right_box);
176            }
177            ImplicationFormula::Var(id) => {
178                if id < 0 {
179                    // Transform Var(-k) into Implies(Var(k), Const(false)).
180                    let positive_id = id.abs();
181                    *formula = ImplicationFormula::Implies(
182                        Box::new(ImplicationFormula::Var(positive_id)),
183                        Box::new(ImplicationFormula::Const(false)),
184                    );
185                    // The children of this new Implies node are Var(positive_id) and Const(false).
186                    // These are already "clean" in terms of Not nodes or negative Vars,
187                    // so no further recursive call on `*formula` itself is needed here.
188                } else {
189                    // Positive literal, restore it as is.
190                    *formula = ImplicationFormula::Var(id);
191                }
192            }
193            ImplicationFormula::Const(constant) => {
194                // Constants are left as is, restore it.
195                *formula = ImplicationFormula::Const(constant);
196            }
197        }
198    }
199
200    /// Get the stored implication form
201    pub fn get_implication_form(&self) -> Option<&ImplicationFormula> {
202        self.implication_form.as_ref()
203    }
204
205    /// Encode the formula into triplets
206    pub fn encode_formula_to_triplets(&mut self) {
207        let root_implication_node = match self.implication_form.as_ref() {
208            Some(form) => form,
209            None => {
210                // If there's no implication form (e.g., formula was empty and translate wasn't called,
211                // or explicitly set to None), then there are no triplets.
212                self.triplets = None;
213                return;
214            }
215        };
216
217        // Create a new TripletFormula instance
218        let mut triplet_formula = TripletFormula {
219            triplets: Vec::new(),
220            bridge_vars: HashMap::new(),
221            next_bridge_var: self.num_vars as i32 + 1000,
222            root_var: None,
223            max_original_var: self.num_vars as i32,
224        };
225
226        // Map to store TripletVar result for each processed ImplicationFormula node
227        let mut processed_results: HashMap<*const ImplicationFormula, TripletVar> = HashMap::new();
228
229        // Stack for iterative post-order traversal
230        let mut traversal_stack: Vec<&ImplicationFormula> = Vec::new();
231
232        // Start traversal from the root of the implication form
233        traversal_stack.push(root_implication_node);
234
235        while let Some(current_node) = traversal_stack.last().cloned() {
236            let current_node_ptr = current_node as *const ImplicationFormula;
237
238            // If result for this node has already been processed, skip it
239            if processed_results.contains_key(&current_node_ptr)
240                && !matches!(current_node, ImplicationFormula::Implies(_, _))
241            {
242                traversal_stack.pop();
243                continue;
244            }
245
246            match current_node {
247                ImplicationFormula::Var(id) => {
248                    // Leaf Node: process immediately
249                    traversal_stack.pop();
250                    let result = TripletVar::Var(*id);
251                    processed_results.insert(current_node_ptr, result);
252                }
253                ImplicationFormula::Const(value) => {
254                    // Leaf Node: process immediately
255                    traversal_stack.pop();
256                    let result = TripletVar::Const(*value);
257                    processed_results.insert(current_node_ptr, result);
258                }
259                ImplicationFormula::Implies(left_expr, right_expr) => {
260                    let left_child = &**left_expr;
261                    let right_child = &**right_expr;
262
263                    let left_child_ptr = left_child as *const ImplicationFormula;
264                    let right_child_ptr = right_child as *const ImplicationFormula;
265
266                    let left_processed = processed_results.contains_key(&left_child_ptr);
267                    let right_processed = processed_results.contains_key(&right_child_ptr);
268
269                    if left_processed && right_processed {
270                        // Both children have been processed, process Implies node
271                        traversal_stack.pop();
272
273                        let var_a = processed_results
274                            .get(&left_child_ptr)
275                            .expect("Left child should be processed")
276                            .clone();
277
278                        let var_b = processed_results
279                            .get(&right_child_ptr)
280                            .expect("Right child should be processed")
281                            .clone();
282
283                        // Create a new bridge variable for this implication
284                        let bridge_id = triplet_formula.next_bridge_var;
285                        triplet_formula.next_bridge_var += 1;
286                        let bridge_var = TripletVar::Var(bridge_id);
287
288                        // Store the bridge variable in the map
289                        triplet_formula
290                            .triplets
291                            .push((bridge_var.clone(), var_a, var_b));
292                        triplet_formula
293                            .bridge_vars
294                            .insert(bridge_var.clone(), triplet_formula.triplets.len() - 1);
295                        processed_results.insert(current_node_ptr, bridge_var);
296                    } else {
297                        // Children are not processed
298                        if !left_processed {
299                            // Push left child onto the stack for later processing
300                            traversal_stack.push(left_child);
301                        }
302
303                        if !right_processed {
304                            // Push right child onto the stack for later processing
305                            traversal_stack.push(right_child);
306                        }
307
308                        // Current Implies node remains on the stack to be revisited
309                    }
310                }
311                ImplicationFormula::Not(_) => {
312                    // Should not happen in the final implication form
313                    unreachable!("Not nodes should have been removed in the implication form");
314                }
315            }
316        }
317
318        // The result for the root of the entire implication form is its TripletVar representation
319        triplet_formula.root_var = processed_results
320            .get(&(root_implication_node as *const ImplicationFormula))
321            .cloned();
322
323        self.triplets = Some(triplet_formula);
324    }
325
326    /// Get the triplets representation
327    pub fn get_triplets(&self) -> Option<&TripletFormula> {
328        self.triplets.as_ref()
329    }
330
331    /// Get the number of variables in the formula
332    pub fn num_variables(&self) -> usize {
333        self.num_vars
334    }
335
336    /// Get the number of clauses in the formula
337    pub fn num_clauses(&self) -> usize {
338        self.clauses.len()
339    }
340
341    /// Get the clauses
342    pub fn get_clauses(&self) -> &[Vec<i32>] {
343        &self.clauses
344    }
345}