Skip to main content

weavatrix_graph/payload/
keyed.rs

1use super::StablePayloadGraph;
2use crate::{GraphError, Result, StableEdgeKey, StableNodeKey};
3#[cfg(not(feature = "std"))]
4use alloc::collections::BTreeMap as KeyMap;
5use core::hash::Hash;
6#[cfg(feature = "std")]
7use std::collections::HashMap as KeyMap;
8
9/// A GraphMap-style key index backed by generation-stable graph handles.
10///
11/// Keys provide domain lookup while algorithms continue to use compact stable
12/// handles from [`StablePayloadGraph`].
13#[derive(Debug, Clone)]
14pub struct KeyedPayloadGraph<Key, NodePayload, EdgePayload> {
15    graph: StablePayloadGraph<NodePayload, EdgePayload>,
16    keys: KeyMap<Key, StableNodeKey>,
17}
18
19impl<Key, NodePayload, EdgePayload> Default for KeyedPayloadGraph<Key, NodePayload, EdgePayload>
20where
21    Key: Clone + Eq + Hash + Ord,
22{
23    fn default() -> Self {
24        Self::new()
25    }
26}
27
28impl<Key, NodePayload, EdgePayload> KeyedPayloadGraph<Key, NodePayload, EdgePayload>
29where
30    Key: Clone + Eq + Hash + Ord,
31{
32    #[must_use]
33    pub fn new() -> Self {
34        Self {
35            graph: StablePayloadGraph::new(),
36            keys: key_map_with_capacity(0),
37        }
38    }
39
40    #[must_use]
41    pub fn with_capacity(nodes: usize, edges: usize) -> Self {
42        Self {
43            graph: StablePayloadGraph::with_capacity(nodes, edges),
44            keys: key_map_with_capacity(nodes),
45        }
46    }
47
48    /// Inserts a key or replaces its node payload without changing its handle.
49    ///
50    /// Returns the stable handle and the previous payload, when one existed.
51    ///
52    /// # Errors
53    ///
54    /// Returns an error when the stable node index space is exhausted.
55    pub fn insert_node(
56        &mut self,
57        key: Key,
58        payload: NodePayload,
59    ) -> Result<(StableNodeKey, Option<NodePayload>)> {
60        if let Some(handle) = self.keys.get(&key).copied() {
61            if let Some(node) = self.graph.node_mut(handle) {
62                return Ok((handle, Some(core::mem::replace(node, payload))));
63            }
64            self.keys.remove(&key);
65        }
66        let handle = self.graph.add_node(payload)?;
67        self.keys.insert(key, handle);
68        Ok((handle, None))
69    }
70
71    #[must_use]
72    pub fn node_key(&self, key: &Key) -> Option<StableNodeKey> {
73        self.keys.get(key).copied()
74    }
75
76    #[must_use]
77    pub fn node(&self, key: &Key) -> Option<&NodePayload> {
78        self.graph.node(self.node_key(key)?)
79    }
80
81    #[must_use]
82    pub fn node_mut(&mut self, key: &Key) -> Option<&mut NodePayload> {
83        let handle = self.keys.get(key).copied()?;
84        self.graph.node_mut(handle)
85    }
86
87    /// Adds a directed edge between two existing domain keys.
88    ///
89    /// # Errors
90    ///
91    /// Returns an error for a missing endpoint or exhausted edge indices.
92    pub fn add_edge(
93        &mut self,
94        source: &Key,
95        target: &Key,
96        payload: EdgePayload,
97    ) -> Result<StableEdgeKey> {
98        let source = self
99            .node_key(source)
100            .ok_or(GraphError::MissingKeyedNode { endpoint: "source" })?;
101        let target = self
102            .node_key(target)
103            .ok_or(GraphError::MissingKeyedNode { endpoint: "target" })?;
104        self.graph.add_edge(source, target, payload)
105    }
106
107    pub fn remove_node(&mut self, key: &Key) -> Option<NodePayload> {
108        let handle = self.keys.remove(key)?;
109        self.graph.remove_node(handle)
110    }
111
112    #[must_use]
113    pub const fn graph(&self) -> &StablePayloadGraph<NodePayload, EdgePayload> {
114        &self.graph
115    }
116
117    #[must_use]
118    pub fn node_count(&self) -> usize {
119        self.graph.node_count()
120    }
121
122    #[must_use]
123    pub fn edge_count(&self) -> usize {
124        self.graph.edge_count()
125    }
126
127    pub fn into_parts(
128        self,
129    ) -> (
130        StablePayloadGraph<NodePayload, EdgePayload>,
131        impl Iterator<Item = (Key, StableNodeKey)>,
132    ) {
133        (self.graph, self.keys.into_iter())
134    }
135}
136
137#[cfg(feature = "std")]
138fn key_map_with_capacity<Key, Value>(capacity: usize) -> KeyMap<Key, Value> {
139    KeyMap::with_capacity(capacity)
140}
141
142#[cfg(not(feature = "std"))]
143fn key_map_with_capacity<Key, Value>(_capacity: usize) -> KeyMap<Key, Value> {
144    KeyMap::new()
145}