gammalooprs/uv/
poset.rs

1#![allow(dead_code)]
2
3use slotmap::{Key, SecondaryMap, SlotMap, new_key_type};
4use std::collections::HashSet;
5use std::{cmp::Ordering, collections::VecDeque, hash::Hash};
6use symbolica::atom::Symbol;
7use symbolica::symbol;
8
9use ahash::AHashSet;
10use color_eyre::Result;
11use eyre::eyre;
12use serde::{Deserialize, Serialize};
13
14// Define a new key type for the Poset
15new_key_type! {
16    pub struct PosetNode;
17    pub struct DagNode;
18    pub struct UnfoldedWoodNode;
19    pub struct CoverSetNode;
20}
21
22impl DagNode {
23    pub fn symbol(&self) -> Symbol {
24        symbol!(format!("node_{:?}", self.0))
25    }
26}
27/// Trait to define DOT attributes for node data.
28pub trait DotAttrs {
29    fn dot_attrs(&self) -> String;
30}
31
32/// A node in the poset, storing generic data, edges to child nodes, and references to parent odes.
33#[derive(Debug, Eq)]
34pub struct SlotNode<T, R: Key> {
35    pub data: T,
36    pub order: Option<u64>,
37    pub id: R,
38    pub parents: Vec<R>,  // References to parent nodes by key
39    pub children: Vec<R>, // Edges to child nodes by key
40}
41
42impl<T: Hash, R: Key> Hash for SlotNode<T, R> {
43    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
44        self.data.hash(state);
45        self.id.hash(state);
46    }
47}
48
49impl<T: PartialEq, R: Key> PartialEq for SlotNode<T, R> {
50    fn eq(&self, other: &Self) -> bool {
51        self.data == other.data && self.id == other.id
52    }
53}
54
55impl<T: PartialEq, R: Key> PartialOrd for SlotNode<T, R> {
56    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
57        self.order.unwrap().partial_cmp(&other.order.unwrap())
58    }
59}
60
61impl<T: PartialEq + Eq, R: Key> Ord for SlotNode<T, R> {
62    fn cmp(&self, other: &Self) -> Ordering {
63        self.order.unwrap().cmp(&other.order.unwrap())
64    }
65}
66
67impl<T, R: Key> SlotNode<T, R> {
68    pub(crate) fn dot_id(&self, shift: u64) -> String {
69        base62::encode(self.id() - shift)
70    }
71
72    pub(crate) fn to_topo_ordered(&self) -> Result<TopoOrdered<T>>
73    where
74        T: Clone,
75    {
76        Ok(TopoOrdered::new(
77            self.data.clone(),
78            self.order
79                .ok_or_else(|| eyre!("Node has no topological order"))?,
80        ))
81    }
82
83    pub(crate) fn in_degree(&self) -> usize {
84        self.parents.len()
85    }
86
87    pub(crate) fn out_degree(&self) -> usize {
88        self.children.len()
89    }
90
91    pub(crate) fn id(&self) -> u64 {
92        self.id.data().as_ffi()
93    }
94
95    pub(crate) fn add_parent(&mut self, parent: R) {
96        self.parents.push(parent);
97    }
98
99    pub(crate) fn add_child(&mut self, child: R) {
100        self.children.push(child);
101    }
102
103    pub(crate) fn remove_child(&mut self, child: R) {
104        self.children.retain(|&c| c != child);
105    }
106
107    pub(crate) fn remove_parent(&mut self, parent: R) {
108        self.parents.retain(|&c| c != parent);
109    }
110
111    pub(crate) fn is_parent_of(&self, child: R) -> bool {
112        self.children.contains(&child)
113    }
114
115    pub(crate) fn new(data: T, id: R) -> Self {
116        SlotNode {
117            data,
118            id,
119            order: None,
120            parents: Vec::new(),
121            children: Vec::new(),
122        }
123    }
124
125    // pub(crate) fn compare_data(&self, other: &Self) -> Option<std::cmp::Ordering>
126    // where
127    //     T: PartialOrd,
128    // {
129    //     self.data.partial_cmp(&other.data)
130    // }
131}
132
133/// A partially ordered set (poset) that can be built from an iterator and a slotmap.
134pub struct DAG<T, R: Key, D = ()> {
135    pub nodes: SlotMap<R, SlotNode<T, R>>,
136    associated_data: SecondaryMap<R, D>,
137}
138
139impl<T, R: Key, D> DAG<T, R, D> {
140    pub(crate) fn n_nodes(&self) -> usize {
141        self.nodes.len()
142    }
143}
144
145pub type Poset<T, D> = DAG<T, PosetNode, D>;
146pub type HasseDiagram<T, D> = DAG<T, CoverSetNode, D>;
147
148#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
149pub struct TopoOrdered<T> {
150    pub data: T,
151    order: u64,
152}
153
154impl<T> TopoOrdered<T> {
155    pub(crate) fn new(data: T, order: u64) -> Self {
156        TopoOrdered { data, order }
157    }
158}
159
160#[allow(clippy::non_canonical_partial_ord_impl)]
161impl<T> PartialOrd for TopoOrdered<T> {
162    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
163        self.order.partial_cmp(&other.order)
164    }
165}
166
167impl<T> PartialEq for TopoOrdered<T> {
168    fn eq(&self, other: &Self) -> bool {
169        self.order == other.order
170    }
171}
172
173impl<T> Eq for TopoOrdered<T> {}
174
175impl<T> Ord for TopoOrdered<T> {
176    fn cmp(&self, other: &Self) -> Ordering {
177        self.order.cmp(&other.order)
178    }
179}
180
181impl<T, R: Key, D> DAG<T, R, D> {
182    pub(crate) fn new() -> Self {
183        DAG {
184            nodes: SlotMap::with_key(),
185            associated_data: SecondaryMap::new(),
186        }
187    }
188
189    pub(crate) fn children(&self, key: R) -> impl Iterator<Item = R> + '_ {
190        self.nodes.get(key).unwrap().children.iter().copied()
191    }
192
193    pub(crate) fn dot_id(&self, key: R) -> String {
194        self.nodes.get(key).unwrap().dot_id(self.shift())
195    }
196
197    pub(crate) fn node_values(&self) -> impl Iterator<Item = &T> {
198        self.nodes.values().map(|node| &node.data)
199    }
200
201    pub(crate) fn add_edge(&mut self, from: R, to: R) {
202        let from_node = self.nodes.get_mut(from).unwrap();
203        from_node.add_child(to);
204        let to_node = self.nodes.get_mut(to).unwrap();
205        to_node.add_parent(from);
206    }
207
208    pub(crate) fn add_edge_if_new(&mut self, from: R, to: R) {
209        if !self.nodes.get(from).unwrap().is_parent_of(to) {
210            self.add_edge(from, to);
211        }
212    }
213
214    pub(crate) fn remove_edge(&mut self, from: R, to: R) {
215        let from_node = self.nodes.get_mut(from).unwrap();
216        from_node.remove_child(to);
217        let to_node = self.nodes.get_mut(to).unwrap();
218        to_node.remove_parent(from);
219    }
220
221    pub(crate) fn maximum(&self) -> Option<R>
222    where
223        T: Eq,
224    {
225        Some(self.nodes.values().max()?.id)
226    }
227
228    pub(crate) fn minimum(&self) -> Option<R>
229    where
230        T: Eq,
231    {
232        Some(self.nodes.values().min()?.id)
233    }
234
235    pub(crate) fn succesors(&self, node_key: R) -> &[R] {
236        &self.nodes.get(node_key).unwrap().children
237    }
238
239    /// Returns an iterator over all paths starting from the root node, traversed in BFS order.
240    pub(crate) fn bfs_paths(&self) -> BfsPaths<'_, T, R>
241    where
242        T: Eq,
243    {
244        let mut queue = VecDeque::new();
245        queue.push_back(vec![self.minimum().unwrap()]);
246        BfsPaths {
247            queue,
248            visited: HashSet::new(),
249            nodes: &self.nodes,
250        }
251    }
252
253    pub(crate) fn invert(&mut self) {
254        self.nodes.iter_mut().for_each(|(_, node)| {
255            std::mem::swap(&mut node.children, &mut node.parents);
256        });
257    }
258
259    pub(crate) fn bfs_paths_inv(&self) -> BfsPaths<'_, T, R> {
260        let mut queue = VecDeque::new();
261        let mut maximal_elements = Vec::new();
262        let mut has_incoming = HashSet::new();
263
264        for node in self.nodes.values() {
265            for &parent in node.parents.iter() {
266                has_incoming.insert(parent);
267            }
268        }
269
270        for (key, _node) in self.nodes.iter() {
271            if !has_incoming.contains(&key) {
272                maximal_elements.push(key);
273            }
274        }
275
276        for &max in &maximal_elements {
277            queue.push_back(vec![max]);
278        }
279
280        BfsPaths {
281            queue,
282            visited: HashSet::new(),
283            nodes: &self.nodes,
284        }
285    }
286
287    pub(crate) fn data(&self, key: R) -> &T {
288        &self.nodes.get(key).unwrap().data
289    }
290
291    pub(crate) fn shift(&self) -> u64 {
292        self.nodes.iter().next().unwrap().1.id()
293    }
294
295    pub(crate) fn to_dot(&self, label: &impl Fn(&T) -> String) -> String {
296        self.to_dot_impl(&|node| label(&node.data))
297    }
298
299    pub(crate) fn to_dot_impl(&self, label: &impl Fn(&SlotNode<T, R>) -> String) -> String {
300        let mut dot = String::new();
301        dot.push_str("digraph Poset {\n");
302        dot.push_str("    node [shape=circle];\n");
303
304        let shift = self.shift();
305
306        for node in self.nodes.values() {
307            let node_id = node.dot_id(shift);
308            dot.push_str(&format!("n{} [{}];\n", node_id, label(node)));
309            for &child in node.children.iter() {
310                dot.push_str(&format!(
311                    "n{} -> n{};\n",
312                    node_id,
313                    self.nodes.get(child).unwrap().dot_id(shift)
314                ));
315            }
316        }
317
318        dot.push_str("}\n");
319        dot
320    }
321
322    pub(crate) fn dot_structure(&self) -> String {
323        let shift = self.shift();
324        self.to_dot_impl(&|n| format!("label={}", n.dot_id(shift)))
325    }
326
327    pub(crate) fn add_node(&mut self, data: T) -> R {
328        self.nodes.insert_with_key(|key| SlotNode::new(data, key))
329    }
330
331    /// Returns all descendants of the node, used for propagating transitive relations.
332    fn get_all_descendants(&self, node_key: R) -> HashSet<R> {
333        let mut descendants = HashSet::new();
334        let mut stack = vec![node_key];
335        while let Some(current_key) = stack.pop() {
336            let current_node = self.nodes.get(current_key).unwrap();
337            for &child_key in current_node.children.iter() {
338                if descendants.insert(child_key) {
339                    stack.push(child_key);
340                }
341            }
342        }
343        descendants
344    }
345
346    /// Returns all ancestors of the node, used for propagating transitive relations.
347    fn get_all_ancestors(&self, node_key: R) -> HashSet<R> {
348        let mut ancestors = HashSet::new();
349        let mut stack = vec![node_key];
350        while let Some(current_key) = stack.pop() {
351            let current_node = self.nodes.get(current_key).unwrap();
352            for &parent_key in current_node.parents.iter() {
353                if ancestors.insert(parent_key) {
354                    stack.push(parent_key);
355                }
356            }
357        }
358        ancestors
359    }
360
361    pub(crate) fn transitive_edges(&self, a: R) -> Vec<(R, R)> {
362        let mut edges = Vec::new();
363
364        let children: AHashSet<_> = self
365            .nodes
366            .get(a)
367            .unwrap()
368            .children
369            .iter()
370            .cloned()
371            .collect();
372
373        for &child in &children {
374            for descendant in self.get_all_descendants(child) {
375                if children.contains(&descendant) {
376                    edges.push((a, descendant));
377                }
378            }
379        }
380        edges
381    }
382
383    pub(crate) fn in_degree(&self, key: R) -> usize {
384        self.nodes.get(key).unwrap().in_degree()
385    }
386
387    pub(crate) fn compute_topological_order(&mut self) -> Vec<R> {
388        // Initialize queue with nodes having in-degree zero
389        let mut queue = VecDeque::new(); //S in the wikipedia article
390
391        let mut indegrees = SecondaryMap::new();
392
393        for (key, node) in self.nodes.iter() {
394            indegrees.insert(key, node.in_degree());
395            if node.in_degree() == 0 {
396                queue.push_back(key);
397            }
398        }
399
400        let mut order = vec![];
401        while let Some(node_key) = queue.pop_front() {
402            // Assign the order number to the node
403            if let Some(node) = self.nodes.get_mut(node_key) {
404                node.order = Some(order.len() as u64);
405                order.push(node_key);
406            }
407
408            // For each child, decrement its in-degree
409            for &child_key in &self.nodes.get(node_key).unwrap().children {
410                indegrees[child_key] -= 1;
411                if indegrees[child_key] == 0 {
412                    queue.push_back(child_key);
413                }
414            }
415        }
416
417        // Optional: Check if graph has cycles
418        if order.len() != self.nodes.len() {
419            panic!("The graph contains a cycle!");
420        }
421        order
422    }
423}
424
425impl<T, R: Key, D> Default for DAG<T, R, D> {
426    fn default() -> Self {
427        Self::new()
428    }
429}
430
431impl<T, D> Poset<T, D> {
432    pub(crate) fn poset_family(&self, data: &T) -> [Vec<PosetNode>; 2]
433    where
434        T: PartialOrd,
435    {
436        let mut parents = Vec::new();
437        let mut children = Vec::new();
438        for (key, node) in self.nodes.iter() {
439            match data.partial_cmp(&node.data) {
440                Some(std::cmp::Ordering::Greater) => {
441                    children.push(key);
442                }
443                Some(std::cmp::Ordering::Less) => {
444                    parents.push(key);
445                }
446                _ => {}
447            }
448        }
449        [parents, children]
450    }
451    pub(crate) fn poset_push(&mut self, data: T, associated_data: D, flip: bool) -> PosetNode
452    where
453        T: PartialOrd,
454    {
455        let id = self.nodes.insert_with_key(|key| SlotNode::new(data, key));
456
457        self.associated_data.insert(id, associated_data);
458        let new_node = self.nodes.get(id).unwrap();
459
460        let [mut parents, mut children] = self.poset_family(&new_node.data);
461
462        if flip {
463            std::mem::swap(&mut parents, &mut children);
464        }
465
466        for &parent_key in &parents {
467            self.add_edge(parent_key, id);
468        }
469
470        for &child_key in &children {
471            self.add_edge(id, child_key);
472        }
473
474        self.update_transitive_closure(id);
475
476        id
477    }
478
479    /// Updates the transitive closure of the poset by propagating the relationships.
480    fn update_transitive_closure(&mut self, new_node_key: PosetNode) {
481        // Propagate relationships for all descendants of new_node
482        let descendants = self.get_all_descendants(new_node_key);
483        for &descendant_key in &descendants {
484            self.add_edge_if_new(new_node_key, descendant_key);
485        }
486
487        // Propagate relationships for all ancestors of new_node
488        let ancestors = self.get_all_ancestors(new_node_key);
489        for &ancestor_key in &ancestors {
490            self.add_edge_if_new(ancestor_key, new_node_key);
491        }
492    }
493
494    pub(crate) fn remove_transitive_edges(mut self) -> HasseDiagram<T, D> {
495        let edges_to_remove: Vec<_> = self
496            .nodes
497            .keys()
498            .flat_map(|node_key| self.transitive_edges(node_key))
499            .collect();
500
501        // let shift = self.shift();
502        for (a, b) in edges_to_remove {
503            // println!(
504            //     "removing edge from {} to {}",
505            //     base62::encode(a.0.as_ffi() - shift),
506            //     base62::encode(b.0.as_ffi() - shift)
507            // );
508            self.remove_edge(a, b);
509        }
510
511        let mut hasse = HasseDiagram::new();
512
513        let mut new_map = SecondaryMap::new();
514
515        for (key, node) in self.nodes.into_iter() {
516            new_map.insert(key, hasse.add_node(node.data));
517            if let Some(d) = self.associated_data.remove(key) {
518                hasse.associated_data.insert(new_map[key], d);
519            }
520            for child in node.children {
521                hasse.add_edge(new_map[key], new_map[child]);
522            }
523        }
524
525        hasse
526    }
527}
528
529impl<T: PartialOrd, D> FromIterator<(T, D)> for Poset<T, D> {
530    fn from_iter<I: IntoIterator<Item = (T, D)>>(iter: I) -> Self {
531        let mut poset = Poset::new();
532        for (data, assoc) in iter {
533            poset.poset_push(data, assoc, false);
534        }
535        poset
536    }
537}
538
539/// An iterator over paths in the poset, traversed in BFS order.
540pub struct BfsPaths<'a, T, R: Key> {
541    queue: VecDeque<Vec<R>>,
542    visited: HashSet<Vec<R>>,
543    nodes: &'a SlotMap<R, SlotNode<T, R>>,
544}
545
546impl<'a, T, R: Key> Iterator for BfsPaths<'a, T, R> {
547    type Item = Vec<R>;
548
549    fn next(&mut self) -> Option<Self::Item> {
550        while let Some(path) = self.queue.pop_front() {
551            if !self.visited.insert(path.clone()) {
552                continue;
553            }
554
555            let last_node_key = path.last().unwrap();
556            let last_node = self.nodes.get(*last_node_key).unwrap();
557
558            for &child_key in last_node.children.iter() {
559                if !path.contains(&child_key) {
560                    let mut new_path = path.clone();
561                    new_path.push(child_key);
562                    self.queue.push_back(new_path);
563                }
564            }
565
566            return Some(path);
567        }
568        None
569    }
570}
gammalooprs/uv/poset.rs

gammalooprs/uv/
poset.rs
