use std::collections::HashMap;
use std::path::{Path, PathBuf};
use parking_lot::RwLock;
use sha2::{Digest, Sha256};
use torvyn_types::ComponentTypeId;
use crate::error::EngineError;
use crate::traits::WasmEngine;
use crate::types::CompiledComponent;
pub struct CompiledComponentCache {
memory: RwLock<HashMap<ComponentTypeId, CompiledComponent>>,
disk_dir: Option<PathBuf>,
}
impl CompiledComponentCache {
pub fn new(disk_dir: Option<PathBuf>) -> Self {
Self {
memory: RwLock::new(HashMap::new()),
disk_dir,
}
}
pub fn compute_type_id(bytes: &[u8]) -> ComponentTypeId {
let mut hasher = Sha256::new();
hasher.update(bytes);
let hash: [u8; 32] = hasher.finalize().into();
ComponentTypeId::new(hash)
}
pub fn get<E: WasmEngine>(
&self,
type_id: &ComponentTypeId,
engine: &E,
) -> Result<Option<CompiledComponent>, EngineError> {
{
let guard = self.memory.read();
if let Some(compiled) = guard.get(type_id) {
return Ok(Some(compiled.clone()));
}
}
if let Some(ref dir) = self.disk_dir {
let path = disk_cache_path(dir, type_id);
if path.exists() {
if let Ok(bytes) = std::fs::read(&path) {
match unsafe { engine.deserialize_component(&bytes) } {
Ok(Some(compiled)) => {
let mut guard = self.memory.write();
guard.insert(*type_id, compiled.clone());
return Ok(Some(compiled));
}
Ok(None) => {
let _ = std::fs::remove_file(&path);
}
Err(_e) => {
let _ = std::fs::remove_file(&path);
}
}
}
}
}
Ok(None)
}
pub fn insert<E: WasmEngine>(
&self,
type_id: ComponentTypeId,
compiled: CompiledComponent,
engine: &E,
) {
{
let mut guard = self.memory.write();
guard.insert(type_id, compiled.clone());
}
if let Some(ref dir) = self.disk_dir {
if let Ok(bytes) = engine.serialize_component(&compiled) {
let path = disk_cache_path(dir, &type_id);
if let Some(parent) = path.parent() {
let _ = std::fs::create_dir_all(parent);
}
let _ = std::fs::write(&path, &bytes);
}
}
}
pub fn len(&self) -> usize {
self.memory.read().len()
}
pub fn is_empty(&self) -> bool {
self.memory.read().is_empty()
}
pub fn clear(&self) {
self.memory.write().clear();
}
pub fn compile_or_get<E: WasmEngine>(
&self,
bytes: &[u8],
engine: &E,
) -> Result<(ComponentTypeId, CompiledComponent), EngineError> {
let type_id = Self::compute_type_id(bytes);
if let Some(compiled) = self.get(&type_id, engine)? {
return Ok((type_id, compiled));
}
let compiled = engine.compile_component(bytes)?;
self.insert(type_id, compiled.clone(), engine);
Ok((type_id, compiled))
}
}
fn disk_cache_path(dir: &Path, type_id: &ComponentTypeId) -> PathBuf {
let hex = format!("{type_id}");
let prefix = &hex[..2.min(hex.len())];
dir.join(prefix).join(format!("{hex}.bin"))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_compute_type_id_deterministic() {
let id1 = CompiledComponentCache::compute_type_id(b"hello world");
let id2 = CompiledComponentCache::compute_type_id(b"hello world");
assert_eq!(id1, id2);
}
#[test]
fn test_compute_type_id_different_inputs() {
let id1 = CompiledComponentCache::compute_type_id(b"hello");
let id2 = CompiledComponentCache::compute_type_id(b"world");
assert_ne!(id1, id2);
}
#[test]
fn test_cache_new_empty() {
let cache = CompiledComponentCache::new(None);
assert_eq!(cache.len(), 0);
assert!(cache.is_empty());
}
#[test]
fn test_cache_clear() {
let cache = CompiledComponentCache::new(None);
cache.clear();
assert!(cache.is_empty());
}
#[test]
fn test_disk_cache_path_format() {
let type_id = ComponentTypeId::new([0xab; 32]);
let path = disk_cache_path(Path::new("/tmp/cache"), &type_id);
let path_str = path.to_string_lossy();
assert!(path_str.contains("cache"));
assert!(path_str.ends_with(".bin"));
}
}