use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::PathBuf;
pub trait NodeEmbedder {
fn embed_nodes(&self, nodes: &[(&str, &str)]) -> HashMap<String, Vec<f32>>;
}
impl NodeEmbedder for crate::embedding::StaticEmbedder {
fn embed_nodes(&self, nodes: &[(&str, &str)]) -> HashMap<String, Vec<f32>> {
crate::embedding::StaticEmbedder::embed_nodes(self, nodes)
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
struct EmbedEntry {
hash: String,
embeddings: HashMap<String, Vec<f32>>,
}
pub struct EmbedCache {
path: PathBuf,
entries: HashMap<String, EmbedEntry>,
}
impl EmbedCache {
pub fn load(path: PathBuf) -> Self {
let entries = std::fs::read_to_string(&path)
.ok()
.and_then(|s| serde_json::from_str(&s).ok())
.unwrap_or_default();
Self { path, entries }
}
pub fn save(&self) -> Result<(), String> {
if let Some(parent) = self.path.parent() {
std::fs::create_dir_all(parent).map_err(|e| e.to_string())?;
}
let json = serde_json::to_string(&self.entries).map_err(|e| e.to_string())?;
std::fs::write(&self.path, json).map_err(|e| e.to_string())
}
pub fn get_if_fresh<'a>(
&'a self,
file_path: &str,
hash: &str,
) -> Option<&'a HashMap<String, Vec<f32>>> {
self.entries
.get(file_path)
.filter(|e| e.hash == hash)
.map(|e| &e.embeddings)
}
pub fn insert(
&mut self,
file_path: String,
hash: String,
embeddings: HashMap<String, Vec<f32>>,
) {
self.entries
.insert(file_path, EmbedEntry { hash, embeddings });
}
}
pub fn embed_file_nodes<E: NodeEmbedder>(
embedder: &E,
nodes: &[(&str, &str)],
file_path: &str,
hash: &str,
cache: &mut EmbedCache,
) -> HashMap<String, Vec<f32>> {
if let Some(cached) = cache.get_if_fresh(file_path, hash) {
return cached.clone();
}
let embs = embedder.embed_nodes(nodes);
cache.insert(file_path.to_string(), hash.to_string(), embs.clone());
embs
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::tempdir;
struct FakeEmbedder;
impl NodeEmbedder for FakeEmbedder {
fn embed_nodes(&self, nodes: &[(&str, &str)]) -> HashMap<String, Vec<f32>> {
nodes
.iter()
.map(|(id, _)| (id.to_string(), vec![1.0f32, 2.0, 3.0]))
.collect()
}
}
#[test]
fn test_embed_file_nodes_computes_embeddings() {
let tmp = tempdir().unwrap();
let cache_path = tmp.path().join("embed_cache.json");
let mut cache = EmbedCache::load(cache_path);
let embedder = FakeEmbedder;
let nodes = vec![("node1", "MyClass"), ("node2", "doThing")];
let result = embed_file_nodes(&embedder, &nodes, "src/foo.rs", "abc123", &mut cache);
assert_eq!(result.len(), 2);
assert!(result.contains_key("node1"));
assert_eq!(result["node1"], vec![1.0f32, 2.0, 3.0]);
}
#[test]
fn test_embed_file_nodes_cache_hit() {
let tmp = tempdir().unwrap();
let cache_path = tmp.path().join("embed_cache.json");
let mut cache = EmbedCache::load(cache_path);
let embedder = FakeEmbedder;
let nodes = vec![("n1", "Foo")];
let first = embed_file_nodes(&embedder, &nodes, "src/bar.rs", "hash1", &mut cache);
cache.insert(
"src/bar.rs".to_string(),
"hash1".to_string(),
[("n1".to_string(), vec![9.0f32])].into_iter().collect(),
);
let second = embed_file_nodes(&embedder, &nodes, "src/bar.rs", "hash1", &mut cache);
assert_eq!(second["n1"], vec![9.0f32]);
let _ = first;
}
#[test]
fn test_embed_file_nodes_cache_miss_on_changed_hash() {
let tmp = tempdir().unwrap();
let cache_path = tmp.path().join("embed_cache.json");
let mut cache = EmbedCache::load(cache_path);
let embedder = FakeEmbedder;
let nodes = vec![("n1", "Foo")];
embed_file_nodes(&embedder, &nodes, "src/baz.rs", "hash1", &mut cache);
cache.insert(
"src/baz.rs".to_string(),
"hash1".to_string(),
[("n1".to_string(), vec![9.0f32])].into_iter().collect(),
);
let result = embed_file_nodes(&embedder, &nodes, "src/baz.rs", "hash2", &mut cache);
assert_eq!(result["n1"], vec![1.0f32, 2.0, 3.0]);
}
#[test]
fn test_embed_cache_save_and_reload() {
let tmp = tempdir().unwrap();
let cache_path = tmp.path().join("embed_cache.json");
let mut cache = EmbedCache::load(cache_path.clone());
let embedder = FakeEmbedder;
let nodes = vec![("x", "Thing")];
embed_file_nodes(&embedder, &nodes, "src/x.rs", "deadbeef", &mut cache);
cache.save().unwrap();
let cache2 = EmbedCache::load(cache_path);
let hit = cache2.get_if_fresh("src/x.rs", "deadbeef");
assert!(hit.is_some());
assert_eq!(hit.unwrap()["x"], vec![1.0f32, 2.0, 3.0]);
}
#[test]
fn test_embed_cache_stale_hash_returns_none() {
let tmp = tempdir().unwrap();
let cache_path = tmp.path().join("embed_cache.json");
let mut cache = EmbedCache::load(cache_path);
cache.insert("f.rs".to_string(), "old".to_string(), Default::default());
assert!(cache.get_if_fresh("f.rs", "new").is_none());
assert!(cache.get_if_fresh("f.rs", "old").is_some());
}
}