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