Skip to main content

blazegraph_io_core/graphs/
serialization.rs

1use crate::types::*;
2use anyhow::Result;
3
4impl DocumentGraph {
5    pub fn to_sequential_format(&self) -> SequentialDocument {
6        // Collect all nodes and sort by text_order, with root node first
7        let mut nodes: Vec<&DocumentNode> = self.nodes.values().collect();
8        nodes.sort_by(|a, b| {
9            // Document root (with text_order = None) should come first
10            match (a.text_order, b.text_order) {
11                (None, None) => std::cmp::Ordering::Equal,
12                (None, Some(_)) => std::cmp::Ordering::Less,
13                (Some(_), None) => std::cmp::Ordering::Greater,
14                (Some(a_order), Some(b_order)) => a_order.cmp(&b_order),
15            }
16        });
17
18        let segments: Vec<SequentialSegment> = nodes
19            .into_iter()
20            .enumerate()
21            .map(|(index, node)| SequentialSegment {
22                id: index,
23                node_type: node.node_type.clone(),
24                text: node.content.text.clone(),
25                location: node.location.clone(),
26                style: node.style_info.clone(),
27                tokens: node.token_count,
28            })
29            .collect();
30
31        SequentialDocument {
32            format: "sequential".to_string(),
33            segments,
34            structural_profile: self.structural_profile.clone(),
35        }
36    }
37
38    pub fn to_flat_format(&self) -> FlatDocument {
39        // Collect all nodes and sort by text_order, with root node first
40        let mut nodes: Vec<&DocumentNode> = self.nodes.values().collect();
41        nodes.sort_by(|a, b| {
42            // Document root (with text_order = None) should come first
43            match (a.text_order, b.text_order) {
44                (None, None) => std::cmp::Ordering::Equal,
45                (None, Some(_)) => std::cmp::Ordering::Less,
46                (Some(_), None) => std::cmp::Ordering::Greater,
47                (Some(a_order), Some(b_order)) => a_order.cmp(&b_order),
48            }
49        });
50
51        let chunks: Vec<String> = nodes
52            .into_iter()
53            .map(|node| node.content.text.clone())
54            .collect();
55
56        FlatDocument {
57            format: "flat".to_string(),
58            chunks,
59        }
60    }
61
62    pub fn save_with_format(&self, path: &str, format: &str) -> Result<()> {
63        match format {
64            "sequential" => {
65                let sequential = self.to_sequential_format();
66                let json = serde_json::to_string_pretty(&sequential)?;
67                std::fs::write(path, json)?;
68            }
69            "flat" => {
70                let flat = self.to_flat_format();
71                let json = serde_json::to_string_pretty(&flat)?;
72                std::fs::write(path, json)?;
73            }
74            _ => {
75                self.save_to_json(path)?;
76            }
77        }
78        Ok(())
79    }
80}