light_cache/policy/
linked_arena.rs

1use hashbrown::HashMap;
2
3use super::Policy;
4use crate::LightCache;
5
6use std::hash::{BuildHasher, Hash};
7
8/// A doubly linked list arena
9///
10/// Useful for creating various caching policies
11pub(crate) struct LinkedArena<I, N> {
12    pub(crate) idx_of: HashMap<I, usize>,
13    pub(crate) nodes: Vec<N>,
14    pub(crate) head: Option<usize>,
15    pub(crate) tail: Option<usize>,
16}
17
18impl<I, N> LinkedArena<I, N> {
19    pub fn new() -> Self {
20        LinkedArena {
21            idx_of: HashMap::new(),
22            nodes: Vec::new(),
23            head: None,
24            tail: None,
25        }
26    }
27}
28
29pub trait LinkedNode<I>
30where
31    Self: Sized,
32    I: Copy + Hash + Eq,
33{
34    fn new(item: I, parent: Option<usize>, child: Option<usize>) -> Self;
35    fn item(&self) -> &I;
36
37    fn prev(&self) -> Option<usize>;
38    fn next(&self) -> Option<usize>;
39
40    fn set_prev(&mut self, parent: Option<usize>);
41    fn set_next(&mut self, child: Option<usize>);
42}
43
44impl<I, N> LinkedArena<I, N>
45where
46    N: LinkedNode<I>,
47    I: Copy + Hash + Eq,
48{
49    /// Insert a new node at the front of the list
50    /// 
51    /// # Panics
52    /// Panics if the key already exists in the list
53    pub(crate) fn insert_head(&mut self, key: I) {
54        let new_head = self.nodes.len();
55        if let Some(_) = self.idx_of.insert(key, new_head) {
56            panic!("Key already exists in LinkedArena");
57        }
58
59        if let Some(old_head) = self.head {
60            self.nodes.push(N::new(key, None, Some(old_head)));
61
62            // saftey: we should have a valid old head idx
63            unsafe {
64                self.nodes
65                    .get_unchecked_mut(old_head)
66                    .set_prev(Some(new_head));
67            }
68        } else {
69            self.nodes.push(N::new(key, None, None));
70            self.tail = Some(new_head);
71        }
72
73        self.head = Some(new_head);
74    }
75
76    /// Ensures the node at idx has the correct parent and child relationships
77    ///
78    /// # Note:
79    /// panics if the idx is out of bounds
80    fn relink(&mut self, idx: usize) {
81        let node = self.nodes.get(idx).unwrap();
82        let prev = node.prev();
83        let next = node.next();
84
85        if let Some(parent) = prev {
86            // saftey: the parent indexes are always valid in arena
87            unsafe {
88                self.nodes.get_unchecked_mut(parent).set_next(Some(idx));
89            }
90        } else {
91            self.head = Some(idx);
92        }
93
94        if let Some(child) = next {
95            // saftey: the child indexes are always valid in arena
96            unsafe {
97                self.nodes.get_unchecked_mut(child).set_prev(Some(idx));
98            }
99        } else {
100            self.tail = Some(idx);
101        }
102    }
103
104    /// Remove the node at idx from the list
105    /// passing through any parent/child relationships
106    ///
107    /// # Panics
108    /// if idx is out of bounds
109    fn unlink(&mut self, idx: usize) {
110        let node = self.nodes.get(idx).expect("Invalid index to unlink");
111        let parent = node.prev();
112        let child = node.next();
113
114        // saftey: we should have a valid parent and child idx
115        if let Some(parent) = parent {
116            unsafe {
117                self.nodes.get_unchecked_mut(parent).set_next(child);
118            }
119        } else {
120            self.head = child;
121        }
122
123        if let Some(child) = child {
124            unsafe {
125                self.nodes.get_unchecked_mut(child).set_prev(parent);
126            }
127        } else {
128            self.tail = parent;
129        }
130    }
131}
132
133/// All methods here should never be able to create an invalid state on the list
134#[allow(unused)]
135impl<I, N> LinkedArena<I, N>
136where
137    N: LinkedNode<I>,
138    I: Copy + Hash + Eq,
139{
140    /// Move the node at idx to the front of the list
141    ///
142    /// # Panics
143    /// Panics if new_head is not in the list
144    pub(crate) fn move_to_head_item(&mut self, key: &I) {
145        if let Some(new_head_idx) = self.idx_of.get(key).copied() {
146            self.move_to_head(new_head_idx)
147        } else {
148            panic!("Invalid key to move to head on LinkedArena");
149        }
150    }
151
152    /// # Panics
153    /// if idx is out of bounds
154    pub(crate) fn move_to_head(&mut self, new_head: usize) {
155        self.unlink(new_head);
156
157        // saftey: unlink checks that the idx is valid
158        let node = unsafe { self.nodes.get_unchecked_mut(new_head) };
159        node.set_next(self.head);
160        node.set_prev(None);
161
162        if let Some(old) = self.head.replace(new_head) {
163            unsafe {
164                // unwrap: we should have a valid old head idx
165                self.nodes.get_unchecked_mut(old).set_prev(Some(new_head));
166            }
167        }
168    }
169
170    /// # WARNING: Modifiying the indexes of prev/next will cause UB
171    pub(crate) fn get_node_mut(&mut self, key: &I) -> Option<(usize, &mut N)> {
172        if let Some(idx) = self.idx_of.get(key).copied() {
173            Some((idx, unsafe { self.nodes.get_unchecked_mut(idx) }))
174        } else {
175            None
176        }
177    }
178
179    /// Remove the node from the given index, updating the start or end bounds as needed
180    /// and returning the removed node.
181    /// 
182    /// Noop if the item doesnt exist
183    pub(crate) fn remove_item(&mut self, item: &I) -> Option<(usize, N)> {
184        if let Some(idx) = self.idx_of.get(item).copied() {
185            Some(self.remove(idx))
186        } else {
187            None
188        }
189    }
190
191    /// Remove the node from the given index, updating the start or end bounds as needed
192    /// and returning the removed node.
193    ///
194    /// This method will also call remove on the cache even if already done
195    ///
196    /// # Panics
197    /// If idx is out of bounds
198    pub(crate) fn remove(&mut self, idx: usize) -> (usize, N) {
199        // panics if the idx is out of bounds
200        self.unlink(idx);
201
202        let removed = self.nodes.swap_remove(idx);
203        self.idx_of.remove(removed.item());
204
205        let len = self.nodes.len();
206        // if the last element was just removed than this index will be out of bounds
207        // and theres nothing to relink cause nothing was moved
208        if idx != len {
209            // unwrap: we should have a valid item at idx now and we should have a valid idx for the item
210            *self
211                .idx_of
212                .get_mut(self.nodes.get(idx).unwrap().item())
213                .unwrap() = idx;
214
215            self.relink(idx);
216        }
217
218        (len, removed)
219    }
220
221    /// Try to remove the last node in the list
222    /// 
223    /// Noop if the list is empty
224    pub(crate) fn remove_end<V, S, P>(&mut self, cache: &LightCache<I, V, S, P>)
225    where
226        V: Clone + Sync,
227        S: BuildHasher,
228        P: Policy<I, V>,
229    {
230        if let Some(tail) = self.tail {
231            let (_, n) = self.remove(tail);
232
233            cache.remove_no_policy(n.item());
234        }
235    }
236
237    pub(crate) fn head(&self) -> Option<(usize, &N)> {
238        // saftey: head is always a valid index
239        self.head.map(|idx| (idx, unsafe { self.nodes.get_unchecked(idx) }))
240    }
241
242    pub(crate) fn tail(&self) -> Option<(usize, &N)> {
243        // saftey: tail is always a valid index
244        self.tail.map(|idx| (idx, unsafe { self.nodes.get_unchecked(idx) }))
245    }
246
247    pub(crate) fn len(&self) -> usize {
248        self.nodes.len()
249    }
250}