blob: afa4bb7c76a667ddc9b289799515ab1afe72a010 [file] [log] [blame]
//! **This module is experimental**
//!
//! This module provides threadsafe versions of FrozenMap and FrozenVec,
//! ideal for use as a cache.
//!
//! These lock internally, however locks only last as long as the method calls
//!
use stable_deref_trait::StableDeref;
use std::alloc::Layout;
use std::borrow::Borrow;
use std::collections::BTreeMap;
use std::collections::HashMap;
use std::hash::Hash;
use std::iter::{FromIterator, IntoIterator};
use std::mem::MaybeUninit;
use std::ops::Index;
use std::sync::atomic::AtomicPtr;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use std::sync::RwLock;
/// Append-only threadsafe version of `std::collections::HashMap` where
/// insertion does not require mutable access
pub struct FrozenMap<K, V> {
map: RwLock<HashMap<K, V>>,
}
impl<K, V> Default for FrozenMap<K, V> {
fn default() -> Self {
Self {
map: Default::default(),
}
}
}
impl<K: Eq + Hash, V: StableDeref> FrozenMap<K, V> {
// these should never return &K or &V
// these should never delete any entries
pub fn new() -> Self {
Self::default()
}
/// If the key exists in the map, returns a reference
/// to the corresponding value, otherwise inserts a
/// new entry in the map for that key and returns a
/// reference to the given value.
///
/// Existing values are never overwritten.
///
/// The key may be any borrowed form of the map's key type, but
/// [`Hash`] and [`Eq`] on the borrowed form *must* match those for
/// the key type.
///
/// # Examples
///
/// ```
/// use elsa::sync::FrozenMap;
///
/// let map = FrozenMap::new();
/// assert_eq!(map.insert(1, Box::new("a")), &"a");
/// assert_eq!(map.insert(1, Box::new("b")), &"a");
/// ```
pub fn insert(&self, k: K, v: V) -> &V::Target {
let mut map = self.map.write().unwrap();
let ret = unsafe {
let inserted = &**map.entry(k).or_insert(v);
&*(inserted as *const _)
};
ret
}
/// If the key exists in the map, returns a reference to the corresponding
/// value, otherwise inserts a new entry in the map for that key and the
/// value returned by the creation function, and returns a reference to the
/// generated value.
///
/// Existing values are never overwritten.
///
/// The key may be any borrowed form of the map's key type, but [`Hash`] and
/// [`Eq`] on the borrowed form *must* match those for the key type.
///
/// **Note** that the write lock is held for the duration of this function’s
/// execution, even while the value creation function is executing (if
/// needed). This will block any concurrent `get` or `insert` calls.
///
/// # Examples
///
/// ```
/// use elsa::sync::FrozenMap;
///
/// let map = FrozenMap::new();
/// assert_eq!(map.insert_with(1, || Box::new("a")), &"a");
/// assert_eq!(map.insert_with(1, || unreachable!()), &"a");
/// ```
pub fn insert_with(&self, k: K, f: impl FnOnce() -> V) -> &V::Target {
let mut map = self.map.write().unwrap();
let ret = unsafe {
let inserted = &**map.entry(k).or_insert_with(f);
&*(inserted as *const _)
};
ret
}
/// If the key exists in the map, returns a reference to the corresponding
/// value, otherwise inserts a new entry in the map for that key and the
/// value returned by the creation function, and returns a reference to the
/// generated value.
///
/// Existing values are never overwritten.
///
/// The key may be any borrowed form of the map's key type, but [`Hash`] and
/// [`Eq`] on the borrowed form *must* match those for the key type.
///
/// **Note** that the write lock is held for the duration of this function’s
/// execution, even while the value creation function is executing (if
/// needed). This will block any concurrent `get` or `insert` calls.
///
/// # Examples
///
/// ```
/// use elsa::sync::FrozenMap;
///
/// let map = FrozenMap::new();
/// assert_eq!(map.insert_with_key(1, |_| Box::new("a")), &"a");
/// assert_eq!(map.insert_with_key(1, |_| unreachable!()), &"a");
/// ```
pub fn insert_with_key(&self, k: K, f: impl FnOnce(&K) -> V) -> &V::Target {
let mut map = self.map.write().unwrap();
let ret = unsafe {
let inserted = &**map.entry(k).or_insert_with_key(f);
&*(inserted as *const _)
};
ret
}
/// Returns a reference to the value corresponding to the key.
///
/// The key may be any borrowed form of the map's key type, but
/// [`Hash`] and [`Eq`] on the borrowed form *must* match those for
/// the key type.
///
/// # Examples
///
/// ```
/// use elsa::sync::FrozenMap;
///
/// let map = FrozenMap::new();
/// map.insert(1, Box::new("a"));
/// assert_eq!(map.get(&1), Some(&"a"));
/// assert_eq!(map.get(&2), None);
/// ```
pub fn get<Q: ?Sized>(&self, k: &Q) -> Option<&V::Target>
where
K: Borrow<Q>,
Q: Hash + Eq,
{
let map = self.map.read().unwrap();
let ret = unsafe { map.get(k).map(|x| &*(&**x as *const V::Target)) };
ret
}
/// Applies a function to the owner of the value corresponding to the key (if any).
///
/// The key may be any borrowed form of the map's key type, but
/// [`Hash`] and [`Eq`] on the borrowed form *must* match those for
/// the key type.
///
/// # Examples
///
/// ```
/// use elsa::sync::FrozenMap;
///
/// let map = FrozenMap::new();
/// map.insert(1, Box::new("a"));
/// assert_eq!(map.map_get(&1, Clone::clone), Some(Box::new("a")));
/// assert_eq!(map.map_get(&2, Clone::clone), None);
/// ```
pub fn map_get<Q: ?Sized, T, F>(&self, k: &Q, f: F) -> Option<T>
where
K: Borrow<Q>,
Q: Hash + Eq,
F: FnOnce(&V) -> T,
{
let map = self.map.read().unwrap();
let ret = map.get(k).map(f);
ret
}
/// # Examples
///
/// ```
/// use elsa::sync::FrozenMap;
///
/// let map = FrozenMap::new();
/// assert_eq!(map.len(), 0);
/// map.insert(1, Box::new("a"));
/// assert_eq!(map.len(), 1);
/// ```
pub fn len(&self) -> usize {
let map = self.map.read().unwrap();
map.len()
}
/// # Examples
///
/// ```
/// use elsa::sync::FrozenMap;
///
/// let map = FrozenMap::new();
/// assert_eq!(map.is_empty(), true);
/// map.insert(1, Box::new("a"));
/// assert_eq!(map.is_empty(), false);
/// ```
pub fn is_empty(&self) -> bool {
let map = self.map.read().unwrap();
map.is_empty()
}
// TODO add more
}
/// Append-only threadsafe version of `std::vec::Vec` where
/// insertion does not require mutable access
pub struct FrozenVec<T> {
vec: RwLock<Vec<T>>,
}
impl<T> Default for FrozenVec<T> {
fn default() -> Self {
Self {
vec: Default::default(),
}
}
}
impl<T: StableDeref> FrozenVec<T> {
pub fn new() -> Self {
Default::default()
}
// these should never return &T
// these should never delete any entries
pub fn push(&self, val: T) {
let mut vec = self.vec.write().unwrap();
vec.push(val);
}
/// Push, immediately getting a reference to the element
pub fn push_get(&self, val: T) -> &T::Target {
let mut vec = self.vec.write().unwrap();
vec.push(val);
unsafe { &*(&**vec.get_unchecked(vec.len() - 1) as *const T::Target) }
}
/// Push, immediately getting a an index of the element
///
/// Index can then be used with the `get` method
///
/// # Examples
///
/// ```
/// use elsa::sync::FrozenVec;
///
/// let map = FrozenVec::new();
/// let idx = map.push_get_index(String::from("a"));
/// assert_eq!(map.get(idx), Some("a"));
/// assert_eq!(idx, 0);
/// assert_eq!(map.push_get_index(String::from("b")), 1);
/// ```
pub fn push_get_index(&self, val: T) -> usize {
let mut vec = self.vec.write().unwrap();
let index = vec.len();
vec.push(val);
return index;
}
pub fn get(&self, index: usize) -> Option<&T::Target> {
let vec = self.vec.read().unwrap();
unsafe { vec.get(index).map(|x| &*(&**x as *const T::Target)) }
}
// TODO add more
}
/// Append-only threadsafe version of `std::vec::Vec` where
/// insertion does not require mutable access.
/// Does not have locks, only allows `Copy` types and will
/// spinlock on contention. The spinlocks are really rare as
/// they only happen on reallocation due to a push going over
/// the capacity.
pub struct LockFreeFrozenVec<T: Copy> {
data: AtomicPtr<T>,
len: AtomicUsize,
cap: AtomicUsize,
}
impl<T: Copy> Drop for LockFreeFrozenVec<T> {
fn drop(&mut self) {
let cap = *self.cap.get_mut();
let layout = self.layout(cap);
unsafe {
std::alloc::dealloc((*self.data.get_mut()).cast(), layout);
}
}
}
impl<T: Copy> Default for LockFreeFrozenVec<T> {
fn default() -> Self {
Self {
// FIXME: use `std::ptr::invalid_mut()` once that is stable.
data: AtomicPtr::new(std::mem::align_of::<T>() as *mut T),
len: AtomicUsize::new(0),
cap: AtomicUsize::new(0),
}
}
}
impl<T: Copy> LockFreeFrozenVec<T> {
pub fn new() -> Self {
Default::default()
}
pub fn with_capacity(cap: usize) -> Self {
Self {
data: AtomicPtr::new(
Box::into_raw(vec![MaybeUninit::<T>::uninit(); cap].into_boxed_slice()).cast(),
),
len: AtomicUsize::new(0),
cap: AtomicUsize::new(cap),
}
}
fn lock<U>(&self, f: impl FnOnce(&mut *mut T) -> U) -> U {
let mut ptr;
loop {
ptr = self.data.swap(std::ptr::null_mut(), Ordering::Acquire);
if !ptr.is_null() {
// Wheeeee spinlock
break;
}
}
let ret = f(&mut ptr);
self.data.store(ptr, Ordering::Release);
ret
}
fn layout(&self, cap: usize) -> Layout {
let num_bytes = std::mem::size_of::<T>() * cap;
let align = std::mem::align_of::<T>();
Layout::from_size_align(num_bytes, align).unwrap()
}
// these should never return &T
// these should never delete any entries
const NOT_ZST: () = if std::mem::size_of::<T>() == 0 {
panic!("`LockFreeFrozenVec` cannot be used with ZSTs");
};
pub fn push(&self, val: T) -> usize {
// This statement actually does something: it evaluates a constant.
#[allow(path_statements)]
{
Self::NOT_ZST
}
self.lock(|ptr| {
// These values must be consistent with the pointer we got.
let len = self.len.load(Ordering::Acquire);
let cap = self.cap.load(Ordering::Acquire);
if len >= cap {
if cap == 0 {
// No memory allocated yet
let layout = self.layout(128);
// SAFETY: `LockFreeFrozenVec` statically rejects zsts
unsafe {
*ptr = std::alloc::alloc(layout).cast::<T>();
}
// This is written before the end of the `lock` closure, so no one will observe this
// until the data pointer has been updated anyway.
self.cap.store(128, Ordering::Release);
} else {
// Out of memory, realloc with double the capacity
let layout = self.layout(cap);
let new_size = layout.size() * 2;
// SAFETY: `LockFreeFrozenVec` statically rejects zsts and the input `ptr` has always been
// allocated at the size stated in `cap`.
unsafe {
*ptr = std::alloc::realloc((*ptr).cast(), layout, new_size).cast::<T>();
}
// This is written before the end of the `lock` closure, so no one will observe this
// until the data pointer has been updated anyway.
self.cap.store(cap * 2, Ordering::Release);
}
assert!(!ptr.is_null());
}
unsafe {
ptr.add(len).write(val);
}
// This is written before updating the data pointer. Other `push` calls cannot observe this,
// because they are blocked on aquiring the data pointer before they ever read the `len`.
// `get` may read the length before actually aquiring the data pointer lock, but that is fine,
// as once it is able to aquire the lock, there will be actually the right number of elements
// stored.
self.len.store(len + 1, Ordering::Release);
len
})
}
pub fn get(&self, index: usize) -> Option<T> {
// The length can only grow, so just doing the length check
// independently of the `lock` and read is fine. Worst case we
// read an old length value and end up returning `None` even if
// another thread already inserted the value.
let len = self.len.load(Ordering::Relaxed);
if index >= len {
return None;
}
self.lock(|ptr| Some(unsafe { ptr.add(index).read() }))
}
}
#[test]
fn test_non_lockfree() {
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
struct Moo(i32);
for vec in [
LockFreeFrozenVec::new(),
LockFreeFrozenVec::with_capacity(1),
LockFreeFrozenVec::with_capacity(2),
LockFreeFrozenVec::with_capacity(1000),
] {
assert_eq!(vec.get(1), None);
vec.push(Moo(1));
let i = vec.push(Moo(2));
vec.push(Moo(3));
assert_eq!(vec.get(i), Some(Moo(2)));
std::thread::scope(|s| {
s.spawn(|| {
for i in 0..1000 {
vec.push(Moo(i));
}
});
s.spawn(|| {
for i in 0..1000 {
vec.push(Moo(i));
}
});
for i in 0..2000 {
while vec.get(i).is_none() {}
}
});
}
}
/// Append-only threadsafe version of `std::collections::BTreeMap` where
/// insertion does not require mutable access
#[derive(Debug)]
pub struct FrozenBTreeMap<K, V>(RwLock<BTreeMap<K, V>>);
impl<K: Clone + Ord, V: StableDeref> FrozenBTreeMap<K, V> {
pub fn new() -> Self {
Self(RwLock::new(BTreeMap::new()))
}
// these should never return &K or &V
// these should never delete any entries
/// Returns a reference to the value corresponding to the key.
///
/// The key may be any borrowed form of the map's key type, but
/// [`Ord`] on the borrowed form *must* match those for
/// the key type.
///
/// # Examples
///
/// ```
/// use elsa::sync::FrozenBTreeMap;
///
/// let map = FrozenBTreeMap::new();
/// map.insert(1, Box::new("a"));
/// assert_eq!(map.get(&1), Some(&"a"));
/// assert_eq!(map.get(&2), None);
/// ```
pub fn get<Q: ?Sized>(&self, k: &Q) -> Option<&V::Target>
where
K: Borrow<Q>,
Q: Ord,
{
let map = self.0.read().unwrap();
let ret = unsafe { map.get(k).map(|x| &*(&**x as *const V::Target)) };
ret
}
/// Insert a new value into the map. Does nothing if the key is already occupied.
///
/// # Examples
///
/// ```
/// use elsa::sync::FrozenBTreeMap;
///
/// let map = FrozenBTreeMap::new();
/// map.insert(1, Box::new("a"));
/// assert_eq!(map.get(&1), Some(&"a"));
/// ```
pub fn insert(&self, k: K, v: V) -> &V::Target {
let mut map = self.0.write().unwrap();
let ret = unsafe {
let inserted = &**map.entry(k).or_insert(v);
&*(inserted as *const _)
};
ret
}
/// Applies a function to the owner of the value corresponding to the key (if any).
///
/// The key may be any borrowed form of the map's key type, but
/// [`Ord`] on the borrowed form *must* match those for
/// the key type.
///
/// # Examples
///
/// ```
/// use elsa::sync::FrozenBTreeMap;
///
/// let map = FrozenBTreeMap::new();
/// map.insert(1, Box::new("a"));
/// assert_eq!(map.map_get(&1, Clone::clone), Some(Box::new("a")));
/// assert_eq!(map.map_get(&2, Clone::clone), None);
/// ```
pub fn map_get<Q: ?Sized, T, F>(&self, k: &Q, f: F) -> Option<T>
where
K: Borrow<Q>,
Q: Ord,
F: FnOnce(&V) -> T,
{
let map = self.0.read().unwrap();
let ret = map.get(k).map(f);
ret
}
/// # Examples
///
/// ```
/// use elsa::sync::FrozenBTreeMap;
///
/// let map = FrozenBTreeMap::new();
/// assert_eq!(map.len(), 0);
/// map.insert(1, Box::new("a"));
/// assert_eq!(map.len(), 1);
/// ```
pub fn len(&self) -> usize {
let map = self.0.read().unwrap();
map.len()
}
/// # Examples
///
/// ```
/// use elsa::sync::FrozenBTreeMap;
///
/// let map = FrozenBTreeMap::new();
/// assert_eq!(map.is_empty(), true);
/// map.insert(1, Box::new("a"));
/// assert_eq!(map.is_empty(), false);
/// ```
pub fn is_empty(&self) -> bool {
let map = self.0.read().unwrap();
map.is_empty()
}
}
impl<K: Clone + Ord, V: StableDeref> From<BTreeMap<K, V>> for FrozenBTreeMap<K, V> {
fn from(map: BTreeMap<K, V>) -> Self {
Self(RwLock::new(map))
}
}
impl<Q: ?Sized, K, V> Index<&Q> for FrozenBTreeMap<K, V>
where
Q: Ord,
K: Clone + Ord + Borrow<Q>,
V: StableDeref,
{
type Output = V::Target;
/// # Examples
///
/// ```
/// use elsa::sync::FrozenBTreeMap;
///
/// let map = FrozenBTreeMap::new();
/// map.insert(1, Box::new("a"));
/// assert_eq!(map[&1], "a");
/// ```
fn index(&self, idx: &Q) -> &V::Target {
self.get(idx)
.expect("attempted to index FrozenBTreeMap with unknown key")
}
}
impl<K: Clone + Ord, V: StableDeref> FromIterator<(K, V)> for FrozenBTreeMap<K, V> {
fn from_iter<T>(iter: T) -> Self
where
T: IntoIterator<Item = (K, V)>,
{
let map: BTreeMap<_, _> = iter.into_iter().collect();
map.into()
}
}
impl<K: Clone + Ord, V: StableDeref> Default for FrozenBTreeMap<K, V> {
fn default() -> Self {
Self::new()
}
}