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::new();
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
37        for node in &self.nodes {
38            bytes.extend_from_slice(bytemuck::bytes_of(node));
39        }
40
41        // Edge count
42        let edge_count: u64 = self.edges.len() as u64;
43        bytes.extend_from_slice(&edge_count.to_le_bytes());
44
45        // Edge records
46        for edge in &self.edges {
47            bytes.extend_from_slice(bytemuck::bytes_of(edge));
48        }
49
50        // Metadata count
51        let metadata_count: u64 = self.metadata.len() as u64;
52        bytes.extend_from_slice(&metadata_count.to_le_bytes());
53
54        // Metadata records
55        for meta in &self.metadata {
56            if let Some(m) = meta {
57                bytes.extend_from_slice(bytemuck::bytes_of(m));
58            } else {
59                // Write zeroed metadata record
60                bytes.extend_from_slice(bytemuck::bytes_of(&MetadataRec::default()));
61            }
62        }
63
64        bytes
65    }
66
67    /// Deserialize graph data from bytes
68    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        // Read node count
76        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        // Validate we have enough data for nodes
84        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        // Read node records
94        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        // Read edge count
105        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        // Validate we have enough data for edges
116        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        // Read edge records
126        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        // Read metadata count
137        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        // Read metadata records
148        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    /// Calculate required capacity for storing this graph data
169    pub fn required_capacity(&self) -> usize {
170        8 // node count
171            + self.nodes.len() * std::mem::size_of::<NodeRec>()
172            + 8 // edge count
173            + self.edges.len() * std::mem::size_of::<EdgeRec>()
174            + 8 // metadata count
175            + self.metadata.len() * std::mem::size_of::<MetadataRec>()
176    }
177}
178
179/// Graph section adapter
180///
181/// Handles reading/writing graph data from/to a GRAPH section.
182pub struct GraphSectionAdapter;
183
184impl GraphSectionAdapter {
185    pub const SECTION_NAME: &'static str = "GRAPH";
186
187    /// Load graph data from the GRAPH section
188    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    /// Save graph data to the GRAPH section
196    ///
197    /// Handles auto-resizing if capacity is exceeded.
198    pub fn save(storage: &mut SectionedStorage, data: &GraphData) -> Result<()> {
199        let bytes = data.to_bytes();
200        let required = bytes.len() as u64;
201
202        // Check if section exists
203        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                // Check if it's a capacity error
208                if e.to_string().contains("overflow") || e.to_string().contains("capacity") {
209                    // Need to resize - use 2x current capacity or enough for data, whichever is larger
210                    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                    // Retry write
216                    storage.write_section(Self::SECTION_NAME, &bytes)?;
217                } else {
218                    return Err(e);
219                }
220            }
221        } else {
222            // Create new section with reasonable initial capacity (1MB)
223            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    /// Initialize an empty GRAPH section with reasonable default capacity
233    pub fn init(storage: &mut SectionedStorage) -> Result<()> {
234        // Create section with default capacity (1MB) for future growth
235        let default_capacity = 1024 * 1024; // 1MB
236        storage.create_section(Self::SECTION_NAME, default_capacity, 0)?;
237
238        // Write empty data
239        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    /// Check if GRAPH section exists
248    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); // At least node count
317
318        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        // Use actual struct sizes in case they change
333        let expected = 8 // node count
334            + 10 * std::mem::size_of::<NodeRec>()
335            + 8 // edge count
336            + 5 * std::mem::size_of::<EdgeRec>()
337            + 8 // metadata count
338            + 10 * std::mem::size_of::<MetadataRec>();
339        assert_eq!(data.required_capacity(), expected);
340    }
341}