Skip to main content

geographdb_core/storage/
sectioned_graph.rs

1//! Graph data adapter for sectioned storage
2//!
3//! Provides serialization/deserialization of graph data (nodes and edges)
4//! to/from a GRAPH section in a sectioned file.
5
6use anyhow::{Context, Result};
7
8use super::data_structures::{EdgeRec, MetadataRec, NodeRec};
9use super::sectioned::SectionedStorage;
10
11/// Graph data container for serialization
12///
13/// Format compatible with StorageManager file layout:
14/// - Node count (u64 LE)
15/// - Node records (N * 72 bytes)
16/// - Edge count (u64 LE)
17/// - Edge records (M * 48 bytes)
18/// - Metadata count (u64 LE)
19/// - Metadata records (K * 176 bytes)
20#[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    /// Serialize graph data to bytes
29    pub fn to_bytes(&self) -> Vec<u8> {
30        let mut bytes = Vec::with_capacity(self.required_capacity());
31
32        // Node count
33        let node_count: u64 = self.nodes.len() as u64;
34        bytes.extend_from_slice(&node_count.to_le_bytes());
35
36        // Node records (bulk-cast via bytemuck instead of per-element copy)
37        bytes.extend_from_slice(bytemuck::cast_slice(&self.nodes));
38
39        // Edge count
40        let edge_count: u64 = self.edges.len() as u64;
41        bytes.extend_from_slice(&edge_count.to_le_bytes());
42
43        // Edge records (bulk-cast via bytemuck)
44        bytes.extend_from_slice(bytemuck::cast_slice(&self.edges));
45
46        // Metadata count
47        let metadata_count: u64 = self.metadata.len() as u64;
48        bytes.extend_from_slice(&metadata_count.to_le_bytes());
49
50        // Metadata records: Option<MetadataRec> is not Pod, so flatten first.
51        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    /// Deserialize graph data from bytes
62    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        // Read node count
70        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        // Validate we have enough data for nodes
78        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        // Read node records in bulk via bytemuck cast_slice
88        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        // Read edge count
94        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        // Validate we have enough data for edges
105        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        // Read edge records in bulk via bytemuck cast_slice
115        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        // Read metadata count
121        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        // Read metadata records in bulk via bytemuck cast_slice
132        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    /// Calculate required capacity for storing this graph data
148    pub fn required_capacity(&self) -> usize {
149        8 // node count
150            + self.nodes.len() * std::mem::size_of::<NodeRec>()
151            + 8 // edge count
152            + self.edges.len() * std::mem::size_of::<EdgeRec>()
153            + 8 // metadata count
154            + self.metadata.len() * std::mem::size_of::<MetadataRec>()
155    }
156}
157
158/// Graph section adapter
159///
160/// Handles reading/writing graph data from/to a GRAPH section.
161pub struct GraphSectionAdapter;
162
163impl GraphSectionAdapter {
164    pub const SECTION_NAME: &'static str = "GRAPH";
165
166    /// Load graph data from the GRAPH section
167    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    /// Save graph data to the GRAPH section
175    ///
176    /// Handles auto-resizing if capacity is exceeded.
177    pub fn save(storage: &mut SectionedStorage, data: &GraphData) -> Result<()> {
178        let bytes = data.to_bytes();
179        let required = bytes.len() as u64;
180
181        // Check if section exists
182        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                // Check if it's a capacity error
187                if e.to_string().contains("overflow") || e.to_string().contains("capacity") {
188                    // Need to resize - use 2x current capacity or enough for data, whichever is larger
189                    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                    // Retry write
195                    storage.write_section(Self::SECTION_NAME, &bytes)?;
196                } else {
197                    return Err(e);
198                }
199            }
200        } else {
201            // Create new section with reasonable initial capacity (1MB)
202            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    /// Initialize an empty GRAPH section with reasonable default capacity
212    pub fn init(storage: &mut SectionedStorage) -> Result<()> {
213        // Create section with default capacity (1MB) for future growth
214        let default_capacity = 1024 * 1024; // 1MB
215        storage.create_section(Self::SECTION_NAME, default_capacity, 0)?;
216
217        // Write empty data
218        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    /// Check if GRAPH section exists
227    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); // At least node count
296
297        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        // Use actual struct sizes in case they change
312        let expected = 8 // node count
313            + 10 * std::mem::size_of::<NodeRec>()
314            + 8 // edge count
315            + 5 * std::mem::size_of::<EdgeRec>()
316            + 8 // metadata count
317            + 10 * std::mem::size_of::<MetadataRec>();
318        assert_eq!(data.required_capacity(), expected);
319    }
320}