Skip to main content

xz_knowledge_graph/traversal/
cache.rs

1use std::collections::HashMap;
2use std::sync::RwLock;
3
4use crate::types::entity::Entity;
5use crate::types::graph::{PathStep, SubGraph};
6use crate::types::relation::Relation;
7
8/// Simple in-memory traversal cache.
9#[derive(Debug, Default)]
10pub struct TraversalCache {
11    neighbors: RwLock<HashMap<String, SubGraph>>,
12    paths: RwLock<HashMap<(String, String), Vec<Vec<PathStep>>>>,
13    entities: RwLock<HashMap<String, Entity>>,
14    relations: RwLock<HashMap<String, Vec<Relation>>>,
15}
16
17impl TraversalCache {
18    pub fn new() -> Self {
19        Self::default()
20    }
21
22    pub fn get_neighbors(&self, entity_id: &str) -> Option<SubGraph> {
23        self.neighbors.read().unwrap().get(entity_id).cloned()
24    }
25
26    pub fn put_neighbors(&self, entity_id: String, subgraph: SubGraph) {
27        self.neighbors.write().unwrap().insert(entity_id, subgraph);
28    }
29
30    pub fn get_path(&self, from: &str, to: &str) -> Option<Vec<Vec<PathStep>>> {
31        self.paths.read().unwrap().get(&(from.to_string(), to.to_string())).cloned()
32    }
33
34    pub fn put_path(&self, from: String, to: String, paths: Vec<Vec<PathStep>>) {
35        self.paths.write().unwrap().insert((from, to), paths);
36    }
37
38    pub fn get_entity(&self, id: &str) -> Option<Entity> {
39        self.entities.read().unwrap().get(id).cloned()
40    }
41
42    pub fn put_entity(&self, id: String, entity: Entity) {
43        self.entities.write().unwrap().insert(id, entity);
44    }
45
46    pub fn get_relations(&self, entity_id: &str) -> Option<Vec<Relation>> {
47        self.relations.read().unwrap().get(entity_id).cloned()
48    }
49
50    pub fn put_relations(&self, entity_id: String, relations: Vec<Relation>) {
51        self.relations.write().unwrap().insert(entity_id, relations);
52    }
53
54    /// Invalidate all cached data.
55    pub fn invalidate_all(&self) {
56        self.neighbors.write().unwrap().clear();
57        self.paths.write().unwrap().clear();
58        self.entities.write().unwrap().clear();
59        self.relations.write().unwrap().clear();
60    }
61}