Skip to main content

flux_common/
cache.rs

1use std::{
2    fs::{self, File},
3    io::BufWriter,
4    path::PathBuf,
5};
6
7use rustc_hash::FxHashMap;
8
9#[derive(Debug, serde::Serialize, serde::Deserialize)]
10struct QueryVal<R> {
11    constr_hash: u64,
12    result: R,
13}
14
15pub struct QueryCache<R> {
16    /// Where the cache is saved. `None` if caching is disabled.
17    path: Option<PathBuf>,
18    entries: FxHashMap<String, QueryVal<R>>,
19}
20
21impl<R> Default for QueryCache<R> {
22    fn default() -> Self {
23        Self::new()
24    }
25}
26
27impl<R> QueryCache<R> {
28    pub fn new() -> Self {
29        QueryCache { path: None, entries: FxHashMap::default() }
30    }
31
32    pub fn insert(&mut self, key: String, constr_hash: u64, result: R) {
33        let val = QueryVal { constr_hash, result };
34        self.entries.insert(key, val);
35    }
36
37    pub fn lookup(&self, key: &String, constr_hash: u64) -> Option<&R> {
38        let val = self.entries.get(key)?;
39        if val.constr_hash == constr_hash { Some(&val.result) } else { None }
40    }
41
42    /// Look up a result by key only (no hash validation).
43    pub fn lookup_by_key(&self, key: &str) -> Option<&R> {
44        self.entries.get(key).map(|e| &e.result)
45    }
46
47    /// Mutate a cached result by key only (no hash validation).
48    pub fn update_result_by_key(&mut self, key: &str, f: impl FnOnce(&mut R)) {
49        if let Some(entry) = self.entries.get_mut(key) {
50            f(&mut entry.result);
51        }
52    }
53}
54
55impl<R: std::fmt::Debug + serde::Serialize + serde::de::DeserializeOwned> QueryCache<R> {
56    /// Saves the cache to the file it was loaded from, overwriting it. Does nothing if the cache
57    /// was not loaded from a file.
58    pub fn save(&self) -> Result<(), std::io::Error> {
59        let Some(path) = &self.path else { return Ok(()) };
60        if let Some(parent) = path.parent() {
61            fs::create_dir_all(parent)?;
62        }
63        let file = File::create(path).unwrap();
64        serde_json::to_writer(BufWriter::new(file), &self.entries).unwrap();
65        Ok(())
66    }
67
68    /// Loads the cache from `path`, starting with an empty cache if the file is missing or
69    /// invalid. The cache is later saved to the same `path`.
70    pub fn load(path: PathBuf) -> Self {
71        let entries = fs::read(&path)
72            .ok()
73            .and_then(|data| serde_json::from_slice(&data).ok())
74            .unwrap_or_default();
75        QueryCache { path: Some(path), entries }
76    }
77}