Skip to content

Commit

Permalink
Merge pull request #5 from emilHof/stack-borrow-nonnull
Browse files Browse the repository at this point in the history 
Change `*mut LruEntry<K,V>` to `NonNull<LruEntry<K,V>>`
  • Loading branch information
@emilHof
emilHof committed Dec 4, 2022
2 parents d82a0da + 133f76b commit 78a87aa
Showing 1 changed file with 41 additions and 35 deletions.
76 changes: 41 additions & 35 deletions src/lib.rs
View file
Expand Up @@ -74,7 +74,7 @@ use core::iter::FusedIterator;
use core::marker::PhantomData;
use core::mem;
use core::num::NonZeroUsize;
use core::ptr;
use core::ptr::{self, NonNull};
use core::usize;

#[cfg(any(test, not(feature = "hashbrown")))]
Expand Down Expand Up @@ -189,7 +189,7 @@ pub type DefaultHasher = std::collections::hash_map::RandomState;

/// An LRU Cache
pub struct LruCache<K, V, S = DefaultHasher> {
map: HashMap<KeyRef<K>, *mut LruEntry<K, V>, S>,
map: HashMap<KeyRef<K>, NonNull<LruEntry<K, V>>, S>,
cap: NonZeroUsize,

// head and tail are sigil nodes to facilitate inserting entries
Expand Down Expand Up @@ -266,7 +266,7 @@ impl<K: Hash + Eq, V, S: BuildHasher> LruCache<K, V, S> {
/// Creates a new LRU Cache with the given capacity.
fn construct(
cap: NonZeroUsize,
map: HashMap<KeyRef<K>, *mut LruEntry<K, V>, S>,
map: HashMap<KeyRef<K>, NonNull<LruEntry<K, V>>, S>,
) -> LruCache<K, V, S> {
// NB: The compiler warns that cache does not need to be marked as mutable if we
// declare it as such since we only mutate it inside the unsafe block.
Expand Down Expand Up @@ -342,21 +342,22 @@ impl<K: Hash + Eq, V, S: BuildHasher> LruCache<K, V, S> {

match node_ref {
Some(node_ref) => {
let node_ptr: *mut LruEntry<K, V> = *node_ref;
let node_ptr: *mut LruEntry<K, V> = node_ref.as_ptr();

// if the key is already in the cache just update its value and move it to the
// front of the list
unsafe { mem::swap(&mut v, &mut (*(*node_ptr).val.as_mut_ptr()) as &mut V) }
unsafe { mem::swap(&mut v, &mut *(*node_ptr).val.as_mut_ptr()) }
self.detach(node_ptr);
self.attach(node_ptr);
Some((k, v))
}
None => {
let (replaced, node) = self.replace_or_create_node(k, v);
let node_ptr: *mut LruEntry<K, V> = node.as_ptr();

self.attach(node);
self.attach(node_ptr);

let keyref = unsafe { (*node).key.as_ptr() };
let keyref = unsafe { (*node_ptr).key.as_ptr() };
self.map.insert(KeyRef { k: keyref }, node);

replaced.filter(|_| capture)
Expand All @@ -367,29 +368,32 @@ impl<K: Hash + Eq, V, S: BuildHasher> LruCache<K, V, S> {
// Used internally to swap out a node if the cache is full or to create a new node if space
// is available. Shared between `put`, `push`, `get_or_insert`, and `get_or_insert_mut`.
#[allow(clippy::type_complexity)]
fn replace_or_create_node(&mut self, k: K, v: V) -> (Option<(K, V)>, *mut LruEntry<K, V>) {
fn replace_or_create_node(&mut self, k: K, v: V) -> (Option<(K, V)>, NonNull<LruEntry<K, V>>) {
if self.len() == self.cap().get() {
// if the cache is full, remove the last entry so we can use it for the new key
let old_key = KeyRef {
k: unsafe { &(*(*(*self.tail).prev).key.as_ptr()) },
};
let old_node = self.map.remove(&old_key).unwrap();
let node_ptr: *mut LruEntry<K, V> = old_node.as_ptr();

// read out the node's old key and value and then replace it
let replaced = unsafe {
(
mem::replace(&mut (*old_node).key, mem::MaybeUninit::new(k)).assume_init(),
mem::replace(&mut (*old_node).val, mem::MaybeUninit::new(v)).assume_init(),
mem::replace(&mut (*node_ptr).key, mem::MaybeUninit::new(k)).assume_init(),
mem::replace(&mut (*node_ptr).val, mem::MaybeUninit::new(v)).assume_init(),
)
};

let node_ptr: *mut LruEntry<K, V> = old_node;
self.detach(node_ptr);

(Some(replaced), old_node)
} else {
// if the cache is not full allocate a new LruEntry
(None, Box::into_raw(Box::new(LruEntry::new(k, v))))
// Safety: We allocate, turn into raw, and get NonNull all in one step.
(None, unsafe {
NonNull::new_unchecked(Box::into_raw(Box::new(LruEntry::new(k, v))))
})
}
}

Expand Down Expand Up @@ -418,12 +422,12 @@ impl<K: Hash + Eq, V, S: BuildHasher> LruCache<K, V, S> {
Q: Hash + Eq + ?Sized,
{
if let Some(node) = self.map.get_mut(k) {
let node_ptr: *mut LruEntry<K, V> = *node;
let node_ptr: *mut LruEntry<K, V> = node.as_ptr();

self.detach(node_ptr);
self.attach(node_ptr);

Some(unsafe { &(*(*node_ptr).val.as_ptr()) as &V })
Some(unsafe { &*(*node_ptr).val.as_ptr() })
} else {
None
}
Expand Down Expand Up @@ -454,12 +458,12 @@ impl<K: Hash + Eq, V, S: BuildHasher> LruCache<K, V, S> {
Q: Hash + Eq + ?Sized,
{
if let Some(node) = self.map.get_mut(k) {
let node_ptr: *mut LruEntry<K, V> = *node;
let node_ptr: *mut LruEntry<K, V> = node.as_ptr();

self.detach(node_ptr);
self.attach(node_ptr);

Some(unsafe { &mut (*(*node_ptr).val.as_mut_ptr()) as &mut V })
Some(unsafe { &mut *(*node_ptr).val.as_mut_ptr() })
} else {
None
}
Expand Down Expand Up @@ -492,21 +496,22 @@ impl<K: Hash + Eq, V, S: BuildHasher> LruCache<K, V, S> {
F: FnOnce() -> V,
{
if let Some(node) = self.map.get_mut(&KeyRef { k: &k }) {
let node_ptr: *mut LruEntry<K, V> = *node;
let node_ptr: *mut LruEntry<K, V> = node.as_ptr();

self.detach(node_ptr);
self.attach(node_ptr);

unsafe { &(*(*node_ptr).val.as_ptr()) as &V }
unsafe { &*(*node_ptr).val.as_ptr() }
} else {
let v = f();
let (_, node) = self.replace_or_create_node(k, v);
let node_ptr: *mut LruEntry<K, V> = node.as_ptr();

self.attach(node);
self.attach(node_ptr);

let keyref = unsafe { (*node).key.as_ptr() };
let keyref = unsafe { (*node_ptr).key.as_ptr() };
self.map.insert(KeyRef { k: keyref }, node);
unsafe { &(*(*node).val.as_ptr()) as &V }
unsafe { &*(*node_ptr).val.as_ptr() }
}
}

Expand Down Expand Up @@ -537,21 +542,22 @@ impl<K: Hash + Eq, V, S: BuildHasher> LruCache<K, V, S> {
F: FnOnce() -> V,
{
if let Some(node) = self.map.get_mut(&KeyRef { k: &k }) {
let node_ptr: *mut LruEntry<K, V> = *node;
let node_ptr: *mut LruEntry<K, V> = node.as_ptr();

self.detach(node_ptr);
self.attach(node_ptr);

unsafe { &mut (*(*node_ptr).val.as_mut_ptr()) as &mut V }
unsafe { &mut *(*node_ptr).val.as_mut_ptr() }
} else {
let v = f();
let (_, node) = self.replace_or_create_node(k, v);
let node_ptr: *mut LruEntry<K, V> = node.as_ptr();

self.attach(node);
self.attach(node_ptr);

let keyref = unsafe { (*node).key.as_ptr() };
let keyref = unsafe { (*node_ptr).key.as_ptr() };
self.map.insert(KeyRef { k: keyref }, node);
unsafe { &mut (*(*node).val.as_mut_ptr()) as &mut V }
unsafe { &mut *(*node_ptr).val.as_mut_ptr() }
}
}

Expand Down Expand Up @@ -579,7 +585,7 @@ impl<K: Hash + Eq, V, S: BuildHasher> LruCache<K, V, S> {
{
self.map
.get(k)
.map(|node| unsafe { &(*(**node).val.as_ptr()) as &V })
.map(|node| unsafe { &*node.as_ref().val.as_ptr() })
}

/// Returns a mutable reference to the value corresponding to the key in the cache or `None`
Expand All @@ -606,7 +612,7 @@ impl<K: Hash + Eq, V, S: BuildHasher> LruCache<K, V, S> {
{
match self.map.get_mut(k) {
None => None,
Some(node) => Some(unsafe { &mut (*(**node).val.as_mut_ptr()) as &mut V }),
Some(node) => Some(unsafe { &mut *(*node.as_ptr()).val.as_mut_ptr() }),
}
}

Expand Down Expand Up @@ -693,7 +699,7 @@ impl<K: Hash + Eq, V, S: BuildHasher> LruCache<K, V, S> {
None => None,
Some(old_node) => {
let mut old_node = unsafe {
let mut old_node = *Box::from_raw(old_node);
let mut old_node = *Box::from_raw(old_node.as_ptr());
ptr::drop_in_place(old_node.key.as_mut_ptr());

old_node
Expand Down Expand Up @@ -734,7 +740,7 @@ impl<K: Hash + Eq, V, S: BuildHasher> LruCache<K, V, S> {
match self.map.remove(k) {
None => None,
Some(old_node) => {
let mut old_node = unsafe { *Box::from_raw(old_node) };
let mut old_node = unsafe { *Box::from_raw(old_node.as_ptr()) };

self.detach(&mut old_node);

Expand Down Expand Up @@ -800,7 +806,7 @@ impl<K: Hash + Eq, V, S: BuildHasher> LruCache<K, V, S> {
Q: Hash + Eq + ?Sized,
{
if let Some(node) = self.map.get_mut(k) {
let node_ptr: *mut LruEntry<K, V> = *node;
let node_ptr: *mut LruEntry<K, V> = node.as_ptr();
self.detach(node_ptr);
self.attach(node_ptr);
}
Expand Down Expand Up @@ -836,7 +842,7 @@ impl<K: Hash + Eq, V, S: BuildHasher> LruCache<K, V, S> {
Q: Hash + Eq + ?Sized,
{
if let Some(node) = self.map.get_mut(k) {
let node_ptr: *mut LruEntry<K, V> = *node;
let node_ptr: *mut LruEntry<K, V> = node.as_ptr();
self.detach(node_ptr);
self.attach_last(node_ptr);
}
Expand Down Expand Up @@ -1026,9 +1032,9 @@ impl<K: Hash + Eq, V, S: BuildHasher> LruCache<K, V, S> {
k: unsafe { &(*(*(*self.tail).prev).key.as_ptr()) },
};
let old_node = self.map.remove(&old_key).unwrap();
let node_ptr: *mut LruEntry<K, V> = old_node;
let node_ptr: *mut LruEntry<K, V> = old_node.as_ptr();
self.detach(node_ptr);
unsafe { Some(Box::from_raw(old_node)) }
unsafe { Some(Box::from_raw(node_ptr)) }
} else {
None
}
Expand Down Expand Up @@ -1065,7 +1071,7 @@ impl<K: Hash + Eq, V, S: BuildHasher> LruCache<K, V, S> {
impl<K, V, S> Drop for LruCache<K, V, S> {
fn drop(&mut self) {
self.map.drain().for_each(|(_, node)| unsafe {
let mut node = *Box::from_raw(node);
let mut node = *Box::from_raw(node.as_ptr());
ptr::drop_in_place((node).key.as_mut_ptr());
ptr::drop_in_place((node).val.as_mut_ptr());
});
Expand Down

0 comments on commit 78a87aa

Please sign in to comment.