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

gammalooprs/uv/
poset.rs
