shiny/crates/std/src/arena/intern.rs
2024-10-16 01:30:01 +03:00

147 lines
3.5 KiB
Rust

//! # Interning arena.
//!
//! Additionally, if same value was present on the arena, doesn't add new object, returns old.
use std::hash::{BuildHasherDefault, Hash};
use hashbrown::hash_map::{HashMap, RawEntryMut};
use super::{ptr::Ptr, wonly};
// No raw_entry API on hashset.
type Bi<P> = HashMap<Mapping<P>, (), BuildHasherDefault<hasher::NoHash>>;
#[derive(Debug, Default, Clone, Copy)]
struct Mapping<P> {
hash: u64,
key: P,
}
impl<P> Hash for Mapping<P> {
fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
// [`hasher::NoHash`] implementation ensures that
// hasher stores `write_u64` result only once.
state.write_u64(self.hash);
}
}
#[derive(Clone)]
pub struct Arena<T, P> {
objects: wonly::Arena<T, P>,
bi: Bi<P>,
}
impl<T, P> Default for Arena<T, P> {
fn default() -> Self {
Self {
objects: wonly::Arena::new(),
bi: Bi::default(),
}
}
}
impl<T, P> Arena<T, P> {
pub fn modify(&mut self, ptr: P, f: impl FnOnce(&mut T)) -> bool
where
T: Hash + Eq,
P: Ptr,
{
let Some(object) = self.objects.get(ptr.clone()) else {
return false;
};
let hash = fxhash::hash64(object);
let entry = self
.bi
.raw_entry_mut()
.from_hash(hash, |m| self.objects.get(m.key.clone()) == Some(object));
// SAFETY: if we got here, the object is in arena.
let object = unsafe { self.objects.get_mut(ptr.clone()).unwrap_unchecked() };
f(object);
let hash = fxhash::hash64(object);
entry.and_modify(|k, _| *k = Mapping { hash, key: ptr });
true
}
pub fn get(&self, ptr: P) -> Option<&T>
where
P: Ptr,
{
self.objects.get(ptr)
}
pub fn insert(&mut self, object: T) -> P
where
P: Ptr,
T: Hash + Eq,
{
let hash = fxhash::hash64(&object);
match self.bi.raw_entry_mut().from_hash(hash, |m| {
self.objects
.get(m.key.clone())
.map_or(false, |lhs| lhs == &object)
}) {
RawEntryMut::Occupied(occup) => occup.into_key().key.clone(),
RawEntryMut::Vacant(vacant) => {
let key = self.objects.insert(object);
vacant.insert_hashed_nocheck(
hash,
Mapping {
hash,
key: key.clone(),
},
(),
);
key
}
}
}
pub fn new() -> Self {
Self::default()
}
}
mod hasher;
#[cfg(test)]
mod tests {
use super::*;
use crate::arena::ptr::define;
define! {
struct Token;
}
#[test]
fn modification_works() {
// Ensure that when modifying the value,
// pointer is still valid.
let mut arena: Arena<u32, Token> = Arena::new();
let key100to200 = arena.insert(100);
let modified = arena.modify(key100to200, |x| *x = 200);
assert!(modified);
assert_eq!(arena.get(key100to200), Some(&200));
let key100new = arena.insert(100);
assert_ne!(key100to200, key100new);
assert_eq!(arena.get(key100new), Some(&100));
assert_eq!(arena.get(key100to200), Some(&200));
}
#[test]
fn it_interns() {
let mut arena: Arena<u32, Token> = Arena::new();
let key100 = arena.insert(100);
let key100again = arena.insert(100);
assert_eq!(key100, key100again);
}
}