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::new();
31
32 let node_count: u64 = self.nodes.len() as u64;
34 bytes.extend_from_slice(&node_count.to_le_bytes());
35
36 for node in &self.nodes {
38 bytes.extend_from_slice(bytemuck::bytes_of(node));
39 }
40
41 let edge_count: u64 = self.edges.len() as u64;
43 bytes.extend_from_slice(&edge_count.to_le_bytes());
44
45 for edge in &self.edges {
47 bytes.extend_from_slice(bytemuck::bytes_of(edge));
48 }
49
50 let metadata_count: u64 = self.metadata.len() as u64;
52 bytes.extend_from_slice(&metadata_count.to_le_bytes());
53
54 for meta in &self.metadata {
56 if let Some(m) = meta {
57 bytes.extend_from_slice(bytemuck::bytes_of(m));
58 } else {
59 bytes.extend_from_slice(bytemuck::bytes_of(&MetadataRec::default()));
61 }
62 }
63
64 bytes
65 }
66
67 pub fn from_bytes(bytes: &[u8]) -> Result<Self> {
69 if bytes.len() < 8 {
70 return Err(anyhow::anyhow!("Graph data too short: {}", bytes.len()));
71 }
72
73 let mut pos = 0;
74
75 let node_count = u64::from_le_bytes(
77 bytes[pos..pos + 8]
78 .try_into()
79 .context("Invalid node count bytes")?,
80 ) as usize;
81 pos += 8;
82
83 let node_bytes = node_count * std::mem::size_of::<NodeRec>();
85 if pos + node_bytes > bytes.len() {
86 return Err(anyhow::anyhow!(
87 "Not enough data for nodes: need {}, have {}",
88 pos + node_bytes,
89 bytes.len()
90 ));
91 }
92
93 let mut nodes = Vec::with_capacity(node_count);
95 for i in 0..node_count {
96 let start = pos + i * std::mem::size_of::<NodeRec>();
97 let end = start + std::mem::size_of::<NodeRec>();
98 let node_rec: &NodeRec = bytemuck::try_from_bytes::<NodeRec>(&bytes[start..end])
99 .map_err(|e| anyhow::anyhow!("Invalid node record: {}", e))?;
100 nodes.push(*node_rec);
101 }
102 pos += node_bytes;
103
104 if pos + 8 > bytes.len() {
106 return Err(anyhow::anyhow!("Not enough data for edge count"));
107 }
108 let edge_count = u64::from_le_bytes(
109 bytes[pos..pos + 8]
110 .try_into()
111 .context("Invalid edge count bytes")?,
112 ) as usize;
113 pos += 8;
114
115 let edge_bytes = edge_count * std::mem::size_of::<EdgeRec>();
117 if pos + edge_bytes > bytes.len() {
118 return Err(anyhow::anyhow!(
119 "Not enough data for edges: need {}, have {}",
120 pos + edge_bytes,
121 bytes.len()
122 ));
123 }
124
125 let mut edges = Vec::with_capacity(edge_count);
127 for i in 0..edge_count {
128 let start = pos + i * std::mem::size_of::<EdgeRec>();
129 let end = start + std::mem::size_of::<EdgeRec>();
130 let edge_rec: &EdgeRec = bytemuck::try_from_bytes::<EdgeRec>(&bytes[start..end])
131 .map_err(|e| anyhow::anyhow!("Invalid edge record: {}", e))?;
132 edges.push(*edge_rec);
133 }
134 pos += edge_bytes;
135
136 if pos + 8 > bytes.len() {
138 return Err(anyhow::anyhow!("Not enough data for metadata count"));
139 }
140 let metadata_count = u64::from_le_bytes(
141 bytes[pos..pos + 8]
142 .try_into()
143 .context("Invalid metadata count bytes")?,
144 ) as usize;
145 pos += 8;
146
147 let mut metadata = Vec::with_capacity(metadata_count);
149 for i in 0..metadata_count {
150 let start = pos + i * std::mem::size_of::<MetadataRec>();
151 let end = start + std::mem::size_of::<MetadataRec>();
152 if end > bytes.len() {
153 return Err(anyhow::anyhow!("Not enough data for metadata records"));
154 }
155 let meta_rec: &MetadataRec =
156 bytemuck::try_from_bytes::<MetadataRec>(&bytes[start..end])
157 .map_err(|e| anyhow::anyhow!("Invalid metadata record: {}", e))?;
158 metadata.push(Some(*meta_rec));
159 }
160
161 Ok(Self {
162 nodes,
163 edges,
164 metadata,
165 })
166 }
167
168 pub fn required_capacity(&self) -> usize {
170 8 + self.nodes.len() * std::mem::size_of::<NodeRec>()
172 + 8 + self.edges.len() * std::mem::size_of::<EdgeRec>()
174 + 8 + self.metadata.len() * std::mem::size_of::<MetadataRec>()
176 }
177}
178
179pub struct GraphSectionAdapter;
183
184impl GraphSectionAdapter {
185 pub const SECTION_NAME: &'static str = "GRAPH";
186
187 pub fn load(storage: &mut SectionedStorage) -> Result<GraphData> {
189 let bytes = storage
190 .read_section(Self::SECTION_NAME)
191 .context("GRAPH section not found or empty")?;
192 GraphData::from_bytes(&bytes).context("Failed to parse GRAPH section")
193 }
194
195 pub fn save(storage: &mut SectionedStorage, data: &GraphData) -> Result<()> {
199 let bytes = data.to_bytes();
200 let required = bytes.len() as u64;
201
202 if storage.get_section(Self::SECTION_NAME).is_some() {
204 let result = storage.write_section(Self::SECTION_NAME, &bytes);
205
206 if let Err(e) = result {
207 if e.to_string().contains("overflow") || e.to_string().contains("capacity") {
209 let current = storage.get_section(Self::SECTION_NAME).unwrap();
211 let new_capacity = (current.capacity * 2).max(required * 2);
212 storage
213 .resize_section(Self::SECTION_NAME, new_capacity)
214 .context("Failed to resize GRAPH section")?;
215 storage.write_section(Self::SECTION_NAME, &bytes)?;
217 } else {
218 return Err(e);
219 }
220 }
221 } else {
222 let section_capacity = (1024 * 1024).max(required * 2);
224 storage.create_section(Self::SECTION_NAME, section_capacity, 0)?;
225 storage.write_section(Self::SECTION_NAME, &bytes)?;
226 }
227
228 storage.flush()?;
229 Ok(())
230 }
231
232 pub fn init(storage: &mut SectionedStorage) -> Result<()> {
234 let default_capacity = 1024 * 1024; storage.create_section(Self::SECTION_NAME, default_capacity, 0)?;
237
238 let empty = GraphData::default();
240 let bytes = empty.to_bytes();
241 storage.write_section(Self::SECTION_NAME, &bytes)?;
242 storage.flush()?;
243
244 Ok(())
245 }
246
247 pub fn exists(storage: &SectionedStorage) -> bool {
249 storage.get_section(Self::SECTION_NAME).is_some()
250 }
251}
252
253#[cfg(test)]
254mod tests {
255 use super::*;
256 use bytemuck::Zeroable;
257
258 #[test]
259 fn test_graph_data_serialization_roundtrip() {
260 let nodes = vec![NodeRec {
261 id: 1,
262 morton_code: 42,
263 x: 1.0,
264 y: 2.0,
265 z: 3.0,
266 edge_off: 0,
267 edge_len: 0,
268 flags: 0,
269 begin_ts: 0,
270 end_ts: 0,
271 tx_id: 0,
272 visibility: 1,
273 _padding: [0; 7],
274 }];
275
276 let edges = vec![EdgeRec {
277 src: 1,
278 dst: 2,
279 w: 1.5,
280 flags: 0,
281 begin_ts: 0,
282 end_ts: 0,
283 tx_id: 0,
284 visibility: 1,
285 _padding: [0; 7],
286 }];
287
288 let metadata = vec![Some(MetadataRec::from_strings(
289 "function", "return", 10, 20, 1, 0, 1, 10,
290 ))];
291
292 let original = GraphData {
293 nodes,
294 edges,
295 metadata,
296 };
297
298 let bytes = original.to_bytes();
299 let restored = GraphData::from_bytes(&bytes).unwrap();
300
301 assert_eq!(restored.nodes.len(), original.nodes.len());
302 assert_eq!(restored.nodes[0].id, 1);
303 assert_eq!(restored.edges.len(), original.edges.len());
304 assert_eq!(restored.edges[0].src, 1);
305 assert_eq!(restored.metadata.len(), original.metadata.len());
306 assert_eq!(
307 restored.metadata[0].as_ref().unwrap().get_block_kind(),
308 "function"
309 );
310 }
311
312 #[test]
313 fn test_empty_graph_data() {
314 let empty = GraphData::default();
315 let bytes = empty.to_bytes();
316 assert!(bytes.len() >= 8); let restored = GraphData::from_bytes(&bytes).unwrap();
319 assert_eq!(restored.nodes.len(), 0);
320 assert_eq!(restored.edges.len(), 0);
321 assert_eq!(restored.metadata.len(), 0);
322 }
323
324 #[test]
325 fn test_required_capacity() {
326 let data = GraphData {
327 nodes: vec![NodeRec::zeroed(); 10],
328 edges: vec![EdgeRec::zeroed(); 5],
329 metadata: vec![Some(MetadataRec::default()); 10],
330 };
331
332 let expected = 8 + 10 * std::mem::size_of::<NodeRec>()
335 + 8 + 5 * std::mem::size_of::<EdgeRec>()
337 + 8 + 10 * std::mem::size_of::<MetadataRec>();
339 assert_eq!(data.required_capacity(), expected);
340 }
341}