use std::sync::atomic::AtomicBool;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use std::sync::Arc;
pub struct MarkOnDrop {
marker: Arc<AtomicBool>,
gc: Arc<AtomicUsize>,
}
impl Drop for MarkOnDrop {
fn drop(&mut self) {
self.marker.store(true, Ordering::SeqCst);
self.gc.fetch_add(1, Ordering::SeqCst);
}
}
pub struct HashMap<K, V> {
v: std::collections::HashMap<K, (V, Arc<AtomicBool>)>,
gc: Arc<AtomicUsize>,
}
impl<K,V> Default for HashMap<K,V>
where K: std::cmp::Eq + std::hash::Hash
{
fn default() -> Self {
HashMap {
v: std::collections::HashMap::new(),
gc: Arc::new(AtomicUsize::new(0)),
}
}
}
impl<K, V> HashMap<K, V>
where K: std::cmp::Eq + std::hash::Hash
{
pub fn new() -> Self {
Self::default()
}
}
impl<K,V> HashMap<K,V>
where K: std::cmp::Eq + std::hash::Hash
{
pub fn insert(&mut self, k: K, v: V) -> (MarkOnDrop, Option<V>)
where K: std::cmp::Eq + std::hash::Hash
{
self.maybe_gc();
let mark = MarkOnDrop {
marker: Arc::new(AtomicBool::new(false)),
gc: self.gc.clone(),
};
let old = match self.v.insert(k, (v, mark.marker.clone())) {
None => None,
Some((v, marker)) => {
if marker.load(Ordering::SeqCst) == false {
Some(v)
} else {
None
}
}
};
(mark, old)
}
pub fn get<Q: ?Sized>(&mut self, k: &Q) -> Option<&V>
where Q: std::cmp::Eq + std::hash::Hash,
K: std::borrow::Borrow<Q>,
{
let remove = if let Some((_, marker)) = self.v.get(k) {
marker.load(Ordering::SeqCst)
} else {
false
};
if remove {
self.v.remove(k);
}
self.v.get(k).map(|(v,_)|v)
}
pub fn get_mut<Q: ?Sized>(&mut self, k: &Q) -> Option<&mut V>
where Q: std::cmp::Eq + std::hash::Hash,
K: std::borrow::Borrow<Q>,
{
let remove = if let Some((_, marker)) = self.v.get(k) {
marker.load(Ordering::SeqCst)
} else {
false
};
if remove {
self.v.remove(k);
}
self.v.get_mut(k).map(|(v,_)|v)
}
pub fn len(&self) -> usize {
self.v.len()
}
fn maybe_gc(&mut self) {
if self.gc.load(Ordering::SeqCst) > self.len() / 2 {
self.gc();
}
}
pub fn gc(&mut self) {
self.gc.store(0, Ordering::SeqCst);
self.v.retain(|_, (_, marker)| {
!marker.load(Ordering::SeqCst)
})
}
pub fn entry(&mut self, k: K) -> Entry<K, V> {
self.maybe_gc();
let remove = if let Some((_, marker)) = self.v.get(&k) {
marker.load(Ordering::SeqCst)
} else {
false
};
if remove {
self.v.remove(&k);
}
match self.v.entry(k) {
std::collections::hash_map::Entry::Occupied(n) => {
Entry::Occupied(OccupiedEntry{n})
},
std::collections::hash_map::Entry::Vacant(n) => {
Entry::Vacant(VacantEntry{n, gc: self.gc.clone()})
},
}
}
pub fn iter(&self) -> impl Iterator<Item = (&K, &V)> {
self.v.iter().filter_map(|(k, (v,marker))|{
if marker.load(Ordering::SeqCst) == true {
None
} else {
Some((k,v))
}
})
}
pub fn iter_mut(&mut self) -> impl Iterator<Item = (&K, &mut V)> {
self.v.iter_mut().filter_map(|(k, (v,marker))|{
if marker.load(Ordering::SeqCst) == true {
None
} else {
Some((k,v))
}
})
}
}
pub struct OccupiedEntry<'a, K: 'a, V: 'a>{
n: std::collections::hash_map::OccupiedEntry<'a, K, (V,Arc<AtomicBool>)>,
}
pub struct VacantEntry<'a, K: 'a, V: 'a>{
n: std::collections::hash_map::VacantEntry<'a, K, (V,Arc<AtomicBool>)>,
gc: Arc<AtomicUsize>,
}
pub enum Entry<'a, K: 'a, V: 'a> {
Occupied(OccupiedEntry<'a, K, V>),
Vacant(VacantEntry<'a, K, V>),
}
impl<'a, K, V> OccupiedEntry<'a, K, V> {
pub fn into_mut(self) -> &'a mut V {
&mut self.n.into_mut().0
}
}
impl<'a, K, V> VacantEntry<'a, K, V> {
pub fn insert_with<F: FnOnce(MarkOnDrop) -> V>(self, value: F) -> &'a mut V {
let mark = MarkOnDrop {
marker: Arc::new(AtomicBool::new(false)),
gc: self.gc.clone(),
};
let marker = mark.marker.clone();
&mut (self.n.insert((value(mark), marker)).0)
}
}
impl<'a, K, V> Entry<'a, K, V> {
pub fn or_insert_with<F: FnOnce(MarkOnDrop) -> V>(self, default: F) -> &'a mut V {
match self {
Entry::Occupied(entry) => entry.into_mut(),
Entry::Vacant(entry) => {
entry.insert_with(default)
}
}
}
}
#[test]
fn iter() {
let mut wm : HashMap<u32, u8> = HashMap::new();
let (mark, _) = wm.insert(1, 7);
drop(mark);
let (mark, _) = wm.insert(2, 8);
let mut iter = wm.iter();
assert_eq!(iter.next(), Some((&2, &8)));
assert_eq!(iter.next(), None);
}
#[test]
fn entry() {
let mut wm : HashMap<u32, u8> = HashMap::new();
let mut holdme = None;
{
let val = wm.entry(1).or_insert_with(|mark|{
holdme = Some(mark);
2
});
*val = 3;
}
assert_eq!(wm.get(&1), Some(&3));
drop(holdme);
assert_eq!(wm.get(&1), None);
{
let val = wm.entry(1).or_insert_with(|mark|2);
*val = 3;
}
assert_eq!(wm.get(&1), None);
}
#[test]
fn foo() {
let mut wm : HashMap<u32, &'static str> = HashMap::new();
let marks : Vec<MarkOnDrop> = (0..100000).map(|i|{
let (mark, _) = wm.insert(i + 100000, "world");
drop(mark);
assert_eq!(wm.get(&(i + 100000)), None);
let (mark, _) = wm.insert(i, "world");
mark
}).collect();
assert_eq!(wm.get(&1), Some(&"world"));
drop(marks);
assert_eq!(wm.get(&1), None);
}