use std::collections::{HashMap, HashSet};
use std::path::PathBuf;
use std::time::Duration;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum RefreshMode {
Manual,
Always,
Files,
#[default]
Auto,
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum MapCacheKey {
Files(FilesKey),
Auto(AutoKey),
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct FilesKey {
pub chat_fnames: Option<Vec<PathBuf>>,
pub other_fnames: Option<Vec<PathBuf>>,
pub max_tokens: usize,
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct AutoKey {
pub chat_fnames: Option<Vec<PathBuf>>,
pub other_fnames: Option<Vec<PathBuf>>,
pub max_tokens: usize,
pub mentioned_fnames: Option<Vec<String>>,
pub mentioned_idents: Option<Vec<String>>,
pub anchor_fnames: Option<Vec<PathBuf>>,
pub anchor_idents: Option<Vec<String>>,
pub anchor_scoped: Option<Vec<(PathBuf, String)>>,
}
pub struct AnchorCacheParams<'a> {
pub anchor_fnames: &'a [PathBuf],
pub anchor_idents: &'a HashSet<String>,
pub anchor_scoped: &'a [(PathBuf, String)],
}
impl MapCacheKey {
pub fn files(chat_fnames: &[PathBuf], other_fnames: &[PathBuf], max_tokens: usize) -> Self {
MapCacheKey::Files(FilesKey {
chat_fnames: non_empty_sorted(chat_fnames),
other_fnames: non_empty_sorted(other_fnames),
max_tokens,
})
}
pub fn auto(
chat_fnames: &[PathBuf],
other_fnames: &[PathBuf],
max_tokens: usize,
mentioned_fnames: &HashSet<String>,
mentioned_idents: &HashSet<String>,
anchors: AnchorCacheParams<'_>,
) -> Self {
MapCacheKey::Auto(AutoKey {
chat_fnames: non_empty_sorted(chat_fnames),
other_fnames: non_empty_sorted(other_fnames),
max_tokens,
mentioned_fnames: non_empty_sorted_set(mentioned_fnames),
mentioned_idents: non_empty_sorted_set(mentioned_idents),
anchor_fnames: non_empty_sorted(anchors.anchor_fnames),
anchor_idents: non_empty_sorted_set(anchors.anchor_idents),
anchor_scoped: non_empty_sorted_pairs(anchors.anchor_scoped),
})
}
}
fn non_empty_sorted<T: Clone + Ord>(items: &[T]) -> Option<Vec<T>> {
if items.is_empty() {
None
} else {
let mut v: Vec<_> = items.to_vec();
v.sort();
Some(v)
}
}
fn non_empty_sorted_pairs<A: Clone + Ord, B: Clone + Ord>(items: &[(A, B)]) -> Option<Vec<(A, B)>> {
if items.is_empty() {
None
} else {
let mut v: Vec<_> = items.to_vec();
v.sort();
Some(v)
}
}
fn non_empty_sorted_set<T: Clone + Ord>(items: &HashSet<T>) -> Option<Vec<T>> {
if items.is_empty() {
None
} else {
let mut v: Vec<_> = items.iter().cloned().collect();
v.sort();
Some(v)
}
}
#[derive(Debug, Default)]
pub struct MapCache {
cache: HashMap<MapCacheKey, String>,
last_map: Option<String>,
last_duration: Duration,
}
impl MapCache {
pub fn new() -> Self {
Self::default()
}
pub fn should_use_cache(&self, mode: RefreshMode, force_refresh: bool) -> bool {
if force_refresh {
return false;
}
match mode {
RefreshMode::Manual => self.last_map.is_some(),
RefreshMode::Always => false,
RefreshMode::Files => true,
RefreshMode::Auto => self.last_duration > Duration::from_secs_f64(1.0),
}
}
pub fn get(
&self,
key: &MapCacheKey,
mode: RefreshMode,
force_refresh: bool,
) -> Option<&String> {
if !self.should_use_cache(mode, force_refresh) {
return None;
}
if mode == RefreshMode::Manual {
return self.last_map.as_ref();
}
self.cache.get(key)
}
pub fn get_last_map(&self) -> Option<&String> {
self.last_map.as_ref()
}
pub fn set(&mut self, key: MapCacheKey, value: String, duration: Duration) {
self.cache.insert(key, value.clone());
self.last_map = Some(value);
self.last_duration = duration;
}
pub fn last_duration(&self) -> Duration {
self.last_duration
}
pub fn clear(&mut self) {
self.cache.clear();
self.last_map = None;
self.last_duration = Duration::ZERO;
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn cache_key_files_mode() {
let key = MapCacheKey::files(
&[PathBuf::from("a.rs"), PathBuf::from("b.rs")],
&[PathBuf::from("c.rs")],
1024,
);
if let MapCacheKey::Files(k) = &key {
assert_eq!(k.max_tokens, 1024);
let chat = k.chat_fnames.as_ref().unwrap();
assert_eq!(chat[0], PathBuf::from("a.rs"));
assert_eq!(chat[1], PathBuf::from("b.rs"));
} else {
panic!("Expected Files key");
}
}
#[test]
fn cache_key_empty_is_none() {
let key = MapCacheKey::files(&[], &[PathBuf::from("c.rs")], 1024);
if let MapCacheKey::Files(k) = &key {
assert!(k.chat_fnames.is_none()); assert!(k.other_fnames.is_some());
} else {
panic!("Expected Files key");
}
}
#[test]
fn cache_set_and_get() {
let mut cache = MapCache::new();
let key = MapCacheKey::files(&[PathBuf::from("a.rs")], &[], 1024);
cache.set(key.clone(), "test map".to_string(), Duration::from_secs(2));
let result = cache.get(&key, RefreshMode::Files, false);
assert!(result.is_some());
assert_eq!(result.unwrap(), "test map");
}
#[test]
fn cache_manual_mode() {
let mut cache = MapCache::new();
let key = MapCacheKey::files(&[PathBuf::from("a.rs")], &[], 1024);
cache.set(key.clone(), "test map".to_string(), Duration::from_secs(1));
let other_key = MapCacheKey::files(&[PathBuf::from("b.rs")], &[], 2048);
let result = cache.get(&other_key, RefreshMode::Manual, false);
assert!(result.is_some());
assert_eq!(result.unwrap(), "test map");
}
#[test]
fn cache_always_mode() {
let mut cache = MapCache::new();
let key = MapCacheKey::files(&[PathBuf::from("a.rs")], &[], 1024);
cache.set(key.clone(), "test map".to_string(), Duration::from_secs(1));
let result = cache.get(&key, RefreshMode::Always, false);
assert!(result.is_none());
}
#[test]
fn cache_auto_mode_threshold() {
let mut cache = MapCache::new();
let key = MapCacheKey::auto(
&[PathBuf::from("a.rs")],
&[],
1024,
&HashSet::new(),
&HashSet::new(),
AnchorCacheParams {
anchor_fnames: &[],
anchor_idents: &HashSet::new(),
anchor_scoped: &[],
},
);
cache.set(
key.clone(),
"fast map".to_string(),
Duration::from_millis(500),
);
let result = cache.get(&key, RefreshMode::Auto, false);
assert!(result.is_none());
cache.set(key.clone(), "slow map".to_string(), Duration::from_secs(2));
let result = cache.get(&key, RefreshMode::Auto, false);
assert!(result.is_some());
}
#[test]
fn cache_force_refresh() {
let mut cache = MapCache::new();
let key = MapCacheKey::files(&[PathBuf::from("a.rs")], &[], 1024);
cache.set(key.clone(), "test map".to_string(), Duration::from_secs(2));
let result = cache.get(&key, RefreshMode::Files, true);
assert!(result.is_none());
}
#[test]
fn cache_always_writes() {
let mut cache = MapCache::new();
let key = MapCacheKey::files(&[PathBuf::from("a.rs")], &[], 1024);
cache.set(key.clone(), "test map".to_string(), Duration::from_secs(1));
let result = cache.get(&key, RefreshMode::Files, false);
assert!(result.is_some());
}
}