1use arcstr::ArcStr;
7use grafeo_common::types::{EdgeId, EpochId, NodeId, PropertyKey, PropertyMap, Value};
8use serde::{Deserialize, Serialize};
9
10#[derive(Debug, Clone)]
31pub struct Edge {
32 pub id: EdgeId,
34 pub src: NodeId,
36 pub dst: NodeId,
38 pub edge_type: ArcStr,
40 pub properties: PropertyMap,
42}
43
44impl Edge {
45 #[must_use]
47 pub fn new(id: EdgeId, src: NodeId, dst: NodeId, edge_type: impl Into<ArcStr>) -> Self {
48 Self {
49 id,
50 src,
51 dst,
52 edge_type: edge_type.into(),
53 properties: PropertyMap::new(),
54 }
55 }
56
57 pub fn set_property(&mut self, key: impl Into<PropertyKey>, value: impl Into<Value>) {
59 self.properties.insert(key.into(), value.into());
60 }
61
62 #[must_use]
64 pub fn get_property(&self, key: &str) -> Option<&Value> {
65 self.properties.get(&PropertyKey::new(key))
66 }
67
68 pub fn remove_property(&mut self, key: &str) -> Option<Value> {
70 self.properties.remove(&PropertyKey::new(key))
71 }
72
73 #[must_use]
75 pub fn properties_as_btree(&self) -> std::collections::BTreeMap<PropertyKey, Value> {
76 self.properties.to_btree_map()
77 }
78
79 #[must_use]
84 pub fn other_endpoint(&self, node: NodeId) -> Option<NodeId> {
85 if node == self.src {
86 Some(self.dst)
87 } else if node == self.dst {
88 Some(self.src)
89 } else {
90 None
91 }
92 }
93}
94
95#[repr(C)]
100#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
101pub struct EdgeRecord {
102 pub id: EdgeId,
104 pub src: NodeId,
106 pub dst: NodeId,
108 pub type_id: u32,
110 pub props_offset: u32,
112 pub props_count: u16,
114 pub flags: EdgeFlags,
116 pub epoch: EpochId,
118}
119
120impl EdgeRecord {
121 pub const FLAG_DELETED: u16 = 1 << 0;
123 pub const FLAG_HAS_VERSION: u16 = 1 << 1;
125
126 #[must_use]
128 pub const fn new(id: EdgeId, src: NodeId, dst: NodeId, type_id: u32, epoch: EpochId) -> Self {
129 Self {
130 id,
131 src,
132 dst,
133 type_id,
134 props_offset: 0,
135 props_count: 0,
136 flags: EdgeFlags(0),
137 epoch,
138 }
139 }
140
141 #[must_use]
143 pub const fn is_deleted(&self) -> bool {
144 self.flags.contains(Self::FLAG_DELETED)
145 }
146
147 pub fn set_deleted(&mut self, deleted: bool) {
149 if deleted {
150 self.flags.set(Self::FLAG_DELETED);
151 } else {
152 self.flags.clear(Self::FLAG_DELETED);
153 }
154 }
155}
156
157#[repr(transparent)]
159#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize)]
160pub struct EdgeFlags(pub u16);
161
162impl EdgeFlags {
163 #[must_use]
165 pub const fn contains(&self, flag: u16) -> bool {
166 (self.0 & flag) != 0
167 }
168
169 pub fn set(&mut self, flag: u16) {
171 self.0 |= flag;
172 }
173
174 pub fn clear(&mut self, flag: u16) {
176 self.0 &= !flag;
177 }
178}
179
180#[cfg(test)]
181mod tests {
182 use super::*;
183
184 #[test]
185 fn test_edge_creation() {
186 let edge = Edge::new(EdgeId::new(1), NodeId::new(10), NodeId::new(20), "KNOWS");
187
188 assert_eq!(edge.id, EdgeId::new(1));
189 assert_eq!(edge.src, NodeId::new(10));
190 assert_eq!(edge.dst, NodeId::new(20));
191 assert_eq!(edge.edge_type.as_str(), "KNOWS");
192 }
193
194 #[test]
195 fn test_edge_properties() {
196 let mut edge = Edge::new(EdgeId::new(1), NodeId::new(10), NodeId::new(20), "KNOWS");
197
198 edge.set_property("since", 2020i64);
199 edge.set_property("weight", 1.5f64);
200
201 assert_eq!(
202 edge.get_property("since").and_then(|v| v.as_int64()),
203 Some(2020)
204 );
205 assert_eq!(
206 edge.get_property("weight").and_then(|v| v.as_float64()),
207 Some(1.5)
208 );
209 }
210
211 #[test]
212 fn test_edge_other_endpoint() {
213 let edge = Edge::new(EdgeId::new(1), NodeId::new(10), NodeId::new(20), "KNOWS");
214
215 assert_eq!(edge.other_endpoint(NodeId::new(10)), Some(NodeId::new(20)));
216 assert_eq!(edge.other_endpoint(NodeId::new(20)), Some(NodeId::new(10)));
217 assert_eq!(edge.other_endpoint(NodeId::new(30)), None);
218 }
219
220 #[test]
221 fn test_edge_record_flags() {
222 let mut record = EdgeRecord::new(
223 EdgeId::new(1),
224 NodeId::new(10),
225 NodeId::new(20),
226 0,
227 EpochId::INITIAL,
228 );
229
230 assert!(!record.is_deleted());
231 record.set_deleted(true);
232 assert!(record.is_deleted());
233 }
234
235 #[test]
236 fn test_edge_record_size() {
237 let size = std::mem::size_of::<EdgeRecord>();
239 assert!(size <= 64, "EdgeRecord is {} bytes", size);
241 }
242}