1use std::collections::HashMap;
10
11use rusqlite::params;
12use serde::{Deserialize, Serialize};
13
14use super::{CodeGraph, Edge, EdgeKind, Node, NodeKind};
15
16pub const SNAPSHOT_VERSION: u32 = 1;
17
18#[derive(Debug, Serialize, Deserialize)]
19struct SnapshotHeader {
20 leanctx_graph_snapshot: u32,
21 nodes: usize,
22 edges: usize,
23}
24
25#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord)]
27struct SnapNode {
28 kind: String,
29 file: String,
30 name: String,
31 #[serde(skip_serializing_if = "Option::is_none")]
32 line_start: Option<usize>,
33 #[serde(skip_serializing_if = "Option::is_none")]
34 line_end: Option<usize>,
35 #[serde(skip_serializing_if = "Option::is_none")]
36 meta: Option<String>,
37}
38
39#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord)]
41struct SnapEdge {
42 kind: String,
43 source: (String, String, String),
44 target: (String, String, String),
45 #[serde(skip_serializing_if = "Option::is_none")]
46 meta: Option<String>,
47}
48
49#[derive(Debug, Default)]
50pub struct ImportStats {
51 pub nodes: usize,
52 pub edges: usize,
53 pub skipped_edges: usize,
54}
55
56#[derive(Debug, Default)]
57pub struct DriftReport {
58 pub only_local: usize,
59 pub only_snapshot: usize,
60 pub common: usize,
61}
62
63impl DriftReport {
64 pub fn in_sync(&self) -> bool {
65 self.only_local == 0 && self.only_snapshot == 0
66 }
67}
68
69fn collect_nodes(graph: &CodeGraph) -> anyhow::Result<Vec<SnapNode>> {
70 let conn = graph.connection();
71 let mut stmt =
72 conn.prepare("SELECT n.kind, p.path, n.name, n.line_start, n.line_end, n.metadata FROM nodes n JOIN paths p ON p.id = n.file_id")?;
73 let mut nodes: Vec<SnapNode> = stmt
74 .query_map([], |row| {
75 Ok(SnapNode {
76 kind: row.get(0)?,
77 file: row.get(1)?,
78 name: row.get(2)?,
79 line_start: row.get::<_, Option<i64>>(3)?.map(|v| v as usize),
80 line_end: row.get::<_, Option<i64>>(4)?.map(|v| v as usize),
81 meta: row.get(5)?,
82 })
83 })?
84 .collect::<Result<_, _>>()?;
85 nodes.sort();
86 nodes.retain(|n| graph.serves_node_parts(&n.kind, &n.file, &n.name, n.meta.as_deref()));
90 Ok(nodes)
91}
92
93fn collect_edges(graph: &CodeGraph, nodes: &[SnapNode]) -> anyhow::Result<Vec<SnapEdge>> {
96 let served: std::collections::HashSet<(&str, &str, &str)> = nodes
97 .iter()
98 .map(|n| (n.kind.as_str(), n.file.as_str(), n.name.as_str()))
99 .collect();
100 let is_served = |end: &(String, String, String)| served.contains(&(&*end.0, &*end.1, &*end.2));
101 let conn = graph.connection();
102 let mut stmt = conn.prepare(
103 "SELECT e.kind, e.metadata,
104 s.kind, sp.path, s.name,
105 t.kind, tp.path, t.name
106 FROM edges e
107 JOIN nodes s ON s.id = e.source_id
108 JOIN nodes t ON t.id = e.target_id
109 JOIN paths sp ON sp.id = s.file_id
110 JOIN paths tp ON tp.id = t.file_id",
111 )?;
112 let mut edges: Vec<SnapEdge> = stmt
113 .query_map([], |row| {
114 Ok(SnapEdge {
115 kind: row.get(0)?,
116 meta: row.get(1)?,
117 source: (row.get(2)?, row.get(3)?, row.get(4)?),
118 target: (row.get(5)?, row.get(6)?, row.get(7)?),
119 })
120 })?
121 .collect::<Result<_, _>>()?;
122 edges.retain(|e| is_served(&e.source) && is_served(&e.target));
123 edges.sort();
124 Ok(edges)
125}
126
127pub fn export_snapshot(graph: &CodeGraph) -> anyhow::Result<String> {
132 let nodes = collect_nodes(graph)?;
133 let edges = collect_edges(graph, &nodes)?;
134
135 let mut out = String::new();
136 out.push_str(&serde_json::to_string(&SnapshotHeader {
137 leanctx_graph_snapshot: SNAPSHOT_VERSION,
138 nodes: nodes.len(),
139 edges: edges.len(),
140 })?);
141 out.push('\n');
142 for n in &nodes {
143 out.push_str(&format!("{{\"n\":{}}}\n", serde_json::to_string(n)?));
144 }
145 for e in &edges {
146 out.push_str(&format!("{{\"e\":{}}}\n", serde_json::to_string(e)?));
147 }
148 Ok(out)
149}
150
151pub fn import_snapshot(graph: &CodeGraph, content: &str) -> anyhow::Result<ImportStats> {
155 let (nodes, edges) = parse_snapshot(content)?;
156 let mut stats = ImportStats::default();
157 let mut id_by_identity: HashMap<(String, String, String), i64> = HashMap::new();
158
159 for n in &nodes {
160 let node = Node {
161 id: None,
162 kind: NodeKind::parse(&n.kind),
163 name: n.name.clone(),
164 file_path: n.file.clone(),
165 line_start: n.line_start,
166 line_end: n.line_end,
167 metadata: n.meta.clone(),
168 };
169 let id = graph.upsert_node(&node)?;
170 id_by_identity.insert((n.kind.clone(), n.file.clone(), n.name.clone()), id);
171 stats.nodes += 1;
172 }
173
174 for e in &edges {
175 let source = resolve_endpoint(graph, &mut id_by_identity, &e.source);
176 let target = resolve_endpoint(graph, &mut id_by_identity, &e.target);
177 match (source, target) {
178 (Some(s), Some(t)) => {
179 graph.upsert_edge(&Edge {
180 id: None,
181 source_id: s,
182 target_id: t,
183 kind: EdgeKind::parse(&e.kind),
184 metadata: e.meta.clone(),
185 })?;
186 stats.edges += 1;
187 }
188 _ => stats.skipped_edges += 1,
189 }
190 }
191
192 Ok(stats)
193}
194
195pub fn check_snapshot(graph: &CodeGraph, content: &str) -> anyhow::Result<DriftReport> {
197 let local = export_snapshot(graph)?;
198 let local_set: std::collections::HashSet<&str> =
199 local.lines().skip(1).filter(|l| !l.is_empty()).collect();
200 let snap_set: std::collections::HashSet<&str> =
201 content.lines().skip(1).filter(|l| !l.is_empty()).collect();
202
203 Ok(DriftReport {
204 only_local: local_set.difference(&snap_set).count(),
205 only_snapshot: snap_set.difference(&local_set).count(),
206 common: local_set.intersection(&snap_set).count(),
207 })
208}
209
210fn resolve_endpoint(
211 graph: &CodeGraph,
212 cache: &mut HashMap<(String, String, String), i64>,
213 identity: &(String, String, String),
214) -> Option<i64> {
215 if let Some(&id) = cache.get(identity) {
216 return Some(id);
217 }
218 let conn = graph.connection();
220 let found: Option<i64> = conn
221 .query_row(
222 "SELECT n.id FROM nodes n JOIN paths p ON p.id = n.file_id
223 WHERE n.kind = ?1 AND p.path = ?2 AND n.name = ?3",
224 params![identity.0, identity.1, identity.2],
225 |row| row.get(0),
226 )
227 .ok();
228 if let Some(id) = found {
229 cache.insert(identity.clone(), id);
230 }
231 found
232}
233
234fn parse_snapshot(content: &str) -> anyhow::Result<(Vec<SnapNode>, Vec<SnapEdge>)> {
235 let mut lines = content.lines().filter(|l| !l.trim().is_empty());
236 let header_line = lines
237 .next()
238 .ok_or_else(|| anyhow::anyhow!("empty snapshot"))?;
239 let header: SnapshotHeader = serde_json::from_str(header_line)
240 .map_err(|e| anyhow::anyhow!("invalid snapshot header: {e}"))?;
241 if header.leanctx_graph_snapshot != SNAPSHOT_VERSION {
242 anyhow::bail!(
243 "unsupported snapshot version {} (supported: {SNAPSHOT_VERSION})",
244 header.leanctx_graph_snapshot
245 );
246 }
247
248 let mut nodes = Vec::new();
249 let mut edges = Vec::new();
250 for line in lines {
251 let v: serde_json::Value = serde_json::from_str(line)
252 .map_err(|e| anyhow::anyhow!("invalid snapshot line: {e}"))?;
253 if let Some(n) = v.get("n") {
254 nodes.push(serde_json::from_value::<SnapNode>(n.clone())?);
255 } else if let Some(e) = v.get("e") {
256 edges.push(serde_json::from_value::<SnapEdge>(e.clone())?);
257 } else {
258 anyhow::bail!("unknown snapshot line shape: {line}");
259 }
260 }
261 Ok((nodes, edges))
262}
263
264#[cfg(test)]
265mod tests {
266 use super::*;
267
268 fn graph_with_data() -> CodeGraph {
269 let g = CodeGraph::open_in_memory().unwrap();
270 let a = g.upsert_node(&Node::file("src/auth.rs")).expect("node a");
271 let b = g.upsert_node(&Node::file("src/db.rs")).expect("node b");
272 g.upsert_edge(&Edge::new(a, b, EdgeKind::Imports))
273 .expect("edge");
274 g
275 }
276
277 #[test]
278 fn export_is_deterministic_and_id_free() {
279 let g = graph_with_data();
280 let s1 = export_snapshot(&g).unwrap();
281 let s2 = export_snapshot(&g).unwrap();
282 assert_eq!(s1, s2);
283 assert!(!s1.contains("\"id\""), "snapshot must not leak local ids");
284 assert!(s1.starts_with("{\"leanctx_graph_snapshot\":1"));
285 }
286
287 #[test]
288 fn roundtrip_export_import_is_lossless() {
289 let g = graph_with_data();
290 let snapshot = export_snapshot(&g).unwrap();
291
292 let fresh = CodeGraph::open_in_memory().unwrap();
293 let stats = import_snapshot(&fresh, &snapshot).unwrap();
294 assert_eq!(stats.nodes, 2);
295 assert_eq!(stats.edges, 1);
296 assert_eq!(stats.skipped_edges, 0);
297
298 let reexported = export_snapshot(&fresh).unwrap();
299 assert_eq!(snapshot, reexported, "roundtrip must be lossless");
300 }
301
302 #[test]
303 fn import_merges_instead_of_replacing() {
304 let g = graph_with_data();
305 let snapshot = export_snapshot(&g).unwrap();
306
307 let local = CodeGraph::open_in_memory().unwrap();
308 local
309 .upsert_node(&Node::file("src/local_only.rs"))
310 .expect("local node");
311
312 import_snapshot(&local, &snapshot).unwrap();
313 let merged = export_snapshot(&local).unwrap();
314 assert!(merged.contains("local_only.rs"), "local data must survive");
315 assert!(merged.contains("auth.rs"), "snapshot data must be merged");
316 }
317
318 #[test]
319 fn check_reports_drift_and_sync() {
320 let g = graph_with_data();
321 let snapshot = export_snapshot(&g).unwrap();
322
323 let synced = check_snapshot(&g, &snapshot).unwrap();
324 assert!(synced.in_sync());
325
326 g.upsert_node(&Node::file("src/new_file.rs")).unwrap();
327 let drifted = check_snapshot(&g, &snapshot).unwrap();
328 assert!(!drifted.in_sync());
329 assert_eq!(drifted.only_local, 1);
330 assert_eq!(drifted.only_snapshot, 0);
331 }
332
333 #[test]
334 fn rejects_wrong_version() {
335 let g = CodeGraph::open_in_memory().unwrap();
336 let err = import_snapshot(
337 &g,
338 "{\"leanctx_graph_snapshot\":99,\"nodes\":0,\"edges\":0}\n",
339 )
340 .unwrap_err();
341 assert!(err.to_string().contains("unsupported snapshot version"));
342 }
343}