weavatrix_graph/payload/
keyed.rs1use 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#[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 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 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}