use hashbrown::{
DefaultHashBuilder, HashTable, TryReserveError,
hash_table::{Entry as RawEntry, Iter as RawIter},
};
use crate::{Allocator, Ephemeron, Finalize, Gc, GcRefCell, Trace, custom_trace};
use std::{fmt, hash::BuildHasher, marker::PhantomData};
#[derive(Clone, Debug, Default, Finalize)]
pub struct WeakMap<K: Trace + ?Sized + 'static, V: Trace + 'static> {
pub(crate) inner: Gc<GcRefCell<RawWeakMap<K, V>>>,
}
unsafe impl<K: Trace + ?Sized + 'static, V: Trace + 'static> Trace for WeakMap<K, V> {
custom_trace!(this, mark, {
mark(&this.inner);
});
}
impl<K: Trace + ?Sized, V: Trace + Clone> WeakMap<K, V> {
#[must_use]
#[inline]
pub fn new() -> Self {
Allocator::alloc_weak_map()
}
#[inline]
pub fn insert(&mut self, key: &Gc<K>, value: V) {
self.inner.borrow_mut().insert(key, value);
}
#[inline]
pub fn remove(&mut self, key: &Gc<K>) -> Option<V> {
self.inner.borrow_mut().remove(key)
}
#[must_use]
#[inline]
pub fn contains_key(&self, key: &Gc<K>) -> bool {
self.inner.borrow().contains_key(key)
}
#[must_use]
#[inline]
pub fn get(&self, key: &Gc<K>) -> Option<V> {
self.inner.borrow().get(key)
}
}
pub(crate) struct RawWeakMap<K, V, S = DefaultHashBuilder>
where
K: Trace + ?Sized + 'static,
V: Trace + 'static,
{
hash_builder: S,
table: HashTable<Ephemeron<K, V>>,
}
impl<K, V, S> Finalize for RawWeakMap<K, V, S>
where
K: Trace + ?Sized + 'static,
V: Trace + 'static,
{
}
unsafe impl<K, V, S> Trace for RawWeakMap<K, V, S>
where
K: Trace + ?Sized + 'static,
V: Trace + 'static,
{
custom_trace!(this, mark, {
for eph in this.iter() {
mark(eph);
}
});
}
impl<K, V, S> Default for RawWeakMap<K, V, S>
where
S: Default,
K: Trace + ?Sized + 'static,
V: Trace + 'static,
{
fn default() -> Self {
Self::with_hasher(Default::default())
}
}
impl<K, V> RawWeakMap<K, V, DefaultHashBuilder>
where
K: Trace + ?Sized + 'static,
V: Trace + 'static,
{
pub(crate) fn new() -> Self {
Self::default()
}
#[allow(unused)]
pub(crate) fn with_capacity(capacity: usize) -> Self {
Self::with_capacity_and_hasher(capacity, DefaultHashBuilder::default())
}
}
impl<K, V, S> RawWeakMap<K, V, S>
where
K: Trace + ?Sized + 'static,
V: Trace + 'static,
{
pub(crate) const fn with_hasher(hash_builder: S) -> Self {
Self {
hash_builder,
table: HashTable::new(),
}
}
pub(crate) fn with_capacity_and_hasher(capacity: usize, hash_builder: S) -> Self {
Self {
hash_builder,
table: HashTable::with_capacity(capacity),
}
}
#[allow(unused)]
pub(crate) const fn hasher(&self) -> &S {
&self.hash_builder
}
#[allow(unused)]
pub(crate) fn capacity(&self) -> usize {
self.table.capacity()
}
pub(crate) fn iter(&self) -> Iter<'_, K, V> {
Iter {
inner: self.table.iter(),
marker: PhantomData,
}
}
#[allow(unused)]
pub(crate) fn len(&self) -> usize {
self.table.len()
}
#[allow(unused)]
pub(crate) fn is_empty(&self) -> bool {
self.len() == 0
}
pub(crate) fn retain<F>(&mut self, mut f: F)
where
F: FnMut(&Ephemeron<K, V>) -> bool,
{
self.table.retain(|item| f(item));
}
#[allow(unused)]
pub(crate) fn clear(&mut self) {
self.table.clear();
}
}
impl<K, V, S> RawWeakMap<K, V, S>
where
K: Trace + ?Sized + 'static,
V: Trace + Clone + 'static,
S: BuildHasher,
{
#[allow(unused)]
pub(crate) fn reserve(&mut self, additional: usize) {
self.table
.reserve(additional, make_hasher(&self.hash_builder));
}
#[allow(unused)]
pub(crate) fn try_reserve(&mut self, additional: usize) -> Result<(), TryReserveError> {
self.table
.try_reserve(additional, make_hasher(&self.hash_builder))
}
#[allow(unused)]
pub(crate) fn shrink_to_fit(&mut self) {
self.table
.shrink_to(0, make_hasher::<_, V, S>(&self.hash_builder));
}
#[allow(unused)]
pub(crate) fn shrink_to(&mut self, min_capacity: usize) {
self.table
.shrink_to(min_capacity, make_hasher::<_, V, S>(&self.hash_builder));
}
pub(crate) fn get(&self, k: &Gc<K>) -> Option<V> {
if self.table.is_empty() {
None
} else {
let hash = make_hash_from_gc(&self.hash_builder, k);
self.table.find(hash, equivalent_key(k))?.value()
}
}
pub(crate) fn contains_key(&self, k: &Gc<K>) -> bool {
self.get(k).is_some()
}
pub(crate) fn insert(&mut self, k: &Gc<K>, v: V) -> Option<Ephemeron<K, V>> {
let hash = make_hash_from_gc(&self.hash_builder, k);
let hasher = make_hasher(&self.hash_builder);
let entry = self.table.entry(hash, equivalent_key(k), hasher);
let (old, slot) = match entry {
RawEntry::Occupied(occupied_entry) => {
let (v, slot) = occupied_entry.remove();
(Some(v), slot)
}
RawEntry::Vacant(vacant_entry) => (None, vacant_entry),
};
slot.insert(Ephemeron::new(k, v));
old
}
pub(crate) fn remove(&mut self, k: &Gc<K>) -> Option<V> {
let hash = make_hash_from_gc(&self.hash_builder, k);
self.table
.find_entry(hash, equivalent_key(k))
.ok()?
.remove()
.0
.value()
}
pub(crate) fn clear_expired(&mut self) {
self.retain(|eph| eph.value().is_some());
}
}
pub(crate) struct Iter<'a, K, V>
where
K: Trace + ?Sized + 'static,
V: Trace + 'static,
{
inner: RawIter<'a, Ephemeron<K, V>>,
marker: PhantomData<&'a Ephemeron<K, V>>,
}
impl<K, V> Clone for Iter<'_, K, V>
where
K: Trace + ?Sized + 'static,
V: Trace + 'static,
{
#[inline]
fn clone(&self) -> Self {
Iter {
inner: self.inner.clone(),
marker: PhantomData,
}
}
}
impl<K, V> fmt::Debug for Iter<'_, K, V>
where
K: Trace + ?Sized + 'static + fmt::Debug,
V: Trace + 'static + fmt::Debug,
{
#[inline]
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_list().entries(self.clone()).finish()
}
}
impl<'a, K, V> Iterator for Iter<'a, K, V>
where
K: Trace + ?Sized + 'static,
V: Trace + 'static,
{
type Item = &'a Ephemeron<K, V>;
#[inline]
fn next(&mut self) -> Option<Self::Item> {
self.inner.next()
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
self.inner.size_hint()
}
}
impl<K, V, S> fmt::Debug for RawWeakMap<K, V, S>
where
K: fmt::Debug + ?Sized + Trace + Finalize,
V: fmt::Debug + Trace + Finalize,
{
#[inline]
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
self.iter().fmt(f)
}
}
fn make_hasher<K, V, S>(hash_builder: &S) -> impl Fn(&Ephemeron<K, V>) -> u64 + '_
where
S: BuildHasher,
K: Trace + ?Sized + 'static,
V: Trace + 'static,
{
move |val| make_hash_from_eph::<K, V, S>(hash_builder, val)
}
fn make_hash_from_eph<K, V, S>(hash_builder: &S, eph: &Ephemeron<K, V>) -> u64
where
S: BuildHasher,
K: Trace + ?Sized + 'static,
V: Trace + 'static,
{
use std::hash::Hasher;
let mut state = hash_builder.build_hasher();
unsafe {
if let Some(val) = eph.inner().key() {
std::ptr::hash(val, &mut state);
} else {
std::ptr::hash(eph.inner_ptr().as_ptr(), &mut state);
}
}
state.finish()
}
fn make_hash_from_gc<K, S>(hash_builder: &S, gc: &Gc<K>) -> u64
where
S: BuildHasher,
K: Trace + ?Sized + 'static,
{
use std::hash::Hasher;
let mut state = hash_builder.build_hasher();
std::ptr::hash(gc.inner_ptr().as_ptr(), &mut state);
state.finish()
}
fn equivalent_key<K, V>(k: &Gc<K>) -> impl Fn(&Ephemeron<K, V>) -> bool + '_
where
K: Trace + ?Sized + 'static,
V: Trace + 'static,
{
move |eph| unsafe {
eph.inner().key().is_some_and(|val| {
let val: *const _ = val;
std::ptr::eq(val, k.inner_ptr().as_ptr())
})
}
}