use std::hash::{Hash, Hasher};
use rustc_hash::FxHasher;
use crate::cell::CopyCell;
use crate::Arena;
use crate::bloom::bloom;
#[derive(Clone, Copy)]
struct MapNode<'arena, K, V> {
pub key: K,
pub hash: u64,
pub value: CopyCell<V>,
pub left: CopyCell<Option<&'arena MapNode<'arena, K, V>>>,
pub right: CopyCell<Option<&'arena MapNode<'arena, K, V>>>,
pub next: CopyCell<Option<&'arena MapNode<'arena, K, V>>>,
}
impl<'arena, K, V> MapNode<'arena, K, V> {
pub const fn new(key: K, hash: u64, value: V) -> Self {
MapNode {
key,
hash,
value: CopyCell::new(value),
left: CopyCell::new(None),
right: CopyCell::new(None),
next: CopyCell::new(None),
}
}
}
#[derive(Clone, Copy)]
pub struct Map<'arena, K, V> {
root: CopyCell<Option<&'arena MapNode<'arena, K, V>>>,
last: CopyCell<Option<&'arena MapNode<'arena, K, V>>>,
}
impl<'arena, K, V> Default for Map<'arena, K, V> {
fn default() -> Self {
Self::new()
}
}
impl<'arena, K, V> Map<'arena, K, V> {
pub const fn new() -> Self {
Map {
root: CopyCell::new(None),
last: CopyCell::new(None),
}
}
}
impl<'arena, K, V> Map<'arena, K, V> {
#[inline]
pub fn iter(&self) -> MapIter<'arena, K, V> {
MapIter {
next: self.root.get()
}
}
#[inline]
pub fn is_empty(&self) -> bool {
self.root.get().is_none()
}
#[inline]
pub fn clear(&self) {
self.root.set(None);
}
}
impl<'arena, K, V> Map<'arena, K, V>
where
K: Eq + Hash + Copy,
V: Copy,
{
#[inline]
fn hash_key(key: &K) -> u64 {
let mut hasher = FxHasher::default();
key.hash(&mut hasher);
hasher.finish()
}
#[inline]
fn find_slot(&self, key: K, hash: u64) -> &CopyCell<Option<&'arena MapNode<'arena, K, V>>> {
let mut node = &self.root;
loop {
match node.get() {
None => return node,
Some(parent) => {
if hash == parent.hash && key == parent.key {
return node;
} else if hash < parent.hash {
node = &parent.left;
} else {
node = &parent.right;
}
}
}
}
}
#[inline]
pub fn insert(&self, arena: &'arena Arena, key: K, value: V) -> Option<V> {
let hash = Self::hash_key(&key);
let node = self.find_slot(key, hash);
match node.get() {
Some(node) => {
let old = node.value.get();
node.value.set(value);
Some(old)
},
None => {
let new = Some(&*arena.alloc(MapNode::new(key, hash, value)));
if let Some(last) = self.last.get() {
last.next.set(new);
}
self.last.set(new);
node.set(new);
None
}
}
}
#[inline]
pub fn get_key(&self, key: K) -> Option<&K> {
let hash = Self::hash_key(&key);
self.find_slot(key, hash).get().map(|node| &node.key)
}
#[inline]
pub fn get(&self, key: K) -> Option<V> {
let hash = Self::hash_key(&key);
self.find_slot(key, hash).get().map(|node| node.value.get())
}
#[inline]
pub fn contains_key(&self, key: K) -> bool {
let hash = Self::hash_key(&key);
self.find_slot(key, hash).get().is_some()
}
}
#[derive(Clone, Copy)]
pub struct BloomMap<'arena, K, V> {
filter: CopyCell<u64>,
inner: Map<'arena, K, V>,
}
impl<'arena, K, V> BloomMap<'arena, K, V> {
pub const fn new() -> Self {
BloomMap {
filter: CopyCell::new(0),
inner: Map::new(),
}
}
}
impl<'arena, K, V: Copy> BloomMap<'arena, K, V> {
#[inline]
pub fn iter(&self) -> MapIter<'arena, K, V> {
self.inner.iter()
}
#[inline]
pub fn is_empty(&self) -> bool {
self.inner.is_empty()
}
#[inline]
pub fn clear(&self) {
self.filter.set(0);
self.inner.clear();
}
}
impl<'arena, K, V> BloomMap<'arena, K, V>
where
K: Eq + Hash + Copy + AsRef<[u8]>,
V: Copy,
{
#[inline]
pub fn insert(&self, arena: &'arena Arena, key: K, value: V) -> Option<V> {
self.filter.set(self.filter.get() | bloom(key));
self.inner.insert(arena, key, value)
}
#[inline]
pub fn get(&self, key: K) -> Option<V> {
let b = bloom(key.as_ref());
if self.filter.get() & b == b {
self.inner.get(key)
} else {
None
}
}
#[inline]
pub fn contains_key(&self, key: K) -> bool {
let b = bloom(key);
self.filter.get() & b == b && self.inner.contains_key(key)
}
}
pub struct MapIter<'arena, K, V> {
next: Option<&'arena MapNode<'arena, K, V>>
}
impl<'arena, K, V: Copy> Iterator for MapIter<'arena, K, V> {
type Item = (&'arena K, V);
#[inline]
fn next(&mut self) -> Option<Self::Item> {
let next = self.next;
next.map(|map_node| {
let item = (&map_node.key, map_node.value.get());
self.next = map_node.next.get();
item
})
}
}
impl<'arena, K, V: Copy> IntoIterator for Map<'arena, K, V> {
type Item = (&'arena K, V);
type IntoIter = MapIter<'arena, K, V>;
#[inline]
fn into_iter(self) -> Self::IntoIter {
self.iter()
}
}
impl<'arena, K, V: Copy> IntoIterator for BloomMap<'arena, K, V> {
type Item = (&'arena K, V);
type IntoIter = MapIter<'arena, K, V>;
#[inline]
fn into_iter(self) -> Self::IntoIter {
self.iter()
}
}
impl<'arena, K, V> From<Map<'arena, K, V>> for BloomMap<'arena, K, V>
where
K: Eq + Hash + Copy + AsRef<[u8]>,
V: Copy,
{
fn from(map: Map<'arena, K, V>) -> BloomMap<'arena, K, V> {
let mut filter = 0;
for (key, _) in map.iter() {
filter |= bloom(key.as_ref());
}
BloomMap {
filter: CopyCell::new(filter),
inner: map,
}
}
}
impl<'arena, K, V> From<BloomMap<'arena, K, V>> for Map<'arena, K, V> {
#[inline]
fn from(bloom_map: BloomMap<'arena, K, V>) -> Map<'arena, K, V> {
bloom_map.inner
}
}
#[cfg(test)]
mod test {
use super::*;
#[test]
fn map() {
let arena = Arena::new();
let map = Map::new();
map.insert(&arena, "foo", 10u64);
map.insert(&arena, "bar", 20);
map.insert(&arena, "doge", 30);
assert_eq!(map.contains_key("foo"), true);
assert_eq!(map.contains_key("bar"), true);
assert_eq!(map.contains_key("doge"), true);
assert_eq!(map.contains_key("moon"), false);
assert_eq!(map.get("foo"), Some(10));
assert_eq!(map.get("bar"), Some(20));
assert_eq!(map.get("doge"), Some(30));
assert_eq!(map.get("moon"), None);
}
#[test]
fn bloom_map() {
let arena = Arena::new();
let map = BloomMap::new();
map.insert(&arena, "foo", 10u64);
map.insert(&arena, "bar", 20);
map.insert(&arena, "doge", 30);
assert_eq!(map.contains_key("foo"), true);
assert_eq!(map.contains_key("bar"), true);
assert_eq!(map.contains_key("doge"), true);
assert_eq!(map.contains_key("moon"), false);
assert_eq!(map.get("foo"), Some(10));
assert_eq!(map.get("bar"), Some(20));
assert_eq!(map.get("doge"), Some(30));
assert_eq!(map.get("moon"), None);
}
#[test]
fn iter() {
let arena = Arena::new();
let map = Map::new();
map.insert(&arena, "foo", 10u64);
map.insert(&arena, "bar", 20);
map.insert(&arena, "doge", 30);
let mut iter = map.iter();
assert_eq!(iter.next(), Some((&"foo", 10)));
assert_eq!(iter.next(), Some((&"bar", 20)));
assert_eq!(iter.next(), Some((&"doge", 30)));
assert_eq!(iter.next(), None);
}
#[test]
fn insert_replace() {
let arena = Arena::new();
let map = Map::new();
map.insert(&arena, "foo", 10u64);
map.insert(&arena, "bar", 20);
map.insert(&arena, "doge", 30);
let mut iter = map.iter();
assert_eq!(iter.next(), Some((&"foo", 10)));
assert_eq!(iter.next(), Some((&"bar", 20)));
assert_eq!(iter.next(), Some((&"doge", 30)));
assert_eq!(iter.next(), None);
map.insert(&arena, "bar", 42);
let mut iter = map.iter();
assert_eq!(iter.next(), Some((&"foo", 10)));
assert_eq!(iter.next(), Some((&"bar", 42)));
assert_eq!(iter.next(), Some((&"doge", 30)));
assert_eq!(iter.next(), None);
}
#[test]
fn from_eq() {
let arena = Arena::new();
let map = Map::new();
map.insert(&arena, "foo", 10);
map.insert(&arena, "bar", 20);
map.insert(&arena, "doge", 30);
let bloom_map = BloomMap::new();
bloom_map.insert(&arena, "foo", 10);
bloom_map.insert(&arena, "bar", 20);
bloom_map.insert(&arena, "doge", 30);
assert_eq!(map, Map::from(bloom_map));
assert_eq!(BloomMap::from(map), bloom_map);
}
}