Skip to main content

weavatrix_graph/working/
freeze.rs

1use super::{StableEdgeKey, StableNodeKey, WorkingGraph};
2use crate::Vec;
3use crate::{EdgeIndex, Graph, NodeIndex, Result};
4#[cfg(not(feature = "std"))]
5use alloc::collections::BTreeMap as StableMap;
6#[cfg(feature = "std")]
7use std::collections::HashMap as StableMap;
8
9#[cfg(feature = "std")]
10fn stable_map_with_capacity<K, V>(capacity: usize) -> StableMap<K, V> {
11    StableMap::with_capacity(capacity)
12}
13
14#[cfg(not(feature = "std"))]
15fn stable_map_with_capacity<K: Ord, V>(_: usize) -> StableMap<K, V> {
16    StableMap::new()
17}
18
19#[derive(Debug, Clone, PartialEq, Eq)]
20pub struct FreezeMap {
21    nodes: StableMap<StableNodeKey, NodeIndex>,
22    edges: StableMap<StableEdgeKey, EdgeIndex>,
23}
24
25impl FreezeMap {
26    #[must_use]
27    pub fn node(&self, key: StableNodeKey) -> Option<NodeIndex> {
28        self.nodes.get(&key).copied()
29    }
30
31    #[must_use]
32    pub fn edge(&self, key: StableEdgeKey) -> Option<EdgeIndex> {
33        self.edges.get(&key).copied()
34    }
35}
36
37#[derive(Debug, Clone, PartialEq, Eq)]
38pub struct FrozenGraph {
39    graph: Graph,
40    indices: FreezeMap,
41}
42
43impl FrozenGraph {
44    #[must_use]
45    pub const fn graph(&self) -> &Graph {
46        &self.graph
47    }
48
49    #[must_use]
50    pub const fn indices(&self) -> &FreezeMap {
51        &self.indices
52    }
53
54    #[must_use]
55    pub fn into_parts(self) -> (Graph, FreezeMap) {
56        (self.graph, self.indices)
57    }
58}
59
60impl WorkingGraph {
61    /// Canonicalizes the working graph and returns stable-to-compact remapping.
62    ///
63    /// # Errors
64    ///
65    /// Returns an error if an internal endpoint invariant is violated.
66    pub fn freeze(self) -> Result<FrozenGraph> {
67        let mut keyed_nodes = self
68            .nodes
69            .into_iter()
70            .enumerate()
71            .filter_map(|(slot, entry)| {
72                Some((
73                    StableNodeKey::new(u32::try_from(slot).ok()?, entry.generation),
74                    entry.value?,
75                ))
76            })
77            .collect::<Vec<_>>();
78        keyed_nodes.sort_unstable_by(|left, right| left.1.id.cmp(&right.1.id));
79
80        let mut node_map = stable_map_with_capacity(keyed_nodes.len());
81        let mut nodes = Vec::with_capacity(keyed_nodes.len());
82        for (index, (key, node)) in keyed_nodes.into_iter().enumerate() {
83            let index =
84                u32::try_from(index).map_err(|_| crate::GraphError::IndexCapacityExceeded {
85                    category: "frozen nodes",
86                    count: nodes.len(),
87                })?;
88            node_map.insert(key, NodeIndex::new(index));
89            nodes.push(node);
90        }
91
92        let mut keyed_edges = self
93            .edges
94            .into_iter()
95            .enumerate()
96            .filter_map(|(slot, entry)| {
97                Some((
98                    StableEdgeKey::new(u32::try_from(slot).ok()?, entry.generation),
99                    entry.value?.value,
100                ))
101            })
102            .collect::<Vec<_>>();
103        keyed_edges.sort_unstable_by(|left, right| left.1.cmp(&right.1));
104
105        let mut edge_map = stable_map_with_capacity(keyed_edges.len());
106        let mut edges = Vec::with_capacity(keyed_edges.len());
107        for (key, edge) in keyed_edges {
108            let index = if edges.last() == Some(&edge) {
109                edges.len() - 1
110            } else {
111                edges.push(edge);
112                edges.len() - 1
113            };
114            let index =
115                u32::try_from(index).map_err(|_| crate::GraphError::IndexCapacityExceeded {
116                    category: "frozen edges",
117                    count: edges.len(),
118                })?;
119            edge_map.insert(key, EdgeIndex::new(index));
120        }
121
122        let graph = Graph::from_validated_sorted_parts(nodes, edges)?;
123        Ok(FrozenGraph {
124            graph,
125            indices: FreezeMap {
126                nodes: node_map,
127                edges: edge_map,
128            },
129        })
130    }
131}