1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
//! Data structures for memoizing computations.
//!
//! Contruct new caches using [`Default::default`], then construct/retrieve
//! elements with [`get`](Cache::get). `get` should only ever be used with one,
//! `compute` function[^inconsistent].
//!
//! In terms of choice,
//! - [`CopyCache`] should be used for expensive computations that create cheap
//!   (i.e. small) values.
//! - [`Cache`] should be used for expensive computations that create expensive
//!   (i.e. large) values.
//!
//! Both types of caches implement **recursion breaking**. In general because
//! caches are supposed to be used as simple `&` (no `mut`) the reference may be
//! freely copied, including into the `compute` closure. What this means is that
//! a `compute` may call [`get`](Cache::get) on the cache again. This is usually
//! safe and can be used to compute data structures that recursively depend on
//! one another, dynamic-programming style. However if a `get` on a key `k`
//! itself calls `get` again on the same `k` this will either create an infinite
//! recursion or an inconsistent cache[^inconsistent].
//!
//! Consider a simple example where we compute the Fibonacci Series with a
//! [`CopyCache`]:
//!
//! ```rs
//! struct Fib(CopyCache<u32, u32>);
//!
//! impl Fib {
//!   fn get(&self, i: u32) -> u32 {
//!     self.0.get(i, |_| {
//!       if this <= 1 {
//!         return this;
//!       }
//!       let fib_1 = self.get(this - 1);
//!       let fib_2 = self.get(this - 2);
//!       fib_1 + fib_2
//!     })
//!   }
//! }
//!
//! let cache = Fib(Default::default());
//! let fib_5 = cache.get(5);
//! ```
//!
//! This use of recursive [`get`](CopyCache::get) calls is perfectly legal.
//! However if we made an error and called `chache.get(this, ...)` (forgetting
//! the decrement) we would have created an inadvertend infinite recursion.
//!
//! To avoid this scenario both caches are implemented to detect when a
//! recursive call as described is performed and `get` will panic. If your code
//! uses recursive construction and would like to handle this case gracefully
//! use [`get_maybe_recursive`](Cache::get_maybe_recursive) instead wich returns
//! `None` from `get(k)` *iff* `k` this call (potentially transitively)
//! originates from another `get(k)` call.
//!
//! [^inconsistent]: For any given cache value `get` should only ever be used
//!     with one, referentially transparent `compute` function. Essentially this
//!     means running `compute(k)` should always return the same value
//!     *independent of the state of it's environment*. Violation of this rule
//!     can introduces non-determinism in your program.
use std::{cell::RefCell, hash::Hash, pin::Pin};

use rustc_data_structures::fx::FxHashMap as HashMap;

/// Cache for non-copyable types.
pub struct Cache<In, Out>(RefCell<HashMap<In, Option<Pin<Box<Out>>>>>);

impl<In, Out> Cache<In, Out>
where
  In: Hash + Eq + Clone,
{
  /// Size of the cache
  pub fn len(&self) -> usize {
    self.0.borrow().len()
  }
  /// Returns the cached value for the given key, or runs `compute` if
  /// the value is not in cache.
  ///
  /// # Panics
  ///
  /// If this is a recursive invocation for this key.
  pub fn get(&self, key: In, compute: impl FnOnce(In) -> Out) -> &Out {
    self
      .get_maybe_recursive(key, compute)
      .unwrap_or_else(recursion_panic)
  }
  /// Returns the cached value for the given key, or runs `compute` if
  /// the value is not in cache.
  ///
  /// Returns `None` if this is a recursive invocation of `get` for key `key`.
  pub fn get_maybe_recursive<'a>(
    &'a self,
    key: In,
    compute: impl FnOnce(In) -> Out,
  ) -> Option<&'a Out> {
    if !self.0.borrow().contains_key(&key) {
      self.0.borrow_mut().insert(key.clone(), None);
      let out = Box::pin(compute(key.clone()));
      self.0.borrow_mut().insert(key.clone(), Some(out));
    }

    let cache = self.0.borrow();
    // Important here to first `unwrap` the `Option` created by `get`, then
    // propagate the potential option stored in the map.
    let entry = cache.get(&key).expect("invariant broken").as_ref()?;

    // SAFETY: because the entry is pinned, it cannot move and this pointer will
    // only be invalidated if Cache is dropped. The returned reference has a lifetime
    // equal to Cache, so Cache cannot be dropped before this reference goes out of scope.
    Some(unsafe { std::mem::transmute::<&'_ Out, &'a Out>(&**entry) })
  }
}

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<In, Out> Default for Cache<In, Out> {
  fn default() -> Self {
    Cache(RefCell::new(HashMap::default()))
  }
}

/// Cache for copyable types.
pub struct CopyCache<In, Out>(RefCell<HashMap<In, Option<Out>>>);

impl<In, Out> CopyCache<In, Out>
where
  In: Hash + Eq + Clone,
  Out: Copy,
{
  /// Size of the cache
  pub fn len(&self) -> usize {
    self.0.borrow().len()
  }
  /// Returns the cached value for the given key, or runs `compute` if
  /// the value is not in cache.
  ///
  /// # Panics
  ///
  /// If this is a recursive invocation for this key.
  pub fn get(&self, key: In, compute: impl FnOnce(In) -> Out) -> Out {
    self
      .get_maybe_recursive(key, compute)
      .unwrap_or_else(recursion_panic)
  }

  /// Returns the cached value for the given key, or runs `compute` if
  /// the value is not in cache.
  ///
  /// Returns `None` if this is a recursive invocation of `get` for key `key`.
  pub fn get_maybe_recursive(
    &self,
    key: In,
    compute: impl FnOnce(In) -> Out,
  ) -> Option<Out> {
    if !self.0.borrow().contains_key(&key) {
      self.0.borrow_mut().insert(key.clone(), None);
      let out = compute(key.clone());
      self.0.borrow_mut().insert(key.clone(), Some(out));
    }

    *self.0.borrow_mut().get(&key).expect("invariant broken")
  }
}

impl<In, Out> Default for CopyCache<In, Out> {
  fn default() -> Self {
    CopyCache(RefCell::new(HashMap::default()))
  }
}

#[cfg(test)]
mod test {
  use super::*;

  #[test]
  fn test_cached() {
    let cache: Cache<usize, usize> = Cache::default();
    let x = cache.get(0, |_| 0);
    let y = cache.get(1, |_| 1);
    let z = cache.get(0, |_| 2);
    assert_eq!(*x, 0);
    assert_eq!(*y, 1);
    assert_eq!(*z, 0);
    assert!(std::ptr::eq(x, z));
  }

  #[test]
  fn test_recursion_breaking() {
    struct RecursiveUse(Cache<i32, i32>);
    impl RecursiveUse {
      fn get_infinite_recursion(&self, i: i32) -> i32 {
        self
          .0
          .get_maybe_recursive(i, |_| i + self.get_infinite_recursion(i))
          .copied()
          .unwrap_or(-18)
      }
      fn get_safe_recursion(&self, i: i32) -> i32 {
        *self.0.get(i, |_| {
          if i == 0 {
            0
          } else {
            self.get_safe_recursion(i - 1) + i
          }
        })
      }
    }

    let cache = RecursiveUse(Default::default());

    assert_eq!(cache.get_infinite_recursion(60), 42);
    assert_eq!(cache.get_safe_recursion(5), 15);
  }
}