flowistry_pdg_construction/
two_level_cache.rsuse std::{
cell::RefCell,
collections::{hash_map::Entry, HashMap},
hash::Hash,
pin::Pin,
};
pub struct TwoLevelCache<K1, K2, Out>(RefCell<TwoLevelCacheInner<K1, K2, Out>>);
#[doc(hidden)]
pub type TwoLevelCacheInner<K1, K2, Out> = HashMap<K1, TwolevelCacheValue<K2, Out>>;
#[doc(hidden)]
pub type TwolevelCacheValue<K2, Out> = (bool, HashMap<K2, Pin<Box<Out>>>);
impl<K1: Eq + Hash + Clone, K2: Eq + Hash + Clone, Out> TwoLevelCache<K1, K2, Out> {
pub fn len(&self) -> usize {
self.0.borrow().values().map(|(_, m)| m.len()).sum()
}
pub fn get(&self, key: (K1, K2), compute: impl FnOnce((K1, K2)) -> Out) -> &Out {
self.get_maybe_recursive(key, compute)
.unwrap_or_else(recursion_panic)
}
#[doc(hidden)]
pub fn borrow(&self) -> std::cell::Ref<'_, TwoLevelCacheInner<K1, K2, Out>> {
self.0.borrow()
}
pub fn get_maybe_recursive(
&self,
key: (K1, K2),
compute: impl FnOnce((K1, K2)) -> Out,
) -> Option<&Out> {
match self.try_retrieve(key, |in_| Some(compute(in_))) {
Retrieval::Recursive => None,
Retrieval::Success(v) => Some(v),
Retrieval::Uncomputable => unreachable!(),
}
}
pub fn try_retrieve<'a>(
&'a self,
key: (K1, K2),
compute: impl FnOnce((K1, K2)) -> Option<Out>,
) -> Retrieval<&'a Out> {
let k1 = key.0;
let k2 = key.1;
let is_present = {
let mut b = self.0.borrow_mut();
match b.entry(k1.clone()) {
Entry::Vacant(v) => {
v.insert((true, HashMap::new()));
false
}
Entry::Occupied(o) if o.get().0 => {
return Retrieval::Recursive;
}
Entry::Occupied(mut o) if !o.get().1.contains_key(&k2) => {
o.get_mut().0 = true;
false
}
_ => true,
}
};
if !is_present {
let result = compute((k1.clone(), k2.clone()));
let mut borrow = self.0.borrow_mut();
let (constructing, entry) = &mut borrow.get_mut(&k1).unwrap();
if let Some(out) = result {
entry.insert(k2.clone(), Box::pin(out));
} else {
entry.remove(&k2);
}
*constructing = false;
}
let cache = self.0.borrow();
let entry = cache.get(&k1).unwrap();
match entry.1.get(&k2) {
None => Retrieval::Uncomputable,
Some(entry) => Retrieval::Success(
unsafe { std::mem::transmute::<&'_ Out, &'a Out>(&**entry) },
),
}
}
pub fn is_in_cache(&self, key: &(K1, K2)) -> bool {
let k1 = &key.0;
let k2 = &key.1;
self.0
.borrow()
.get(k1)
.is_some_and(|v| v.1.contains_key(k2))
}
#[allow(dead_code)]
pub(crate) unsafe fn clear(&self) {
self.0.borrow_mut().clear()
}
}
pub enum Retrieval<T> {
Success(T),
Recursive,
Uncomputable,
}
impl<T> Retrieval<T> {
pub fn into_success(self) -> Option<T> {
match self {
Retrieval::Success(v) => Some(v),
_ => None,
}
}
}
fn recursion_panic<A>() -> A {
panic!("Recursion detected! The computation of a value tried to retrieve the same from the cache. Using `get_maybe_recursive` to handle this case gracefully.")
}
impl<K1, K2, Out> Default for TwoLevelCache<K1, K2, Out> {
fn default() -> Self {
Self(RefCell::new(HashMap::default()))
}
}