compressible_map/
local_cache.rs

1use core::hash::{BuildHasher, Hash};
2use std::cell::UnsafeCell;
3use std::collections::{hash_map, HashMap};
4use std::pin::Pin;
5
6/// When immutable cache access is required, use this `LocalCache` to store evicted values. Then
7/// when you get mutable cache access, call `into_iter` to update the cache manually.
8///
9/// This cache comes with a price: the values will be boxed in order to provide stable borrows.
10/// Ideally, it should be cheap to box your values, i.e. most of their data should already be on the
11/// heap.
12#[derive(Default)]
13pub struct LocalCache<K, V, H> {
14    accesses: UnsafeCell<HashMap<K, LocalAccess<Pin<Box<V>>>, H>>,
15}
16
17pub enum LocalAccess<V> {
18    /// Represents a global cache hit that we want to remember so we can update the LRU order later.
19    Cached,
20    /// Represents a miss of the global cache that required us to cache the value locally.
21    /// `Pin<Box>` is used to maintain a stable address for the value, even if the the map it lives
22    /// in is mutated.
23    Missed(V),
24}
25
26impl<V> LocalAccess<V> {
27    fn unwrap_ref(&self) -> &V {
28        match self {
29            LocalAccess::Cached => panic!("Tried to unwrap access without value"),
30            LocalAccess::Missed(value) => &value,
31        }
32    }
33
34    fn map<T>(self, f: impl FnOnce(V) -> T) -> LocalAccess<T> {
35        match self {
36            LocalAccess::Cached => LocalAccess::Cached,
37            LocalAccess::Missed(value) => LocalAccess::Missed(f(value)),
38        }
39    }
40}
41
42impl<K, V, H> LocalCache<K, V, H>
43where
44    K: Eq + Hash,
45    H: Default + BuildHasher,
46{
47    pub fn new() -> Self {
48        LocalCache {
49            accesses: UnsafeCell::new(HashMap::with_hasher(Default::default())),
50        }
51    }
52
53    // SAFE: We guarantee in these APIs that all references returned are valid for the lifetime of
54    // the `LocalCache`, even as new values are added to the map. The invariants are:
55    //   1. Once a value is placed here, it will never get dropped or moved until calling
56    //      `into_iter`.
57    //   2. The values are placed into `Pin<Box<V>>` so the memory address is guaranteed stable.
58    //   3. Returned references must be dropped before calling `into_iter`.
59
60    pub fn remember_cached_access(&self, key: K) {
61        let mut_accesses = unsafe { &mut *self.accesses.get() };
62        mut_accesses.entry(key).or_insert(LocalAccess::Cached);
63    }
64
65    pub fn get_or_insert_with(&self, key: K, f: impl FnOnce() -> V) -> &V {
66        let mut_accesses = unsafe { &mut *self.accesses.get() };
67        match mut_accesses.entry(key) {
68            hash_map::Entry::Occupied(occupied) => {
69                let access_ref = occupied.into_mut();
70                match access_ref {
71                    LocalAccess::Cached => {
72                        *access_ref = LocalAccess::Missed(Box::pin(f()));
73
74                        access_ref.unwrap_ref()
75                    }
76                    LocalAccess::Missed(value) => value,
77                }
78            }
79            hash_map::Entry::Vacant(vacant) => {
80                let access_ref = vacant.insert(LocalAccess::Missed(Box::pin(f())));
81
82                access_ref.unwrap_ref()
83            }
84        }
85    }
86
87    pub fn into_iter(self) -> impl Iterator<Item = (K, LocalAccess<V>)> {
88        self.accesses.into_inner().into_iter().map(|(k, access)| {
89            (
90                k,
91                access.map(|value| unsafe { *Pin::into_inner_unchecked(value) }),
92            )
93        })
94    }
95}
compressible_map/local_cache.rs

compressible_map/
local_cache.rs
