xlog_prob/
cnf.rs

1//! CNF emission for PIR via Tseitin encoding (DIMACS).
2
3use std::collections::{BTreeMap, HashMap, HashSet};
4
5use xlog_core::{Result, XlogError};
6
7use crate::pir::{ChoiceVarId, LeafId, PirGraph, PirNode, PirNodeId};
8
9/// A CNF formula in DIMACS-style representation.
10#[derive(Debug, Clone, Default)]
11pub struct CnfFormula {
12    num_vars: u32,
13    clauses: Vec<Vec<i32>>,
14}
15
16impl CnfFormula {
17    /// Returns the highest DIMACS variable identifier used by the formula.
18    pub fn num_vars(&self) -> u32 {
19        self.num_vars
20    }
21
22    /// Returns the formula's clauses as signed DIMACS literals.
23    pub fn clauses(&self) -> &[Vec<i32>] {
24        &self.clauses
25    }
26
27    /// Serializes the formula in deterministic DIMACS CNF format.
28    pub fn to_dimacs(&self) -> String {
29        let mut out = String::new();
30        out.push_str("c xlog-prob cnf\n");
31        out.push_str(&format!("p cnf {} {}\n", self.num_vars, self.clauses.len()));
32        for clause in &self.clauses {
33            for lit in clause {
34                out.push_str(&format!("{} ", lit));
35            }
36            out.push_str("0\n");
37        }
38        out
39    }
40}
41
42#[derive(Debug, Clone)]
43/// A Tseitin encoding together with maps from PIR identities to CNF variables.
44pub struct CnfEncoding {
45    /// The encoded conjunctive-normal-form formula.
46    pub cnf: CnfFormula,
47    /// CNF variable assigned to each encoded PIR node.
48    pub node_var: BTreeMap<PirNodeId, u32>,
49    /// CNF variable assigned to each probabilistic leaf.
50    pub leaf_var: BTreeMap<LeafId, u32>,
51    /// CNF variable assigned to each exclusive-choice variable.
52    pub choice_var: BTreeMap<ChoiceVarId, u32>,
53}
54
55/// FNV-1a 64-bit — deterministic, process-independent.
56fn fnv1a(bytes: &[u8]) -> u64 {
57    const OFFSET: u64 = 0xcbf29ce484222325;
58    const PRIME: u64 = 0x100000001b3;
59    let mut h = OFFSET;
60    for &b in bytes {
61        h ^= b as u64;
62        h = h.wrapping_mul(PRIME);
63    }
64    h
65}
66
67/// Compute a canonical hash for each reachable PIR node.
68/// The hash depends only on node structure (variant + leaf/choice IDs + children's
69/// canonical hashes), not on PirNodeId numeric values.
70fn canonical_hashes(pir: &PirGraph, levels: &[Vec<PirNodeId>]) -> HashMap<PirNodeId, u64> {
71    let mut canon: HashMap<PirNodeId, u64> = HashMap::new();
72    for level in levels {
73        for &id in level {
74            let node = pir.node(id).unwrap();
75            let h = match node {
76                PirNode::Const(b) => fnv1a(&[0, *b as u8]),
77                PirNode::Lit { leaf } => {
78                    let mut buf = [0u8; 5];
79                    buf[0] = 1;
80                    buf[1..5].copy_from_slice(&leaf.as_u32().to_le_bytes());
81                    fnv1a(&buf)
82                }
83                PirNode::NegLit { leaf } => {
84                    let mut buf = [0u8; 5];
85                    buf[0] = 2;
86                    buf[1..5].copy_from_slice(&leaf.as_u32().to_le_bytes());
87                    fnv1a(&buf)
88                }
89                PirNode::And { children } => {
90                    let mut child_h: Vec<u64> = children.iter().map(|c| canon[c]).collect();
91                    child_h.sort();
92                    let mut buf = vec![3u8];
93                    for h in child_h {
94                        buf.extend_from_slice(&h.to_le_bytes());
95                    }
96                    fnv1a(&buf)
97                }
98                PirNode::Or { children } => {
99                    let mut child_h: Vec<u64> = children.iter().map(|c| canon[c]).collect();
100                    child_h.sort();
101                    let mut buf = vec![4u8];
102                    for h in child_h {
103                        buf.extend_from_slice(&h.to_le_bytes());
104                    }
105                    fnv1a(&buf)
106                }
107                PirNode::Decision {
108                    var,
109                    child_false,
110                    child_true,
111                } => {
112                    let mut buf = vec![5u8];
113                    buf.extend_from_slice(&var.as_u32().to_le_bytes());
114                    buf.extend_from_slice(&canon[child_false].to_le_bytes());
115                    buf.extend_from_slice(&canon[child_true].to_le_bytes());
116                    fnv1a(&buf)
117                }
118            };
119            canon.insert(id, h);
120        }
121    }
122    canon
123}
124
125/// Compute a process-independent canonical hash of the PIR structure.
126///
127/// This hash depends only on the semantic structure of the PIR graph (node types,
128/// leaf/choice IDs, children's structural hashes), NOT on `PirNodeId` numeric
129/// values. Two processes that build the same XLOG program will produce the same
130/// canonical hash even if `HashMap`-based interning assigns different node IDs.
131///
132/// Used as the `cnf_hash` component of the disk cache key to enable cross-process
133/// cache hits.
134pub fn canonical_pir_hash(pir: &PirGraph, roots: &[PirNodeId]) -> Result<u64> {
135    if roots.is_empty() {
136        return Err(XlogError::Compilation(
137            "Cannot compute canonical hash for empty PIR root set".to_string(),
138        ));
139    }
140    let levels = pir.levelize(roots)?;
141    let canon = canonical_hashes(pir, &levels);
142
143    // Combine root hashes in a deterministic order (sorted by canonical hash).
144    let mut root_hashes: Vec<u64> = roots
145        .iter()
146        .map(|r| canon.get(r).copied().unwrap_or(0))
147        .collect();
148    root_hashes.sort();
149
150    let mut buf = Vec::with_capacity(1 + root_hashes.len() * 8);
151    buf.push(0xFFu8); // tag to distinguish from node hashes
152    for h in root_hashes {
153        buf.extend_from_slice(&h.to_le_bytes());
154    }
155    Ok(fnv1a(&buf))
156}
157
158/// Encodes the subgraph reachable from `roots` as a deterministic Tseitin CNF.
159pub fn encode_cnf(pir: &PirGraph, roots: &[PirNodeId]) -> Result<CnfEncoding> {
160    if roots.is_empty() {
161        return Err(XlogError::Compilation(
162            "Cannot encode CNF for empty PIR root set".to_string(),
163        ));
164    }
165
166    let mut visited: HashSet<PirNodeId> = HashSet::new();
167    let mut leaf_ids: HashSet<LeafId> = HashSet::new();
168    let mut choice_ids: HashSet<ChoiceVarId> = HashSet::new();
169
170    let mut stack: Vec<PirNodeId> = roots.to_vec();
171    while let Some(node_id) = stack.pop() {
172        if !visited.insert(node_id) {
173            continue;
174        }
175
176        let node = pir.node(node_id).ok_or_else(|| {
177            XlogError::Compilation(format!(
178                "Invalid PIR node id while encoding CNF: {:?}",
179                node_id
180            ))
181        })?;
182
183        match node {
184            PirNode::Const(_) => {}
185            PirNode::Lit { leaf } | PirNode::NegLit { leaf } => {
186                leaf_ids.insert(*leaf);
187            }
188            PirNode::And { children } | PirNode::Or { children } => {
189                stack.extend(children.iter().copied());
190            }
191            PirNode::Decision {
192                var,
193                child_false,
194                child_true,
195            } => {
196                choice_ids.insert(*var);
197                stack.push(*child_false);
198                stack.push(*child_true);
199            }
200        }
201    }
202
203    let mut leaf_list: Vec<LeafId> = leaf_ids.into_iter().collect();
204    leaf_list.sort();
205    let mut choice_list: Vec<ChoiceVarId> = choice_ids.into_iter().collect();
206    choice_list.sort();
207
208    let mut next_var: u32 = 1;
209    let mut leaf_var: BTreeMap<LeafId, u32> = BTreeMap::new();
210    for leaf in leaf_list {
211        leaf_var.insert(leaf, next_var);
212        next_var += 1;
213    }
214
215    let mut choice_var: BTreeMap<ChoiceVarId, u32> = BTreeMap::new();
216    for choice in choice_list {
217        choice_var.insert(choice, next_var);
218        next_var += 1;
219    }
220
221    let mut node_ids: Vec<PirNodeId> = visited.into_iter().collect();
222    node_ids.sort();
223
224    let mut node_var: BTreeMap<PirNodeId, u32> = BTreeMap::new();
225    for node_id in node_ids {
226        let node = pir.node(node_id).ok_or_else(|| {
227            XlogError::Compilation(format!(
228                "Invalid PIR node id while encoding CNF: {:?}",
229                node_id
230            ))
231        })?;
232
233        let var_id = match node {
234            PirNode::Lit { leaf } => *leaf_var.get(leaf).ok_or_else(|| {
235                XlogError::Compilation(format!(
236                    "Missing CNF var for PIR leaf {:?} referenced by node {:?}",
237                    leaf, node_id
238                ))
239            })?,
240            PirNode::NegLit { .. } => {
241                // NegLit gets its own variable, which will be constrained to !leaf_var
242                let v = next_var;
243                next_var += 1;
244                v
245            }
246            _ => {
247                let v = next_var;
248                next_var += 1;
249                v
250            }
251        };
252
253        node_var.insert(node_id, var_id);
254    }
255
256    let num_vars = next_var - 1;
257    let mut clauses: Vec<Vec<i32>> = Vec::new();
258
259    let levels = pir.levelize(roots)?;
260    for level in &levels {
261        for &node_id in level {
262            let node = pir.node(node_id).ok_or_else(|| {
263                XlogError::Compilation(format!(
264                    "Invalid PIR node id while emitting CNF clauses: {:?}",
265                    node_id
266                ))
267            })?;
268            let v = *node_var.get(&node_id).ok_or_else(|| {
269                XlogError::Compilation(format!(
270                    "Missing CNF var for PIR node {:?} while emitting clauses",
271                    node_id
272                ))
273            })? as i32;
274
275            match node {
276                PirNode::Const(true) => clauses.push(vec![v]),
277                PirNode::Const(false) => clauses.push(vec![-v]),
278                PirNode::Lit { .. } => {}
279                PirNode::NegLit { leaf } => {
280                    // NegLit uses opposite polarity: node_var <-> !leaf_var
281                    let leaf_v = *leaf_var.get(leaf).ok_or_else(|| {
282                        XlogError::Compilation(format!(
283                            "Missing CNF var for NegLit leaf {:?} at node {:?}",
284                            leaf, node_id
285                        ))
286                    })? as i32;
287                    // v <-> !leaf_v  means:  (v | leaf_v) & (!v | !leaf_v)
288                    clauses.push(vec![v, leaf_v]); // v=false -> leaf_v=true (contrapositive: !leaf_v -> v)
289                    clauses.push(vec![-v, -leaf_v]); // v=true -> leaf_v=false
290                }
291                PirNode::And { children } => {
292                    if children.is_empty() {
293                        clauses.push(vec![v]);
294                        continue;
295                    }
296                    for &child in children {
297                        let c = *node_var.get(&child).ok_or_else(|| {
298                            XlogError::Compilation(format!(
299                                "Missing CNF var for AND child {:?} of {:?}",
300                                child, node_id
301                            ))
302                        })? as i32;
303                        clauses.push(vec![-v, c]);
304                    }
305                    let mut clause = Vec::with_capacity(children.len() + 1);
306                    for &child in children {
307                        let c = *node_var.get(&child).ok_or_else(|| {
308                            XlogError::Compilation(format!(
309                                "Missing CNF var for AND child {:?} of {:?}",
310                                child, node_id
311                            ))
312                        })? as i32;
313                        clause.push(-c);
314                    }
315                    clause.push(v);
316                    clauses.push(clause);
317                }
318                PirNode::Or { children } => {
319                    if children.is_empty() {
320                        clauses.push(vec![-v]);
321                        continue;
322                    }
323                    for &child in children {
324                        let c = *node_var.get(&child).ok_or_else(|| {
325                            XlogError::Compilation(format!(
326                                "Missing CNF var for OR child {:?} of {:?}",
327                                child, node_id
328                            ))
329                        })? as i32;
330                        clauses.push(vec![-c, v]);
331                    }
332                    let mut clause = Vec::with_capacity(children.len() + 1);
333                    clause.push(-v);
334                    for &child in children {
335                        let c = *node_var.get(&child).ok_or_else(|| {
336                            XlogError::Compilation(format!(
337                                "Missing CNF var for OR child {:?} of {:?}",
338                                child, node_id
339                            ))
340                        })? as i32;
341                        clause.push(c);
342                    }
343                    clauses.push(clause);
344                }
345                PirNode::Decision {
346                    var,
347                    child_false,
348                    child_true,
349                } => {
350                    let x = *choice_var.get(var).ok_or_else(|| {
351                        XlogError::Compilation(format!(
352                            "Missing CNF var for decision variable {:?} at node {:?}",
353                            var, node_id
354                        ))
355                    })? as i32;
356
357                    let f = *node_var.get(child_false).ok_or_else(|| {
358                        XlogError::Compilation(format!(
359                            "Missing CNF var for decision false child {:?} at node {:?}",
360                            child_false, node_id
361                        ))
362                    })? as i32;
363                    let t = *node_var.get(child_true).ok_or_else(|| {
364                        XlogError::Compilation(format!(
365                            "Missing CNF var for decision true child {:?} at node {:?}",
366                            child_true, node_id
367                        ))
368                    })? as i32;
369
370                    // v <-> (x ? t : f)
371                    clauses.push(vec![-x, -t, v]); // (x & t) -> v
372                    clauses.push(vec![x, -f, v]); // (!x & f) -> v
373                    clauses.push(vec![-x, t, -v]); // (v & x) -> t
374                    clauses.push(vec![x, f, -v]); // (v & !x) -> f
375                }
376            }
377        }
378    }
379
380    Ok(CnfEncoding {
381        cnf: CnfFormula { num_vars, clauses },
382        node_var,
383        leaf_var,
384        choice_var,
385    })
386}
387
388#[cfg(test)]
389mod tests {
390    use super::*;
391    use crate::pir::{ChoiceVarId, LeafId, PirGraph};
392
393    fn is_sat_with_unit_clauses(cnf: &CnfFormula, units: &[i32]) -> bool {
394        let num_vars = cnf.num_vars() as usize;
395        assert!(
396            num_vars <= 20,
397            "test sat checker only supports small CNFs (vars={})",
398            num_vars
399        );
400
401        let mut clauses: Vec<&[i32]> = cnf.clauses().iter().map(|c| c.as_slice()).collect();
402        let unit_clauses: Vec<Vec<i32>> = units.iter().map(|&u| vec![u]).collect();
403        for uc in &unit_clauses {
404            clauses.push(uc.as_slice());
405        }
406
407        'assign: for mask in 0..(1u64 << num_vars) {
408            for clause in &clauses {
409                let mut clause_sat = false;
410                for &lit in *clause {
411                    let var = lit.unsigned_abs() as usize;
412                    assert!(var >= 1 && var <= num_vars);
413                    let bit = (mask >> (var - 1)) & 1;
414                    let val = bit == 1;
415                    let lit_sat = if lit > 0 { val } else { !val };
416                    if lit_sat {
417                        clause_sat = true;
418                        break;
419                    }
420                }
421                if !clause_sat {
422                    continue 'assign;
423                }
424            }
425            return true;
426        }
427        false
428    }
429
430    #[test]
431    fn test_encode_cnf_does_not_force_root_assignment() {
432        let mut pir = PirGraph::new();
433        let a = pir.lit(LeafId::new(0));
434
435        let encoding = encode_cnf(&pir, &[a]).unwrap();
436        let var_a = *encoding.leaf_var.get(&LeafId::new(0)).unwrap() as i32;
437
438        assert!(is_sat_with_unit_clauses(&encoding.cnf, &[-var_a]));
439    }
440
441    #[test]
442    fn test_tseitin_and_requires_both_children() {
443        let mut pir = PirGraph::new();
444        let a = pir.lit(LeafId::new(0));
445        let b = pir.lit(LeafId::new(1));
446        let root = pir.and(vec![a, b]);
447
448        let encoding = encode_cnf(&pir, &[root]).unwrap();
449        let var_root = *encoding.node_var.get(&root).unwrap() as i32;
450        let var_a = *encoding.leaf_var.get(&LeafId::new(0)).unwrap() as i32;
451        let var_b = *encoding.leaf_var.get(&LeafId::new(1)).unwrap() as i32;
452
453        assert!(is_sat_with_unit_clauses(
454            &encoding.cnf,
455            &[var_root, var_a, var_b]
456        ));
457        assert!(!is_sat_with_unit_clauses(
458            &encoding.cnf,
459            &[var_root, var_a, -var_b]
460        ));
461        assert!(!is_sat_with_unit_clauses(
462            &encoding.cnf,
463            &[var_root, -var_a, var_b]
464        ));
465        assert!(!is_sat_with_unit_clauses(
466            &encoding.cnf,
467            &[var_root, -var_a, -var_b]
468        ));
469    }
470
471    #[test]
472    fn test_tseitin_or_requires_one_child() {
473        let mut pir = PirGraph::new();
474        let a = pir.lit(LeafId::new(0));
475        let b = pir.lit(LeafId::new(1));
476        let root = pir.or(vec![a, b]);
477
478        let encoding = encode_cnf(&pir, &[root]).unwrap();
479        let var_root = *encoding.node_var.get(&root).unwrap() as i32;
480        let var_a = *encoding.leaf_var.get(&LeafId::new(0)).unwrap() as i32;
481        let var_b = *encoding.leaf_var.get(&LeafId::new(1)).unwrap() as i32;
482
483        assert!(is_sat_with_unit_clauses(
484            &encoding.cnf,
485            &[var_root, var_a, var_b]
486        ));
487        assert!(is_sat_with_unit_clauses(
488            &encoding.cnf,
489            &[var_root, var_a, -var_b]
490        ));
491        assert!(is_sat_with_unit_clauses(
492            &encoding.cnf,
493            &[var_root, -var_a, var_b]
494        ));
495        assert!(!is_sat_with_unit_clauses(
496            &encoding.cnf,
497            &[var_root, -var_a, -var_b]
498        ));
499    }
500
501    #[test]
502    fn test_tseitin_decision_mux_matches_choice_var() {
503        let mut pir = PirGraph::new();
504        let t = pir.const_true();
505        let f = pir.const_false();
506        let root = pir.decision(ChoiceVarId::new(0), f, t);
507
508        let encoding = encode_cnf(&pir, &[root]).unwrap();
509        let var_root = *encoding.node_var.get(&root).unwrap() as i32;
510        let x = *encoding.choice_var.get(&ChoiceVarId::new(0)).unwrap() as i32;
511
512        assert!(is_sat_with_unit_clauses(&encoding.cnf, &[var_root, x]));
513        assert!(!is_sat_with_unit_clauses(&encoding.cnf, &[var_root, -x]));
514    }
515
516    #[test]
517    fn test_dimacs_is_well_formed() {
518        let mut pir = PirGraph::new();
519        let a = pir.lit(LeafId::new(0));
520        let root = pir.or(vec![a]);
521
522        let encoding = encode_cnf(&pir, &[root]).unwrap();
523        let dimacs = encoding.cnf.to_dimacs();
524
525        assert!(dimacs.contains("\np cnf "));
526        assert!(dimacs.lines().any(|l| l.ends_with('0')));
527    }
528
529    #[test]
530    fn test_tseitin_neg_lit_uses_negated_polarity() {
531        let mut pir = PirGraph::new();
532        let a = pir.neg_lit(LeafId::new(0)); // NegLit instead of Lit
533        let root = pir.or(vec![a]);
534
535        let encoding = encode_cnf(&pir, &[root]).unwrap();
536        let var_root = *encoding.node_var.get(&root).unwrap() as i32;
537        let var_a = *encoding.leaf_var.get(&LeafId::new(0)).unwrap() as i32;
538
539        // When NegLit node is true, the underlying leaf var should be FALSE
540        // (because NegLit represents "not leaf")
541        // So if root is true and NegLit(a) makes it true, leaf a must be false
542        assert!(
543            is_sat_with_unit_clauses(&encoding.cnf, &[var_root, -var_a]),
544            "root=true with leaf=false should be SAT"
545        );
546
547        // If leaf is true, then NegLit(leaf) is false, so root (which is Or of NegLit) should be false
548        // Actually wait - the Or node can be true or false. Let's think more carefully.
549        //
550        // The CNF encodes: root <-> Or(neg_lit_node)
551        //                  neg_lit_node <-> !leaf
552        // So: root=true requires neg_lit_node=true requires leaf=false
553        assert!(
554            !is_sat_with_unit_clauses(&encoding.cnf, &[var_root, var_a]),
555            "root=true with leaf=true should be UNSAT (NegLit of true leaf is false)"
556        );
557    }
558}
xlog_prob/cnf.rs

xlog_prob/
cnf.rs
