1use crate::cache::GraphCacheKey;
2use crate::types::PreprocessorOutput;
3use anyhow::{anyhow, Result};
4use sha2::{Digest, Sha256};
5use std::fs;
6use std::path::Path;
7
8#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
17pub enum CachePoint {
18 C0,
20 C1,
22 C2,
24 C3,
26}
27
28impl CachePoint {
29 pub fn all() -> &'static [CachePoint] {
31 &[
32 CachePoint::C0,
33 CachePoint::C1,
34 CachePoint::C2,
35 CachePoint::C3,
36 ]
37 }
38
39 pub fn cascade(&self) -> Vec<CachePoint> {
41 CachePoint::all()
42 .iter()
43 .copied()
44 .filter(|p| p >= self)
45 .collect()
46 }
47
48 pub fn dir_name(&self) -> &'static str {
50 match self {
51 CachePoint::C0 => "c0-pdf",
52 CachePoint::C1 => "c1-xhtml",
53 CachePoint::C2 => "c2-preprocessor",
54 CachePoint::C3 => "c3-graph",
55 }
56 }
57
58 pub fn from_str_with_all(s: &str) -> Result<Option<Self>> {
60 match s.to_lowercase().as_str() {
61 "c0" => Ok(Some(CachePoint::C0)),
62 "c1" => Ok(Some(CachePoint::C1)),
63 "c2" => Ok(Some(CachePoint::C2)),
64 "c3" => Ok(Some(CachePoint::C3)),
65 "all" => Ok(None), _ => Err(anyhow!(
67 "Invalid cache point: '{}'. Use c0, c1, c2, c3, or all",
68 s
69 )),
70 }
71 }
72}
73
74impl std::fmt::Display for CachePoint {
75 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
76 write!(
77 f,
78 "{}",
79 match self {
80 CachePoint::C0 => "C0 (PDF)",
81 CachePoint::C1 => "C1 (XHTML)",
82 CachePoint::C2 => "C2 (Preprocessor)",
83 CachePoint::C3 => "C3 (Graph)",
84 }
85 )
86 }
87}
88
89#[derive(Debug, Clone, Copy, PartialEq, Eq)]
92pub enum FreshFrom {
93 None,
95 C0,
97 C1,
99 C2,
101 C3,
103}
104
105impl FreshFrom {
106 pub fn should_use_cache(&self, point: CachePoint) -> bool {
108 match self {
109 FreshFrom::None => true,
110 FreshFrom::C0 => false,
111 FreshFrom::C1 => point < CachePoint::C1,
112 FreshFrom::C2 => point < CachePoint::C2,
113 FreshFrom::C3 => point < CachePoint::C3,
114 }
115 }
116
117 pub fn parse(s: &str) -> Result<Self> {
119 match s.to_lowercase().as_str() {
120 "c0" => Ok(FreshFrom::C0),
121 "c1" => Ok(FreshFrom::C1),
122 "c2" => Ok(FreshFrom::C2),
123 "c3" => Ok(FreshFrom::C3),
124 _ => Err(anyhow!(
125 "Invalid fresh-from value: '{}'. Use c0, c1, c2, or c3",
126 s
127 )),
128 }
129 }
130}
131
132#[derive(Debug, Clone)]
134pub struct CacheDefaults {
135 pub c0_pdf: bool,
136 pub c1_xhtml: bool,
137 pub c2_preprocessor: bool,
138 pub c3_graph: bool,
139}
140
141impl Default for CacheDefaults {
142 fn default() -> Self {
143 Self {
144 c0_pdf: false,
145 c1_xhtml: true,
146 c2_preprocessor: true,
147 c3_graph: false,
148 }
149 }
150}
151
152impl CacheDefaults {
153 pub fn should_write(&self, point: CachePoint) -> bool {
155 match point {
156 CachePoint::C0 => self.c0_pdf,
157 CachePoint::C1 => self.c1_xhtml,
158 CachePoint::C2 => self.c2_preprocessor,
159 CachePoint::C3 => self.c3_graph,
160 }
161 }
162}
163
164pub struct CacheClearResult {
166 pub deleted: Vec<(CachePoint, usize)>,
167}
168
169pub trait DocumentStorage {
175 fn get_pdf(&self, hash: &str) -> Result<Option<Vec<u8>>>;
177 fn store_pdf(&self, hash: &str, data: &[u8]) -> Result<()>;
178
179 fn get_xhtml(&self, pdf_hash: &str) -> Result<Option<String>>;
181 fn store_xhtml(&self, pdf_hash: &str, xhtml: &str) -> Result<()>;
182
183 fn get_preprocessor_output(&self, pdf_hash: &str) -> Result<Option<PreprocessorOutput>>;
185 fn store_preprocessor_output(&self, pdf_hash: &str, output: &PreprocessorOutput) -> Result<()>;
186
187 fn get_graph_output(
189 &self,
190 cache_key: &GraphCacheKey,
191 ) -> Result<Option<crate::cache::GraphCacheValue>>;
192 fn store_graph_output(
193 &self,
194 cache_key: &GraphCacheKey,
195 cache_value: &crate::cache::GraphCacheValue,
196 ) -> Result<()>;
197
198 fn store_stat(&self, pdf_hash: &str, stat_name: &str, json: &str) -> Result<()>;
205
206 fn clear_cache(&self, from_point: Option<CachePoint>) -> Result<CacheClearResult>;
208}
209
210pub struct FileStorage {
216 cache_dir: String,
217}
218
219impl FileStorage {
220 pub fn new(cache_dir: &str) -> Result<Self> {
221 fs::create_dir_all(cache_dir)?;
222 for point in CachePoint::all() {
223 fs::create_dir_all(format!("{}/{}", cache_dir, point.dir_name()))?;
224 }
225 fs::create_dir_all(format!("{cache_dir}/debug"))?;
226
227 Ok(Self {
228 cache_dir: cache_dir.to_string(),
229 })
230 }
231
232 pub fn cache_dir(&self) -> &str {
233 &self.cache_dir
234 }
235
236 fn pdf_path(&self, hash: &str) -> String {
237 format!("{}/c0-pdf/{}.pdf", self.cache_dir, hash)
238 }
239
240 fn xhtml_path(&self, hash: &str) -> String {
241 format!("{}/c1-xhtml/{}.xhtml", self.cache_dir, hash)
242 }
243
244 fn preprocessor_path(&self, hash: &str) -> String {
245 format!("{}/c2-preprocessor/{}.json", self.cache_dir, hash)
246 }
247
248 fn graph_path(&self, cache_key: &GraphCacheKey) -> String {
249 format!(
250 "{}/c3-graph/{}.json",
251 self.cache_dir,
252 cache_key.to_cache_hash()
253 )
254 }
255
256 fn clear_dir(&self, point: CachePoint) -> Result<usize> {
258 let dir = format!("{}/{}", self.cache_dir, point.dir_name());
259 let path = Path::new(&dir);
260 if !path.exists() {
261 return Ok(0);
262 }
263 let mut count = 0;
264 for entry in fs::read_dir(path)? {
265 let entry = entry?;
266 if entry.file_type()?.is_file() {
267 fs::remove_file(entry.path())?;
268 count += 1;
269 }
270 }
271 Ok(count)
272 }
273}
274
275impl DocumentStorage for FileStorage {
276 fn get_pdf(&self, hash: &str) -> Result<Option<Vec<u8>>> {
278 let path = self.pdf_path(hash);
279 if Path::new(&path).exists() {
280 Ok(Some(fs::read(path)?))
281 } else {
282 Ok(None)
283 }
284 }
285
286 fn store_pdf(&self, hash: &str, data: &[u8]) -> Result<()> {
287 let path = self.pdf_path(hash);
288 fs::write(path, data)?;
289 Ok(())
290 }
291
292 fn get_xhtml(&self, pdf_hash: &str) -> Result<Option<String>> {
294 let path = self.xhtml_path(pdf_hash);
295 if Path::new(&path).exists() {
296 Ok(Some(fs::read_to_string(path)?))
297 } else {
298 Ok(None)
299 }
300 }
301
302 fn store_xhtml(&self, pdf_hash: &str, xhtml: &str) -> Result<()> {
303 let path = self.xhtml_path(pdf_hash);
304 fs::write(path, xhtml)?;
305 Ok(())
306 }
307
308 fn get_preprocessor_output(&self, pdf_hash: &str) -> Result<Option<PreprocessorOutput>> {
310 let path = self.preprocessor_path(pdf_hash);
311 if Path::new(&path).exists() {
312 let json_str = fs::read_to_string(path)?;
313 let output: PreprocessorOutput = serde_json::from_str(&json_str)
314 .map_err(|e| anyhow!("Failed to deserialize cached PreprocessorOutput: {}", e))?;
315 Ok(Some(output))
316 } else {
317 Ok(None)
318 }
319 }
320
321 fn store_preprocessor_output(&self, pdf_hash: &str, output: &PreprocessorOutput) -> Result<()> {
322 let path = self.preprocessor_path(pdf_hash);
323 let json_str = serde_json::to_string_pretty(output)
324 .map_err(|e| anyhow!("Failed to serialize PreprocessorOutput: {}", e))?;
325 fs::write(path, json_str)?;
326 Ok(())
327 }
328
329 fn get_graph_output(
331 &self,
332 cache_key: &GraphCacheKey,
333 ) -> Result<Option<crate::cache::GraphCacheValue>> {
334 let path = self.graph_path(cache_key);
335 if Path::new(&path).exists() {
336 let json_str = fs::read_to_string(path)?;
337 let cache_value: crate::cache::GraphCacheValue = serde_json::from_str(&json_str)
338 .map_err(|e| anyhow!("Failed to deserialize cached GraphCacheValue: {}", e))?;
339 Ok(Some(cache_value))
340 } else {
341 Ok(None)
342 }
343 }
344
345 fn store_graph_output(
346 &self,
347 cache_key: &GraphCacheKey,
348 cache_value: &crate::cache::GraphCacheValue,
349 ) -> Result<()> {
350 let path = self.graph_path(cache_key);
351 let json_str = serde_json::to_string_pretty(cache_value)
352 .map_err(|e| anyhow!("Failed to serialize GraphCacheValue: {}", e))?;
353 fs::write(path, json_str)?;
354 Ok(())
355 }
356
357 fn store_stat(&self, pdf_hash: &str, stat_name: &str, json: &str) -> Result<()> {
359 let dir = format!("{}/stat/{}", self.cache_dir, stat_name);
360 fs::create_dir_all(&dir)?;
361 let path = format!("{dir}/{pdf_hash}.json");
362 fs::write(path, json)?;
363 Ok(())
364 }
365
366 fn clear_cache(&self, from_point: Option<CachePoint>) -> Result<CacheClearResult> {
368 let points_to_clear: Vec<CachePoint> = match from_point {
369 Some(point) => point.cascade(),
370 None => CachePoint::all().to_vec(), };
372
373 let mut deleted = Vec::new();
374 for point in points_to_clear {
375 let count = self.clear_dir(point)?;
376 if count > 0 {
377 deleted.push((point, count));
378 }
379 }
380
381 if from_point.is_none() {
383 let debug_dir = format!("{}/debug", self.cache_dir);
384 if Path::new(&debug_dir).exists() {
385 let mut count = 0;
386 for entry in fs::read_dir(&debug_dir)? {
387 let entry = entry?;
388 if entry.file_type()?.is_file() {
389 fs::remove_file(entry.path())?;
390 count += 1;
391 }
392 }
393 if count > 0 {
394 println!(" Deleted: {} files from debug/", count);
396 }
397 }
398 }
399
400 Ok(CacheClearResult { deleted })
401 }
402}
403
404pub fn calculate_pdf_hash(pdf_bytes: &[u8]) -> String {
410 let chunk_size = 1024; let mut hasher = Sha256::new();
412
413 hasher.update(pdf_bytes.len().to_le_bytes());
415
416 let start_end = std::cmp::min(chunk_size, pdf_bytes.len());
418 hasher.update(&pdf_bytes[0..start_end]);
419
420 if pdf_bytes.len() > chunk_size {
422 let end_start = pdf_bytes.len() - chunk_size;
423 hasher.update(&pdf_bytes[end_start..]);
424 }
425
426 format!("{:x}", hasher.finalize())
427}
428
429pub fn calculate_config_hash<T: serde::Serialize>(config: &T) -> Result<String> {
431 let config_json = serde_json::to_string(config)
432 .map_err(|e| anyhow!("Failed to serialize config for hashing: {}", e))?;
433
434 let mut hasher = Sha256::new();
435 hasher.update(config_json.as_bytes());
436 Ok(format!("{:x}", hasher.finalize()))
437}
438
439pub fn calculate_xhtml_hash(xhtml: &str) -> String {
441 let mut hasher = Sha256::new();
442 hasher.update(xhtml.as_bytes());
443 format!("{:x}", hasher.finalize())
444}
445
446pub struct NoOpStorage;
451
452impl Default for NoOpStorage {
453 fn default() -> Self {
454 Self::new()
455 }
456}
457
458impl NoOpStorage {
459 pub fn new() -> Self {
460 Self
461 }
462}
463
464impl DocumentStorage for NoOpStorage {
465 fn get_pdf(&self, _hash: &str) -> Result<Option<Vec<u8>>> {
466 Ok(None)
467 }
468 fn store_pdf(&self, _hash: &str, _data: &[u8]) -> Result<()> {
469 Ok(())
470 }
471 fn get_xhtml(&self, _pdf_hash: &str) -> Result<Option<String>> {
472 Ok(None)
473 }
474 fn store_xhtml(&self, _pdf_hash: &str, _xhtml: &str) -> Result<()> {
475 Ok(())
476 }
477 fn get_preprocessor_output(&self, _pdf_hash: &str) -> Result<Option<PreprocessorOutput>> {
478 Ok(None)
479 }
480 fn store_preprocessor_output(
481 &self,
482 _pdf_hash: &str,
483 _output: &PreprocessorOutput,
484 ) -> Result<()> {
485 Ok(())
486 }
487 fn get_graph_output(
488 &self,
489 _cache_key: &GraphCacheKey,
490 ) -> Result<Option<crate::cache::GraphCacheValue>> {
491 Ok(None)
492 }
493 fn store_graph_output(
494 &self,
495 _cache_key: &GraphCacheKey,
496 _cache_value: &crate::cache::GraphCacheValue,
497 ) -> Result<()> {
498 Ok(())
499 }
500 fn store_stat(&self, _pdf_hash: &str, _stat_name: &str, _json: &str) -> Result<()> {
501 Ok(())
502 }
503 fn clear_cache(&self, _from_point: Option<CachePoint>) -> Result<CacheClearResult> {
504 Ok(CacheClearResult { deleted: vec![] })
505 }
506}
507
508#[cfg(test)]
513mod tests {
514 use super::*;
515
516 #[test]
517 fn test_pdf_hash_consistency() {
518 let pdf_data = b"test pdf content with some data";
519 let hash1 = calculate_pdf_hash(pdf_data);
520 let hash2 = calculate_pdf_hash(pdf_data);
521 assert_eq!(hash1, hash2);
522 }
523
524 #[test]
525 fn test_pdf_hash_uniqueness() {
526 let pdf1 = b"test pdf content 1";
527 let pdf2 = b"test pdf content 2";
528 let hash1 = calculate_pdf_hash(pdf1);
529 let hash2 = calculate_pdf_hash(pdf2);
530 assert_ne!(hash1, hash2);
531 }
532
533 #[test]
534 fn test_file_storage_roundtrip() {
535 let temp_dir = std::env::temp_dir().join("blazegraph_test_cache_cr11");
536 let _ = std::fs::remove_dir_all(&temp_dir); let storage = FileStorage::new(temp_dir.to_str().unwrap()).unwrap();
538
539 let test_data = b"test pdf data";
540 let hash = "test_hash";
541
542 storage.store_pdf(hash, test_data).unwrap();
544 let retrieved = storage.get_pdf(hash).unwrap();
545 assert_eq!(retrieved, Some(test_data.to_vec()));
546
547 let xhtml = "<html><body>test</body></html>";
549 storage.store_xhtml(hash, xhtml).unwrap();
550 let retrieved_xhtml = storage.get_xhtml(hash).unwrap();
551 assert_eq!(retrieved_xhtml, Some(xhtml.to_string()));
552
553 let xhtml_path = format!("{}/c1-xhtml/{}.xhtml", temp_dir.display(), hash);
555 assert!(Path::new(&xhtml_path).exists());
556
557 std::fs::remove_dir_all(temp_dir).ok();
559 }
560
561 #[test]
562 fn test_cache_point_ordering() {
563 assert!(CachePoint::C0 < CachePoint::C1);
564 assert!(CachePoint::C1 < CachePoint::C2);
565 assert!(CachePoint::C2 < CachePoint::C3);
566 }
567
568 #[test]
569 fn test_cache_point_cascade() {
570 assert_eq!(
571 CachePoint::C0.cascade(),
572 vec![
573 CachePoint::C0,
574 CachePoint::C1,
575 CachePoint::C2,
576 CachePoint::C3
577 ]
578 );
579 assert_eq!(
580 CachePoint::C1.cascade(),
581 vec![CachePoint::C1, CachePoint::C2, CachePoint::C3]
582 );
583 assert_eq!(
584 CachePoint::C2.cascade(),
585 vec![CachePoint::C2, CachePoint::C3]
586 );
587 assert_eq!(CachePoint::C3.cascade(), vec![CachePoint::C3]);
588 }
589
590 #[test]
591 fn test_fresh_from_cache_bypass() {
592 assert!(FreshFrom::None.should_use_cache(CachePoint::C0));
594 assert!(FreshFrom::None.should_use_cache(CachePoint::C3));
595
596 assert!(!FreshFrom::C0.should_use_cache(CachePoint::C0));
598 assert!(!FreshFrom::C0.should_use_cache(CachePoint::C3));
599
600 assert!(FreshFrom::C2.should_use_cache(CachePoint::C0));
602 assert!(FreshFrom::C2.should_use_cache(CachePoint::C1));
603 assert!(!FreshFrom::C2.should_use_cache(CachePoint::C2));
604 assert!(!FreshFrom::C2.should_use_cache(CachePoint::C3));
605 }
606
607 #[test]
608 fn test_clear_cache_cascade() {
609 let temp_dir = std::env::temp_dir().join("blazegraph_test_clear_cr11");
610 let _ = std::fs::remove_dir_all(&temp_dir);
611 let storage = FileStorage::new(temp_dir.to_str().unwrap()).unwrap();
612
613 storage.store_xhtml("hash1", "<html>test</html>").unwrap();
615 storage.store_xhtml("hash2", "<html>test2</html>").unwrap();
616
617 let result = storage.clear_cache(Some(CachePoint::C1)).unwrap();
619 assert!(result
620 .deleted
621 .iter()
622 .any(|(p, c)| *p == CachePoint::C1 && *c == 2));
623
624 assert!(storage.get_xhtml("hash1").unwrap().is_none());
626 assert!(storage.get_xhtml("hash2").unwrap().is_none());
627
628 std::fs::remove_dir_all(temp_dir).ok();
629 }
630
631 #[test]
632 fn test_cache_defaults() {
633 let defaults = CacheDefaults::default();
634 assert!(!defaults.should_write(CachePoint::C0));
635 assert!(defaults.should_write(CachePoint::C1));
636 assert!(defaults.should_write(CachePoint::C2));
637 assert!(!defaults.should_write(CachePoint::C3));
638 }
639}