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