xlog_prob/
pir.rs

1//! Provenance IR (PIR) for probabilistic compilation.
2
3use xlog_core::{Result, XlogError};
4
5#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
6/// Stable index of a node in a [`PirGraph`].
7pub struct PirNodeId(u32);
8
9impl PirNodeId {
10    /// Constructs an identifier from its serialized integer representation.
11    pub fn from_u32(id: u32) -> Self {
12        Self(id)
13    }
14
15    /// Returns the serialized integer representation.
16    pub fn as_u32(self) -> u32 {
17        self.0
18    }
19}
20
21#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
22/// Identifier of an independent probabilistic leaf.
23pub struct LeafId(u32);
24
25impl LeafId {
26    /// Constructs a leaf identifier.
27    pub fn new(id: u32) -> Self {
28        Self(id)
29    }
30
31    /// Returns the underlying integer identifier.
32    pub fn as_u32(self) -> u32 {
33        self.0
34    }
35}
36
37#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
38/// Identifier of a mutually exclusive probabilistic choice.
39pub struct ChoiceVarId(u32);
40
41impl ChoiceVarId {
42    /// Constructs a choice-variable identifier.
43    pub fn new(id: u32) -> Self {
44        Self(id)
45    }
46
47    /// Returns the underlying integer identifier.
48    pub fn as_u32(self) -> u32 {
49        self.0
50    }
51}
52
53#[derive(Debug, Clone, PartialEq)]
54/// A node in the probabilistic intermediate representation.
55pub enum PirNode {
56    /// Boolean constant.
57    Const(bool),
58    /// Positive probabilistic literal.
59    Lit {
60        /// Referenced probabilistic leaf.
61        leaf: LeafId,
62    },
63    /// Negated probabilistic literal with complementary weight.
64    NegLit {
65        /// Referenced probabilistic leaf.
66        leaf: LeafId,
67    },
68    /// Conjunction of child nodes.
69    And {
70        /// Conjunct node identifiers.
71        children: Vec<PirNodeId>,
72    },
73    /// Disjunction of child nodes.
74    Or {
75        /// Disjunct node identifiers.
76        children: Vec<PirNodeId>,
77    },
78    /// Binary branch on an exclusive-choice variable.
79    Decision {
80        /// Variable controlling the branch.
81        var: ChoiceVarId,
82        /// Child selected when the variable is false.
83        child_false: PirNodeId,
84        /// Child selected when the variable is true.
85        child_true: PirNodeId,
86    },
87}
88
89#[derive(Debug, Default, Clone)]
90/// Arena-backed directed acyclic graph of probabilistic expressions.
91pub struct PirGraph {
92    nodes: Vec<PirNode>,
93}
94
95impl PirGraph {
96    /// Creates an empty graph.
97    pub fn new() -> Self {
98        Self { nodes: Vec::new() }
99    }
100
101    /// Returns the number of nodes in the graph.
102    pub fn len(&self) -> usize {
103        self.nodes.len()
104    }
105
106    /// Returns whether the graph contains no nodes.
107    pub fn is_empty(&self) -> bool {
108        self.nodes.is_empty()
109    }
110
111    /// Returns all nodes in identifier order.
112    pub fn nodes(&self) -> &[PirNode] {
113        &self.nodes
114    }
115
116    pub(crate) fn nodes_mut(&mut self) -> &mut [PirNode] {
117        &mut self.nodes
118    }
119
120    /// Returns a node by identifier, or `None` when the identifier is invalid.
121    pub fn node(&self, id: PirNodeId) -> Option<&PirNode> {
122        self.nodes.get(id.0 as usize)
123    }
124
125    fn push_node(&mut self, node: PirNode) -> PirNodeId {
126        let id = PirNodeId(u32::try_from(self.nodes.len()).expect("PIR node id overflow"));
127        self.nodes.push(node);
128        id
129    }
130
131    /// Appends a true constant and returns its identifier.
132    pub fn const_true(&mut self) -> PirNodeId {
133        self.push_node(PirNode::Const(true))
134    }
135
136    /// Appends a false constant and returns its identifier.
137    pub fn const_false(&mut self) -> PirNodeId {
138        self.push_node(PirNode::Const(false))
139    }
140
141    /// Appends a positive leaf literal and returns its node identifier.
142    pub fn lit(&mut self, leaf: LeafId) -> PirNodeId {
143        self.push_node(PirNode::Lit { leaf })
144    }
145
146    /// Appends a negated leaf literal and returns its node identifier.
147    pub fn neg_lit(&mut self, leaf: LeafId) -> PirNodeId {
148        self.push_node(PirNode::NegLit { leaf })
149    }
150
151    /// Appends a conjunction and returns its node identifier.
152    pub fn and(&mut self, children: Vec<PirNodeId>) -> PirNodeId {
153        self.push_node(PirNode::And { children })
154    }
155
156    /// Appends a disjunction and returns its node identifier.
157    pub fn or(&mut self, children: Vec<PirNodeId>) -> PirNodeId {
158        self.push_node(PirNode::Or { children })
159    }
160
161    /// Appends a binary choice node and returns its identifier.
162    pub fn decision(
163        &mut self,
164        var: ChoiceVarId,
165        child_false: PirNodeId,
166        child_true: PirNodeId,
167    ) -> PirNodeId {
168        self.push_node(PirNode::Decision {
169            var,
170            child_false,
171            child_true,
172        })
173    }
174
175    /// Groups nodes reachable from `roots` into dependency-first evaluation levels.
176    pub fn levelize(&self, roots: &[PirNodeId]) -> Result<Vec<Vec<PirNodeId>>> {
177        use std::collections::{HashMap, HashSet};
178
179        let mut visiting: HashSet<PirNodeId> = HashSet::new();
180        let mut levels: HashMap<PirNodeId, u32> = HashMap::new();
181
182        fn compute_level(
183            graph: &PirGraph,
184            id: PirNodeId,
185            visiting: &mut HashSet<PirNodeId>,
186            levels: &mut HashMap<PirNodeId, u32>,
187        ) -> Result<u32> {
188            if let Some(&lvl) = levels.get(&id) {
189                return Ok(lvl);
190            }
191            if !visiting.insert(id) {
192                return Err(XlogError::Compilation(format!(
193                    "Cycle detected while levelizing PIR at node {:?}",
194                    id
195                )));
196            }
197
198            let node = graph.node(id).ok_or_else(|| {
199                XlogError::Compilation(format!("Invalid PIR node id while levelizing: {:?}", id))
200            })?;
201
202            let lvl = match node {
203                PirNode::Const(_) | PirNode::Lit { .. } | PirNode::NegLit { .. } => 0,
204                PirNode::And { children } | PirNode::Or { children } => {
205                    let mut max_child = 0u32;
206                    for &child in children {
207                        let child_lvl = compute_level(graph, child, visiting, levels)?;
208                        max_child = max_child.max(child_lvl);
209                    }
210                    max_child + 1
211                }
212                PirNode::Decision {
213                    child_false,
214                    child_true,
215                    ..
216                } => {
217                    let lf = compute_level(graph, *child_false, visiting, levels)?;
218                    let lt = compute_level(graph, *child_true, visiting, levels)?;
219                    lf.max(lt) + 1
220                }
221            };
222
223            visiting.remove(&id);
224            levels.insert(id, lvl);
225            Ok(lvl)
226        }
227
228        for &root in roots {
229            compute_level(self, root, &mut visiting, &mut levels)?;
230        }
231
232        let max_level = levels.values().copied().max().unwrap_or(0);
233        let mut buckets: Vec<Vec<PirNodeId>> = vec![Vec::new(); (max_level as usize) + 1];
234
235        let mut ids: Vec<PirNodeId> = levels.keys().copied().collect();
236        ids.sort_by_key(|id| id.0);
237        for id in ids {
238            let lvl = levels[&id] as usize;
239            buckets[lvl].push(id);
240        }
241
242        Ok(buckets)
243    }
244}
245
246#[cfg(test)]
247mod tests {
248    use super::*;
249
250    #[test]
251    fn test_levelize_simple_dag() {
252        let mut pir = PirGraph::new();
253        let a = pir.lit(LeafId(0));
254        let b = pir.lit(LeafId(1));
255        let and_ab = pir.and(vec![a, b]);
256        let root = pir.or(vec![and_ab, a]);
257
258        let levels = pir.levelize(&[root]).unwrap();
259
260        assert!(!levels.is_empty());
261        assert!(levels.iter().any(|lvl| lvl.contains(&root)));
262        assert!(levels.iter().any(|lvl| lvl.contains(&a)));
263        assert!(levels.iter().any(|lvl| lvl.contains(&b)));
264    }
265
266    #[test]
267    fn test_levelize_decision_node() {
268        let mut pir = PirGraph::new();
269        let t = pir.const_true();
270        let f = pir.const_false();
271        let d = pir.decision(ChoiceVarId(0), f, t);
272
273        let levels = pir.levelize(&[d]).unwrap();
274        assert!(levels.iter().any(|lvl| lvl.contains(&d)));
275        assert!(levels.iter().any(|lvl| lvl.contains(&t)));
276        assert!(levels.iter().any(|lvl| lvl.contains(&f)));
277    }
278}
xlog_prob/pir.rs

xlog_prob/
pir.rs
