geographdb_core/storage/
sectioned_graph.rs1use anyhow::{Context, Result};
7
8use super::data_structures::{EdgeRec, MetadataRec, NodeRec};
9use super::sectioned::SectionedStorage;
10
11#[derive(Debug, Clone, Default)]
21pub struct GraphData {
22 pub nodes: Vec<NodeRec>,
23 pub edges: Vec<EdgeRec>,
24 pub metadata: Vec<Option<MetadataRec>>,
25}
26
27impl GraphData {
28 pub fn to_bytes(&self) -> Vec<u8> {
30 let mut bytes = Vec::with_capacity(self.required_capacity());
31
32 let node_count: u64 = self.nodes.len() as u64;
34 bytes.extend_from_slice(&node_count.to_le_bytes());
35
36 bytes.extend_from_slice(bytemuck::cast_slice(&self.nodes));
38
39 let edge_count: u64 = self.edges.len() as u64;
41 bytes.extend_from_slice(&edge_count.to_le_bytes());
42
43 bytes.extend_from_slice(bytemuck::cast_slice(&self.edges));
45
46 let metadata_count: u64 = self.metadata.len() as u64;
48 bytes.extend_from_slice(&metadata_count.to_le_bytes());
49
50 let metadata_recs: Vec<MetadataRec> = self
52 .metadata
53 .iter()
54 .map(|m| m.unwrap_or_else(MetadataRec::default))
55 .collect();
56 bytes.extend_from_slice(bytemuck::cast_slice(&metadata_recs));
57
58 bytes
59 }
60
61 pub fn from_bytes(bytes: &[u8]) -> Result<Self> {
63 if bytes.len() < 8 {
64 return Err(anyhow::anyhow!("Graph data too short: {}", bytes.len()));
65 }
66
67 let mut pos = 0;
68
69 let node_count = u64::from_le_bytes(
71 bytes[pos..pos + 8]
72 .try_into()
73 .context("Invalid node count bytes")?,
74 ) as usize;
75 pos += 8;
76
77 let node_bytes = node_count * std::mem::size_of::<NodeRec>();
79 if pos + node_bytes > bytes.len() {
80 return Err(anyhow::anyhow!(
81 "Not enough data for nodes: need {}, have {}",
82 pos + node_bytes,
83 bytes.len()
84 ));
85 }
86
87 let node_slice: &[NodeRec] = bytemuck::try_cast_slice(&bytes[pos..pos + node_bytes])
89 .map_err(|e| anyhow::anyhow!("Invalid node record slice: {}", e))?;
90 let nodes = node_slice.to_vec();
91 pos += node_bytes;
92
93 if pos + 8 > bytes.len() {
95 return Err(anyhow::anyhow!("Not enough data for edge count"));
96 }
97 let edge_count = u64::from_le_bytes(
98 bytes[pos..pos + 8]
99 .try_into()
100 .context("Invalid edge count bytes")?,
101 ) as usize;
102 pos += 8;
103
104 let edge_bytes = edge_count * std::mem::size_of::<EdgeRec>();
106 if pos + edge_bytes > bytes.len() {
107 return Err(anyhow::anyhow!(
108 "Not enough data for edges: need {}, have {}",
109 pos + edge_bytes,
110 bytes.len()
111 ));
112 }
113
114 let edge_slice: &[EdgeRec] = bytemuck::try_cast_slice(&bytes[pos..pos + edge_bytes])
116 .map_err(|e| anyhow::anyhow!("Invalid edge record slice: {}", e))?;
117 let edges = edge_slice.to_vec();
118 pos += edge_bytes;
119
120 if pos + 8 > bytes.len() {
122 return Err(anyhow::anyhow!("Not enough data for metadata count"));
123 }
124 let metadata_count = u64::from_le_bytes(
125 bytes[pos..pos + 8]
126 .try_into()
127 .context("Invalid metadata count bytes")?,
128 ) as usize;
129 pos += 8;
130
131 let meta_bytes = metadata_count * std::mem::size_of::<MetadataRec>();
133 if pos + meta_bytes > bytes.len() {
134 return Err(anyhow::anyhow!("Not enough data for metadata records"));
135 }
136 let meta_slice: &[MetadataRec] = bytemuck::try_cast_slice(&bytes[pos..pos + meta_bytes])
137 .map_err(|e| anyhow::anyhow!("Invalid metadata record slice: {}", e))?;
138 let metadata: Vec<Option<MetadataRec>> = meta_slice.iter().copied().map(Some).collect();
139
140 Ok(Self {
141 nodes,
142 edges,
143 metadata,
144 })
145 }
146
147 pub fn required_capacity(&self) -> usize {
149 8 + self.nodes.len() * std::mem::size_of::<NodeRec>()
151 + 8 + self.edges.len() * std::mem::size_of::<EdgeRec>()
153 + 8 + self.metadata.len() * std::mem::size_of::<MetadataRec>()
155 }
156}
157
158pub struct GraphSectionAdapter;
162
163impl GraphSectionAdapter {
164 pub const SECTION_NAME: &'static str = "GRAPH";
165
166 pub fn load(storage: &mut SectionedStorage) -> Result<GraphData> {
168 let bytes = storage
169 .read_section(Self::SECTION_NAME)
170 .context("GRAPH section not found or empty")?;
171 GraphData::from_bytes(&bytes).context("Failed to parse GRAPH section")
172 }
173
174 pub fn save(storage: &mut SectionedStorage, data: &GraphData) -> Result<()> {
178 let bytes = data.to_bytes();
179 let required = bytes.len() as u64;
180
181 if storage.get_section(Self::SECTION_NAME).is_some() {
183 let result = storage.write_section(Self::SECTION_NAME, &bytes);
184
185 if let Err(e) = result {
186 if e.to_string().contains("overflow") || e.to_string().contains("capacity") {
188 let current = storage.get_section(Self::SECTION_NAME).unwrap();
190 let new_capacity = (current.capacity * 2).max(required * 2);
191 storage
192 .resize_section(Self::SECTION_NAME, new_capacity)
193 .context("Failed to resize GRAPH section")?;
194 storage.write_section(Self::SECTION_NAME, &bytes)?;
196 } else {
197 return Err(e);
198 }
199 }
200 } else {
201 let section_capacity = (1024 * 1024).max(required * 2);
203 storage.create_section(Self::SECTION_NAME, section_capacity, 0)?;
204 storage.write_section(Self::SECTION_NAME, &bytes)?;
205 }
206
207 storage.flush()?;
208 Ok(())
209 }
210
211 pub fn init(storage: &mut SectionedStorage) -> Result<()> {
213 let default_capacity = 1024 * 1024; storage.create_section(Self::SECTION_NAME, default_capacity, 0)?;
216
217 let empty = GraphData::default();
219 let bytes = empty.to_bytes();
220 storage.write_section(Self::SECTION_NAME, &bytes)?;
221 storage.flush()?;
222
223 Ok(())
224 }
225
226 pub fn exists(storage: &SectionedStorage) -> bool {
228 storage.get_section(Self::SECTION_NAME).is_some()
229 }
230}
231
232#[cfg(test)]
233mod tests {
234 use super::*;
235 use bytemuck::Zeroable;
236
237 #[test]
238 fn test_graph_data_serialization_roundtrip() {
239 let nodes = vec![NodeRec {
240 id: 1,
241 morton_code: 42,
242 x: 1.0,
243 y: 2.0,
244 z: 3.0,
245 edge_off: 0,
246 edge_len: 0,
247 flags: 0,
248 begin_ts: 0,
249 end_ts: 0,
250 tx_id: 0,
251 visibility: 1,
252 _padding: [0; 7],
253 }];
254
255 let edges = vec![EdgeRec {
256 src: 1,
257 dst: 2,
258 w: 1.5,
259 flags: 0,
260 begin_ts: 0,
261 end_ts: 0,
262 tx_id: 0,
263 visibility: 1,
264 _padding: [0; 7],
265 }];
266
267 let metadata = vec![Some(MetadataRec::from_strings(
268 "function", "return", 10, 20, 1, 0, 1, 10,
269 ))];
270
271 let original = GraphData {
272 nodes,
273 edges,
274 metadata,
275 };
276
277 let bytes = original.to_bytes();
278 let restored = GraphData::from_bytes(&bytes).unwrap();
279
280 assert_eq!(restored.nodes.len(), original.nodes.len());
281 assert_eq!(restored.nodes[0].id, 1);
282 assert_eq!(restored.edges.len(), original.edges.len());
283 assert_eq!(restored.edges[0].src, 1);
284 assert_eq!(restored.metadata.len(), original.metadata.len());
285 assert_eq!(
286 restored.metadata[0].as_ref().unwrap().get_block_kind(),
287 "function"
288 );
289 }
290
291 #[test]
292 fn test_empty_graph_data() {
293 let empty = GraphData::default();
294 let bytes = empty.to_bytes();
295 assert!(bytes.len() >= 8); let restored = GraphData::from_bytes(&bytes).unwrap();
298 assert_eq!(restored.nodes.len(), 0);
299 assert_eq!(restored.edges.len(), 0);
300 assert_eq!(restored.metadata.len(), 0);
301 }
302
303 #[test]
304 fn test_required_capacity() {
305 let data = GraphData {
306 nodes: vec![NodeRec::zeroed(); 10],
307 edges: vec![EdgeRec::zeroed(); 5],
308 metadata: vec![Some(MetadataRec::default()); 10],
309 };
310
311 let expected = 8 + 10 * std::mem::size_of::<NodeRec>()
314 + 8 + 5 * std::mem::size_of::<EdgeRec>()
316 + 8 + 10 * std::mem::size_of::<MetadataRec>();
318 assert_eq!(data.required_capacity(), expected);
319 }
320}