1use std::fmt::{Display, Formatter};
2use std::path::Path;
3
4use graphyn_core::graph::GraphynGraph;
5use graphyn_core::ir::{
6 Language, ReExportEntry, Relationship, RelationshipKind, Symbol, SymbolKind,
7};
8use graphyn_core::resolver::{AliasEntry, AliasScope};
9use rocksdb::{Options, DB};
10
11const KEY_GRAPH_SNAPSHOT: &[u8] = b"graph_snapshot_v1";
12const SNAPSHOT_VERSION: u8 = 2;
13
14#[derive(Debug)]
15pub enum StoreError {
16 RocksDb(String),
17 Serialization(String),
18 SnapshotNotFound,
19}
20
21impl Display for StoreError {
22 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
23 match self {
24 Self::RocksDb(err) => write!(f, "rocksdb error: {err}"),
25 Self::Serialization(err) => write!(f, "serialization error: {err}"),
26 Self::SnapshotNotFound => write!(f, "snapshot not found"),
27 }
28 }
29}
30
31impl std::error::Error for StoreError {}
32
33#[derive(Debug, Clone)]
34pub struct GraphSnapshot {
35 pub symbols: Vec<Symbol>,
36 pub relationships: Vec<Relationship>,
37 pub alias_chains: Vec<(String, Vec<AliasEntry>)>,
38 pub file_reexports: Vec<(String, Vec<ReExportEntry>)>,
39}
40
41pub struct RocksGraphStore {
42 db: DB,
43}
44
45impl RocksGraphStore {
46 pub fn open(path: &Path) -> Result<Self, StoreError> {
47 let mut options = Options::default();
48 options.create_if_missing(true);
49 let db = DB::open(&options, path).map_err(|err| StoreError::RocksDb(err.to_string()))?;
50 Ok(Self { db })
51 }
52
53 pub fn save_graph(&self, graph: &GraphynGraph) -> Result<(), StoreError> {
54 let snapshot = GraphSnapshot::from_graph(graph)?;
55 self.save_snapshot(&snapshot)
56 }
57
58 pub fn load_graph(&self) -> Result<GraphynGraph, StoreError> {
59 let snapshot = self.load_snapshot()?;
60 snapshot.into_graph()
61 }
62
63 pub fn save_snapshot(&self, snapshot: &GraphSnapshot) -> Result<(), StoreError> {
64 let bytes = snapshot.to_bytes()?;
65 self.db
66 .put(KEY_GRAPH_SNAPSHOT, bytes)
67 .map_err(|err| StoreError::RocksDb(err.to_string()))
68 }
69
70 pub fn load_snapshot(&self) -> Result<GraphSnapshot, StoreError> {
71 let bytes = self
72 .db
73 .get(KEY_GRAPH_SNAPSHOT)
74 .map_err(|err| StoreError::RocksDb(err.to_string()))?
75 .ok_or(StoreError::SnapshotNotFound)?;
76
77 GraphSnapshot::from_bytes(&bytes)
78 }
79}
80
81impl GraphSnapshot {
82 pub fn from_graph(graph: &GraphynGraph) -> Result<Self, StoreError> {
83 let mut symbols: Vec<Symbol> = graph
84 .symbols
85 .iter()
86 .map(|entry| entry.value().clone())
87 .collect();
88 symbols.sort_by(|a, b| a.id.cmp(&b.id));
89
90 let mut relationships = Vec::new();
91 for edge_id in graph.graph.edge_indices() {
92 let (source_idx, target_idx) = graph
93 .graph
94 .edge_endpoints(edge_id)
95 .ok_or_else(|| StoreError::Serialization("missing edge endpoints".to_string()))?;
96 let from = graph
97 .graph
98 .node_weight(source_idx)
99 .cloned()
100 .ok_or_else(|| StoreError::Serialization("missing source node".to_string()))?;
101 let to = graph
102 .graph
103 .node_weight(target_idx)
104 .cloned()
105 .ok_or_else(|| StoreError::Serialization("missing target node".to_string()))?;
106 let meta = graph
107 .graph
108 .edge_weight(edge_id)
109 .ok_or_else(|| StoreError::Serialization("missing edge metadata".to_string()))?;
110
111 relationships.push(Relationship {
112 from,
113 to,
114 kind: meta.kind.clone(),
115 alias: meta.alias.clone(),
116 properties_accessed: meta.properties_accessed.clone(),
117 context: meta.context.clone(),
118 file: meta.file.clone(),
119 line: meta.line,
120 });
121 }
122 relationships.sort_by(|a, b| {
123 a.file
124 .cmp(&b.file)
125 .then(a.line.cmp(&b.line))
126 .then(a.from.cmp(&b.from))
127 .then(a.to.cmp(&b.to))
128 });
129
130 let mut alias_chains: Vec<(String, Vec<AliasEntry>)> = graph
131 .alias_chains
132 .iter()
133 .map(|entry| {
134 let mut aliases = entry.value().clone();
135 aliases.sort_by(|a, b| {
136 a.defined_in_file
137 .cmp(&b.defined_in_file)
138 .then(a.alias_name.cmp(&b.alias_name))
139 });
140 (entry.key().clone(), aliases)
141 })
142 .collect();
143 alias_chains.sort_by(|a, b| a.0.cmp(&b.0));
144
145 let mut file_reexports: Vec<(String, Vec<ReExportEntry>)> = graph
146 .file_reexports
147 .iter()
148 .map(|entry| {
149 let mut re_exports = entry.value().clone();
150 re_exports.sort_by(|a, b| {
151 a.exported_name
152 .cmp(&b.exported_name)
153 .then(a.source_module.cmp(&b.source_module))
154 });
155 (entry.key().clone(), re_exports)
156 })
157 .collect();
158 file_reexports.sort_by(|a, b| a.0.cmp(&b.0));
159
160 Ok(Self {
161 symbols,
162 relationships,
163 alias_chains,
164 file_reexports,
165 })
166 }
167
168 pub fn into_graph(self) -> Result<GraphynGraph, StoreError> {
169 let mut graph = GraphynGraph::new();
170
171 for symbol in self.symbols {
172 graph.add_symbol(symbol);
173 }
174
175 for relationship in &self.relationships {
176 graph.add_relationship(relationship);
177 }
178
179 for (canonical_id, aliases) in self.alias_chains {
180 graph.alias_chains.insert(canonical_id, aliases);
181 }
182
183 for (file, re_exports) in self.file_reexports {
184 graph.file_reexports.insert(file, re_exports);
185 }
186
187 Ok(graph)
188 }
189
190 fn to_bytes(&self) -> Result<Vec<u8>, StoreError> {
191 let mut out = Vec::new();
192
193 write_u8(&mut out, SNAPSHOT_VERSION);
194
195 write_u32(&mut out, self.symbols.len() as u32);
196 for symbol in &self.symbols {
197 write_string(&mut out, &symbol.id)?;
198 write_string(&mut out, &symbol.name)?;
199 write_u8(&mut out, symbol_kind_to_u8(&symbol.kind));
200 write_u8(&mut out, language_to_u8(&symbol.language));
201 write_string(&mut out, &symbol.file)?;
202 write_u32(&mut out, symbol.line_start);
203 write_u32(&mut out, symbol.line_end);
204 write_optional_string(&mut out, symbol.signature.as_deref())?;
205 }
206
207 write_u32(&mut out, self.relationships.len() as u32);
208 for relationship in &self.relationships {
209 write_string(&mut out, &relationship.from)?;
210 write_string(&mut out, &relationship.to)?;
211 write_u8(&mut out, relationship_kind_to_u8(&relationship.kind));
212 write_optional_string(&mut out, relationship.alias.as_deref())?;
213 write_u32(&mut out, relationship.properties_accessed.len() as u32);
214 for prop in &relationship.properties_accessed {
215 write_string(&mut out, prop)?;
216 }
217 write_string(&mut out, &relationship.context)?;
218 write_string(&mut out, &relationship.file)?;
219 write_u32(&mut out, relationship.line);
220 }
221
222 write_u32(&mut out, self.alias_chains.len() as u32);
223 for (canonical, entries) in &self.alias_chains {
224 write_string(&mut out, canonical)?;
225 write_u32(&mut out, entries.len() as u32);
226 for entry in entries {
227 write_string(&mut out, &entry.alias_name)?;
228 write_string(&mut out, &entry.defined_in_file)?;
229 write_u8(&mut out, alias_scope_to_u8(&entry.scope));
230 }
231 }
232
233 write_u32(&mut out, self.file_reexports.len() as u32);
234 for (file, entries) in &self.file_reexports {
235 write_string(&mut out, file)?;
236 write_u32(&mut out, entries.len() as u32);
237 for entry in entries {
238 write_string(&mut out, &entry.exported_name)?;
239 write_string(&mut out, &entry.source_module)?;
240 }
241 }
242
243 Ok(out)
244 }
245
246 fn from_bytes(bytes: &[u8]) -> Result<Self, StoreError> {
247 let mut cursor = ByteCursor::new(bytes);
248
249 let version = cursor.read_u8()?;
250 if version != 1 && version != SNAPSHOT_VERSION {
251 return Err(StoreError::Serialization(format!(
252 "unsupported snapshot version: {version}"
253 )));
254 }
255
256 let symbol_count = cursor.read_u32()? as usize;
257 let mut symbols = Vec::with_capacity(symbol_count);
258 for _ in 0..symbol_count {
259 symbols.push(Symbol {
260 id: cursor.read_string()?,
261 name: cursor.read_string()?,
262 kind: u8_to_symbol_kind(cursor.read_u8()?)?,
263 language: u8_to_language(cursor.read_u8()?)?,
264 file: cursor.read_string()?,
265 line_start: cursor.read_u32()?,
266 line_end: cursor.read_u32()?,
267 signature: cursor.read_optional_string()?,
268 });
269 }
270
271 let rel_count = cursor.read_u32()? as usize;
272 let mut relationships = Vec::with_capacity(rel_count);
273 for _ in 0..rel_count {
274 let from = cursor.read_string()?;
275 let to = cursor.read_string()?;
276 let kind = u8_to_relationship_kind(cursor.read_u8()?)?;
277 let alias = cursor.read_optional_string()?;
278 let prop_count = cursor.read_u32()? as usize;
279 let mut properties_accessed = Vec::with_capacity(prop_count);
280 for _ in 0..prop_count {
281 properties_accessed.push(cursor.read_string()?);
282 }
283 let context = cursor.read_string()?;
284 let file = cursor.read_string()?;
285 let line = cursor.read_u32()?;
286
287 relationships.push(Relationship {
288 from,
289 to,
290 kind,
291 alias,
292 properties_accessed,
293 context,
294 file,
295 line,
296 });
297 }
298
299 let alias_chain_count = cursor.read_u32()? as usize;
300 let mut alias_chains = Vec::with_capacity(alias_chain_count);
301 for _ in 0..alias_chain_count {
302 let canonical = cursor.read_string()?;
303 let entry_count = cursor.read_u32()? as usize;
304 let mut entries = Vec::with_capacity(entry_count);
305 for _ in 0..entry_count {
306 entries.push(AliasEntry {
307 alias_name: cursor.read_string()?,
308 defined_in_file: cursor.read_string()?,
309 scope: u8_to_alias_scope(cursor.read_u8()?)?,
310 });
311 }
312 alias_chains.push((canonical, entries));
313 }
314
315 let mut file_reexports = Vec::new();
316 if version >= 2 {
317 let reexport_file_count = cursor.read_u32()? as usize;
318 file_reexports.reserve(reexport_file_count);
319 for _ in 0..reexport_file_count {
320 let file = cursor.read_string()?;
321 let entry_count = cursor.read_u32()? as usize;
322 let mut entries = Vec::with_capacity(entry_count);
323 for _ in 0..entry_count {
324 entries.push(ReExportEntry {
325 exported_name: cursor.read_string()?,
326 source_module: cursor.read_string()?,
327 });
328 }
329 file_reexports.push((file, entries));
330 }
331 }
332
333 if !cursor.is_at_end() {
334 return Err(StoreError::Serialization(
335 "trailing bytes found in snapshot".to_string(),
336 ));
337 }
338
339 Ok(Self {
340 symbols,
341 relationships,
342 alias_chains,
343 file_reexports,
344 })
345 }
346}
347
348struct ByteCursor<'a> {
349 bytes: &'a [u8],
350 pos: usize,
351}
352
353impl<'a> ByteCursor<'a> {
354 fn new(bytes: &'a [u8]) -> Self {
355 Self { bytes, pos: 0 }
356 }
357
358 fn read_u8(&mut self) -> Result<u8, StoreError> {
359 if self.pos >= self.bytes.len() {
360 return Err(StoreError::Serialization(
361 "unexpected EOF reading u8".to_string(),
362 ));
363 }
364 let v = self.bytes[self.pos];
365 self.pos += 1;
366 Ok(v)
367 }
368
369 fn read_u32(&mut self) -> Result<u32, StoreError> {
370 if self.pos + 4 > self.bytes.len() {
371 return Err(StoreError::Serialization(
372 "unexpected EOF reading u32".to_string(),
373 ));
374 }
375 let mut arr = [0u8; 4];
376 arr.copy_from_slice(&self.bytes[self.pos..self.pos + 4]);
377 self.pos += 4;
378 Ok(u32::from_le_bytes(arr))
379 }
380
381 fn read_string(&mut self) -> Result<String, StoreError> {
382 let len = self.read_u32()? as usize;
383 if self.pos + len > self.bytes.len() {
384 return Err(StoreError::Serialization(
385 "unexpected EOF reading string".to_string(),
386 ));
387 }
388 let slice = &self.bytes[self.pos..self.pos + len];
389 self.pos += len;
390 String::from_utf8(slice.to_vec())
391 .map_err(|err| StoreError::Serialization(format!("invalid UTF-8 string: {err}")))
392 }
393
394 fn read_optional_string(&mut self) -> Result<Option<String>, StoreError> {
395 let has = self.read_u8()?;
396 if has == 0 {
397 Ok(None)
398 } else {
399 Ok(Some(self.read_string()?))
400 }
401 }
402
403 fn is_at_end(&self) -> bool {
404 self.pos == self.bytes.len()
405 }
406}
407
408fn write_u8(out: &mut Vec<u8>, value: u8) {
409 out.push(value);
410}
411
412fn write_u32(out: &mut Vec<u8>, value: u32) {
413 out.extend_from_slice(&value.to_le_bytes());
414}
415
416fn write_string(out: &mut Vec<u8>, value: &str) -> Result<(), StoreError> {
417 let bytes = value.as_bytes();
418 let len = u32::try_from(bytes.len())
419 .map_err(|_| StoreError::Serialization("string too large".to_string()))?;
420 write_u32(out, len);
421 out.extend_from_slice(bytes);
422 Ok(())
423}
424
425fn write_optional_string(out: &mut Vec<u8>, value: Option<&str>) -> Result<(), StoreError> {
426 match value {
427 Some(value) => {
428 write_u8(out, 1);
429 write_string(out, value)
430 }
431 None => {
432 write_u8(out, 0);
433 Ok(())
434 }
435 }
436}
437
438fn symbol_kind_to_u8(kind: &SymbolKind) -> u8 {
439 match kind {
440 SymbolKind::Class => 1,
441 SymbolKind::Interface => 2,
442 SymbolKind::TypeAlias => 3,
443 SymbolKind::Function => 4,
444 SymbolKind::Method => 5,
445 SymbolKind::Property => 6,
446 SymbolKind::Variable => 7,
447 SymbolKind::Module => 8,
448 SymbolKind::Enum => 9,
449 SymbolKind::EnumVariant => 10,
450 SymbolKind::ExternalPackage => 11,
451 }
452}
453
454fn u8_to_symbol_kind(input: u8) -> Result<SymbolKind, StoreError> {
455 match input {
456 1 => Ok(SymbolKind::Class),
457 2 => Ok(SymbolKind::Interface),
458 3 => Ok(SymbolKind::TypeAlias),
459 4 => Ok(SymbolKind::Function),
460 5 => Ok(SymbolKind::Method),
461 6 => Ok(SymbolKind::Property),
462 7 => Ok(SymbolKind::Variable),
463 8 => Ok(SymbolKind::Module),
464 9 => Ok(SymbolKind::Enum),
465 10 => Ok(SymbolKind::EnumVariant),
466 11 => Ok(SymbolKind::ExternalPackage),
467 other => Err(StoreError::Serialization(format!(
468 "unknown symbol kind code: {other}"
469 ))),
470 }
471}
472
473fn language_to_u8(language: &Language) -> u8 {
474 match language {
475 Language::TypeScript => 1,
476 Language::JavaScript => 2,
477 Language::Python => 3,
478 Language::Rust => 4,
479 Language::Go => 5,
480 Language::Java => 6,
481 }
482}
483
484fn u8_to_language(input: u8) -> Result<Language, StoreError> {
485 match input {
486 1 => Ok(Language::TypeScript),
487 2 => Ok(Language::JavaScript),
488 3 => Ok(Language::Python),
489 4 => Ok(Language::Rust),
490 5 => Ok(Language::Go),
491 6 => Ok(Language::Java),
492 other => Err(StoreError::Serialization(format!(
493 "unknown language code: {other}"
494 ))),
495 }
496}
497
498fn relationship_kind_to_u8(kind: &RelationshipKind) -> u8 {
499 match kind {
500 RelationshipKind::Imports => 1,
501 RelationshipKind::Calls => 2,
502 RelationshipKind::Extends => 3,
503 RelationshipKind::Implements => 4,
504 RelationshipKind::UsesType => 5,
505 RelationshipKind::AccessesProperty => 6,
506 RelationshipKind::ReExports => 7,
507 RelationshipKind::Instantiates => 8,
508 }
509}
510
511fn u8_to_relationship_kind(input: u8) -> Result<RelationshipKind, StoreError> {
512 match input {
513 1 => Ok(RelationshipKind::Imports),
514 2 => Ok(RelationshipKind::Calls),
515 3 => Ok(RelationshipKind::Extends),
516 4 => Ok(RelationshipKind::Implements),
517 5 => Ok(RelationshipKind::UsesType),
518 6 => Ok(RelationshipKind::AccessesProperty),
519 7 => Ok(RelationshipKind::ReExports),
520 8 => Ok(RelationshipKind::Instantiates),
521 other => Err(StoreError::Serialization(format!(
522 "unknown relationship kind code: {other}"
523 ))),
524 }
525}
526
527fn alias_scope_to_u8(scope: &AliasScope) -> u8 {
528 match scope {
529 AliasScope::ImportAlias => 1,
530 AliasScope::ReExport => 2,
531 AliasScope::BarrelReExport => 3,
532 AliasScope::DefaultImport => 4,
533 }
534}
535
536fn u8_to_alias_scope(input: u8) -> Result<AliasScope, StoreError> {
537 match input {
538 1 => Ok(AliasScope::ImportAlias),
539 2 => Ok(AliasScope::ReExport),
540 3 => Ok(AliasScope::BarrelReExport),
541 4 => Ok(AliasScope::DefaultImport),
542 other => Err(StoreError::Serialization(format!(
543 "unknown alias scope code: {other}"
544 ))),
545 }
546}