Skip to main content

weavatrix_graph/payload/stable/
freeze.rs

1use super::StablePayloadGraph;
2use crate::Vec;
3use crate::{
4    EdgeEndpoints, EdgeIndex, NodeIndex, PayloadGraph, Result, StableEdgeKey, StableNodeKey,
5};
6
7#[derive(Debug, Clone, PartialEq, Eq)]
8pub struct PayloadFreezeMap {
9    nodes: Vec<Option<(u32, NodeIndex)>>,
10    edges: Vec<Option<(u32, EdgeIndex)>>,
11}
12
13impl PayloadFreezeMap {
14    pub(crate) const fn from_slots(
15        nodes: Vec<Option<(u32, NodeIndex)>>,
16        edges: Vec<Option<(u32, EdgeIndex)>>,
17    ) -> Self {
18        Self { nodes, edges }
19    }
20
21    #[must_use]
22    pub fn node(&self, key: StableNodeKey) -> Option<NodeIndex> {
23        let (generation, index) = self.nodes.get(key.index())?.as_ref()?;
24        (*generation == key.generation()).then_some(*index)
25    }
26
27    #[must_use]
28    pub fn edge(&self, key: StableEdgeKey) -> Option<EdgeIndex> {
29        let (generation, index) = self.edges.get(key.index())?.as_ref()?;
30        (*generation == key.generation()).then_some(*index)
31    }
32}
33
34#[derive(Debug, Clone, PartialEq, Eq)]
35pub struct FrozenPayloadGraph<NodePayload, EdgePayload> {
36    graph: PayloadGraph<NodePayload, EdgePayload>,
37    indices: PayloadFreezeMap,
38}
39
40impl<NodePayload, EdgePayload> FrozenPayloadGraph<NodePayload, EdgePayload> {
41    #[must_use]
42    pub const fn graph(&self) -> &PayloadGraph<NodePayload, EdgePayload> {
43        &self.graph
44    }
45
46    #[must_use]
47    pub const fn indices(&self) -> &PayloadFreezeMap {
48        &self.indices
49    }
50
51    #[must_use]
52    pub fn into_parts(self) -> (PayloadGraph<NodePayload, EdgePayload>, PayloadFreezeMap) {
53        (self.graph, self.indices)
54    }
55}
56
57impl<NodePayload, EdgePayload> StablePayloadGraph<NodePayload, EdgePayload> {
58    /// Compacts live slots into an immutable CSR-backed payload graph.
59    ///
60    /// # Errors
61    ///
62    /// Returns an error if compact index capacity is exhausted.
63    pub fn freeze(self) -> Result<FrozenPayloadGraph<NodePayload, EdgePayload>> {
64        let mut node_map = vec![None; self.nodes.len()];
65        let mut nodes = Vec::with_capacity(self.node_count);
66        for (slot, entry) in self.nodes.into_iter().enumerate() {
67            let Some(payload) = entry.value else {
68                continue;
69            };
70            let compact = NodeIndex::new(u32::try_from(nodes.len()).map_err(|_| {
71                crate::GraphError::IndexCapacityExceeded {
72                    category: "frozen payload nodes",
73                    count: nodes.len(),
74                }
75            })?);
76            node_map[slot] = Some((entry.generation, compact));
77            nodes.push(payload);
78        }
79
80        let mut edge_map = vec![None; self.edges.len()];
81        let mut edges = Vec::with_capacity(self.edge_count);
82        for (slot, entry) in self.edges.into_iter().enumerate() {
83            let Some(payload) = entry.value else {
84                continue;
85            };
86            let source = mapped_node(&node_map, entry.source)?;
87            let target = mapped_node(&node_map, entry.target)?;
88            let compact = EdgeIndex::new(u32::try_from(edges.len()).map_err(|_| {
89                crate::GraphError::IndexCapacityExceeded {
90                    category: "frozen payload edges",
91                    count: edges.len(),
92                }
93            })?);
94            edge_map[slot] = Some((entry.generation, compact));
95            edges.push((EdgeEndpoints::new(source, target), payload));
96        }
97        let graph = PayloadGraph::try_from_edges(nodes, edges)?;
98        Ok(FrozenPayloadGraph {
99            graph,
100            indices: PayloadFreezeMap::from_slots(node_map, edge_map),
101        })
102    }
103}
104
105pub(crate) fn mapped_node(nodes: &[Option<(u32, NodeIndex)>], slot: u32) -> Result<NodeIndex> {
106    nodes
107        .get(slot as usize)
108        .and_then(Option::as_ref)
109        .copied()
110        .map(|(_, index)| index)
111        .ok_or(crate::GraphError::InvalidNodeIndex {
112            node: slot as usize,
113            node_count: nodes.len(),
114        })
115}