weavatrix_graph/working/
freeze.rs1use 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 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}