use std::{
hash::{BuildHasher, Hash},
mem::MaybeUninit,
};
use hashbrown::{DefaultHashBuilder, HashMap, hash_map::RawEntryMut};
#[derive(Debug, Default)]
pub struct OverlayMap<K, V, S = DefaultHashBuilder>
where
K: Eq + Hash,
{
map: HashMap<K, Overlay<V>, S>,
}
unsafe impl<K, V, S> Sync for OverlayMap<K, V, S>
where
K: Eq + Hash + Sync,
S: Sync,
{
}
impl<K, V> OverlayMap<K, V, DefaultHashBuilder>
where
K: Eq + Hash,
{
pub fn new() -> Self {
Self::with_hasher(DefaultHashBuilder::default())
}
}
impl<K, V, S> OverlayMap<K, V, S>
where
K: Eq + Hash,
S: BuildHasher + Default,
{
pub fn with_capacity(capacity: usize) -> Self {
Self::with_capacity_and_hasher(capacity, Default::default())
}
pub fn with_hasher(hasher: S) -> Self {
Self {
map: HashMap::with_hasher(hasher),
}
}
pub fn with_capacity_and_hasher(capacity: usize, hasher: S) -> Self {
Self {
map: HashMap::with_capacity_and_hasher(capacity, hasher),
}
}
pub fn len(&self) -> usize {
self.map.len()
}
pub fn is_empty(&self) -> bool {
self.map.is_empty()
}
#[inline]
pub fn fg(&self, key: &K) -> Option<&V> {
self.map.get(key).map(|entry| entry.fg_unchecked())
}
#[inline]
pub fn bg(&self, key: &K) -> Option<&V> {
self.map.get(key).and_then(|entry| entry.bg())
}
#[inline]
pub fn push(&mut self, key: K, value: V) -> bool {
match self.map.raw_entry_mut().from_key(&key) {
RawEntryMut::Occupied(mut occupied) => {
occupied.get_mut().push(value);
true
}
RawEntryMut::Vacant(vacant) => {
vacant.insert(key, Overlay::new_fg(value));
false
}
}
}
pub fn push_if<F>(&mut self, key: &K, predicate: F) -> bool
where
F: FnOnce(&V) -> Option<V>,
{
let entry = match self.map.get_mut(key) {
Some(e) => e,
None => return false,
};
match predicate(entry.fg_unchecked()) {
Some(new) => {
entry.push(new);
true
}
None => false,
}
}
#[inline]
pub fn pull(&mut self, key: &K) -> Option<V> {
match self.map.raw_entry_mut().from_key(key) {
RawEntryMut::Occupied(mut occupied) => {
let entry = occupied.get_mut();
let evicted = entry.pull_unchecked();
if entry.is_empty() {
occupied.remove();
}
Some(evicted)
}
RawEntryMut::Vacant(_) => None,
}
}
pub fn pull_if<F>(&mut self, key: &K, predicate: F) -> Option<V>
where
F: FnOnce(&V) -> bool,
{
match self.map.raw_entry_mut().from_key(key) {
RawEntryMut::Occupied(mut occupied) => {
let entry = occupied.get_mut();
if predicate(entry.fg_unchecked()) {
let evicted = entry.pull_unchecked();
if entry.is_empty() {
occupied.remove();
}
Some(evicted)
} else {
None
}
}
RawEntryMut::Vacant(_) => None,
}
}
#[inline]
pub fn swap(&mut self, key: K, value: V) -> Option<V> {
match self.map.raw_entry_mut().from_key(&key) {
RawEntryMut::Occupied(mut occupied) => occupied.get_mut().swap(value),
RawEntryMut::Vacant(vacant) => {
vacant.insert(key, Overlay::new_fg(value));
None
}
}
}
pub fn swap_if<F>(&mut self, key: &K, predicate: F) -> Option<V>
where
F: FnOnce(&V) -> Option<V>,
{
let entry = self.map.get_mut(key)?;
match predicate(entry.fg_unchecked()) {
Some(new) => entry.swap(new),
None => None,
}
}
pub fn flip(&mut self, key: &K) {
if let Some(entry) = self.map.get_mut(key) {
entry.flip();
}
}
pub fn extend_count<I>(&mut self, iter: I) -> usize
where
I: IntoIterator<Item = (K, V)>,
{
let mut replaced = 0;
for (key, val) in iter {
replaced += self.push(key, val) as usize;
}
replaced
}
}
impl<K, V, S> Clone for OverlayMap<K, V, S>
where
K: Clone + Eq + Hash,
V: Clone,
S: Clone + BuildHasher,
{
fn clone(&self) -> Self {
Self {
map: self.map.clone(),
}
}
}
impl<K, V, S> PartialEq for OverlayMap<K, V, S>
where
K: Eq + Hash,
V: PartialEq,
S: BuildHasher,
{
fn eq(&self, other: &Self) -> bool {
self.map == other.map
}
}
impl<K, V, S> Eq for OverlayMap<K, V, S>
where
K: Eq + Hash,
V: Eq,
S: BuildHasher,
{
}
impl<K, V, S> Extend<(K, V)> for OverlayMap<K, V, S>
where
K: Eq + Hash,
S: BuildHasher + Default,
{
fn extend<I: IntoIterator<Item = (K, V)>>(&mut self, iter: I) {
for (k, v) in iter {
self.push(k, v);
}
}
}
impl<K, V, S> IntoIterator for OverlayMap<K, V, S>
where
K: Eq + Hash,
S: BuildHasher,
{
type Item = (K, Overlay<V>);
type IntoIter = hashbrown::hash_map::IntoIter<K, Overlay<V>>;
fn into_iter(self) -> Self::IntoIter {
self.map.into_iter()
}
}
const SLOT0_PRESENT: u8 = 1 << 0;
const SLOT1_PRESENT: u8 = 1 << 1;
const SLOT_MASK: u8 = SLOT0_PRESENT | SLOT1_PRESENT;
const FG_SLOT: u8 = 1 << 2;
#[derive(Debug)]
pub struct Overlay<T> {
bits: u8,
slots: [MaybeUninit<T>; 2],
}
impl<T> Overlay<T> {
pub fn new_empty() -> Self {
Self {
bits: 0,
slots: [MaybeUninit::uninit(), MaybeUninit::uninit()],
}
}
pub fn new_fg(val: T) -> Self {
Self {
bits: SLOT0_PRESENT,
slots: [MaybeUninit::new(val), MaybeUninit::uninit()],
}
}
pub fn new_both(fg: T, bg: T) -> Self {
Self {
bits: SLOT0_PRESENT | SLOT1_PRESENT,
slots: [MaybeUninit::new(fg), MaybeUninit::new(bg)],
}
}
#[inline]
pub fn fg(&self) -> Option<&T> {
let idx = self.fg_index();
if self.is_slot_present(idx) {
Some(unsafe { self.slots[idx].assume_init_ref() })
} else {
None
}
}
#[inline]
pub fn fg_unchecked(&self) -> &T {
let idx = self.fg_index();
unsafe { self.slots[idx].assume_init_ref() }
}
#[inline]
pub fn bg(&self) -> Option<&T> {
let idx = self.bg_index();
if self.is_slot_present(idx) {
Some(unsafe { self.slots[idx].assume_init_ref() })
} else {
None
}
}
#[inline]
pub fn bg_unchecked(&self) -> &T {
let idx = self.bg_index();
unsafe { self.slots[idx].assume_init_ref() }
}
#[inline]
pub fn is_empty(&self) -> bool {
(self.bits & SLOT_MASK) == 0
}
#[inline]
pub fn is_full(&self) -> bool {
(self.bits & SLOT_MASK) == SLOT_MASK
}
#[inline]
pub fn clear(&mut self) {
if (self.bits & SLOT0_PRESENT) != 0 {
unsafe { self.slots[0].assume_init_drop() };
}
if (self.bits & SLOT1_PRESENT) != 0 {
unsafe { self.slots[1].assume_init_drop() };
}
self.bits = 0;
}
#[inline]
pub fn clear_unchecked(&mut self) {
unsafe {
self.slots[0].assume_init_drop();
self.slots[1].assume_init_drop();
}
self.bits = 0;
}
#[inline]
pub fn push(&mut self, val: T) {
self.push_fg_to_bg();
let idx = self.fg_index();
self.slots[idx] = MaybeUninit::new(val);
self.bits |= 1 << idx;
}
#[inline]
pub fn pull(&mut self) -> Option<T> {
let fgi = self.fg_index();
if self.is_slot_present(fgi) {
self.bits ^= FG_SLOT | (1 << fgi);
Some(unsafe { self.slots[fgi].assume_init_read() })
} else {
None
}
}
#[inline]
pub fn pull_unchecked(&mut self) -> T {
let fgi = self.fg_index();
self.bits ^= FG_SLOT | (1 << fgi);
unsafe { self.slots[fgi].assume_init_read() }
}
#[inline]
pub fn swap(&mut self, val: T) -> Option<T> {
let bgi = self.bg_index();
if self.is_slot_present(bgi) {
let evicted = unsafe { self.slots[bgi].assume_init_read() };
self.slots[bgi] = MaybeUninit::new(val);
self.flip_unchecked();
Some(evicted)
} else {
self.push(val);
None
}
}
pub fn iter(&self) -> impl Iterator<Item = &T> {
self.fg().into_iter().chain(self.bg())
}
#[inline]
pub fn flip(&mut self) {
if (self.bits & SLOT_MASK) == SLOT_MASK {
self.bits ^= FG_SLOT;
}
}
#[inline]
pub fn flip_unchecked(&mut self) {
self.bits ^= FG_SLOT;
}
#[inline]
fn fg_index(&self) -> usize {
((self.bits & FG_SLOT) >> 2) as usize
}
#[inline]
fn bg_index(&self) -> usize {
self.fg_index() ^ 1
}
#[inline]
fn is_slot_present(&self, idx: usize) -> bool {
(self.bits & (1 << idx)) != 0
}
#[inline]
fn push_fg_to_bg(&mut self) {
let bgi = self.bg_index();
if self.is_slot_present(bgi) {
unsafe {
self.slots[bgi].assume_init_drop();
}
}
self.flip_unchecked();
}
}
impl<T> Default for Overlay<T> {
fn default() -> Self {
Self::new_empty()
}
}
impl<T> From<T> for Overlay<T> {
fn from(value: T) -> Self {
Self::new_fg(value)
}
}
impl<T: Clone> Clone for Overlay<T> {
fn clone(&self) -> Self {
let mut clone = Self::new_empty();
clone.bits = self.bits;
if self.is_slot_present(0) {
clone.slots[0] = MaybeUninit::new(unsafe { self.slots[0].assume_init_ref().clone() });
}
if self.is_slot_present(1) {
clone.slots[1] = MaybeUninit::new(unsafe { self.slots[1].assume_init_ref().clone() });
}
clone
}
}
impl<T: PartialEq> PartialEq for Overlay<T> {
fn eq(&self, other: &Self) -> bool {
if (self.bits & SLOT_MASK) != (other.bits & SLOT_MASK) {
return false;
}
self.fg() == other.fg() && self.bg() == other.bg()
}
}
impl<T: Eq> Eq for Overlay<T> {}
impl<V> Drop for Overlay<V> {
fn drop(&mut self) {
if (self.bits & SLOT0_PRESENT) != 0 {
unsafe { self.slots[0].assume_init_drop() };
}
if (self.bits & SLOT1_PRESENT) != 0 {
unsafe { self.slots[1].assume_init_drop() };
}
}
}
pub struct OverlayIntoIter<T> {
overlay: Overlay<T>,
}
impl<T> Iterator for OverlayIntoIter<T> {
type Item = T;
fn next(&mut self) -> Option<Self::Item> {
self.overlay.pull()
}
}
impl<T> IntoIterator for Overlay<T> {
type Item = T;
type IntoIter = OverlayIntoIter<T>;
fn into_iter(self) -> Self::IntoIter {
OverlayIntoIter { overlay: self }
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn integration_push_pull_cycle() {
let mut map = OverlayMap::<&str, i32>::new();
map.push("x", 1);
map.push("x", 2);
map.push("x", 3);
assert_eq!(map.fg(&"x"), Some(&3));
assert_eq!(map.bg(&"x"), Some(&2));
let pulled = map.pull(&"x");
assert_eq!(pulled, Some(3));
assert_eq!(map.fg(&"x"), Some(&2));
let pulled = map.pull(&"x");
assert_eq!(pulled, Some(2));
assert_eq!(map.fg(&"x"), None);
assert!(map.is_empty());
}
#[test]
fn extend_count_overlapping_and_new() {
let mut map = OverlayMap::<&str, i32>::new();
map.push("a", 10);
map.push("b", 20);
let replaced = map.extend_count([("a", 100), ("c", 300), ("b", 200)]);
assert_eq!(replaced, 2);
assert_eq!(map.fg(&"a"), Some(&100));
assert_eq!(map.bg(&"a"), Some(&10));
assert_eq!(map.fg(&"b"), Some(&200));
assert_eq!(map.bg(&"b"), Some(&20));
assert_eq!(map.fg(&"c"), Some(&300));
assert_eq!(map.bg(&"c"), None);
}
#[test]
fn push_if_and_swap_if_logic() {
let mut map = OverlayMap::<&str, i32>::new();
map.push("key", 1);
let pushed = map.push_if(&"key", |v| if *v < 5 { Some(*v + 10) } else { None });
assert!(pushed);
assert_eq!(map.fg(&"key"), Some(&11));
assert_eq!(map.bg(&"key"), Some(&1));
let evicted = map.swap_if(&"key", |v| if *v == 11 { Some(42) } else { None });
assert_eq!(evicted, Some(1));
assert_eq!(map.fg(&"key"), Some(&42));
assert_eq!(map.bg(&"key"), Some(&11));
}
}