#![feature(box_into_raw_non_null)]
use core::borrow::Borrow;
use core::cmp::Ordering::{Less, Equal, Greater};
use core::fmt;
use core::iter::FusedIterator;
use core::ptr::{self, NonNull};
use core::mem;
use core::marker::PhantomData;
const MAX_HEIGHT: usize = 12;
type Tower<K, V> = Vec<Option<NonNull<Node<K, V>>>>;
pub struct SkipMap<K: Ord, V> {
len: usize,
tower: Tower<K, V>,
_marker: PhantomData<Box<Node<K, V>>>,
}
impl<K: Ord, V> SkipMap<K, V> {
pub fn new() -> Self {
Self {
len: 0,
tower: vec![None; MAX_HEIGHT],
_marker: PhantomData
}
}
pub fn len(&self) -> usize {
self.len
}
pub fn is_empty(&self) -> bool {
self.len == 0
}
pub fn iter(&self) -> Iter<K, V> {
Iter {
head: self.tower[0],
len: self.len,
_marker: PhantomData
}
}
}
impl<K: Ord, V> SkipMap<K, V> {
pub fn insert(&mut self, key: K, value: V) -> Option<V> {
let mut cur_tower = &self.tower as *const _ as *mut Tower<K, V>;
let mut h = self.tower.len() - 1;
let new_height = next_height();
let mut to_modify = Vec::with_capacity(new_height);
loop {
let mut seek_down = false;
if let Some(ref mut next_ptr) = unsafe { &mut *cur_tower }[h] {
let next = unsafe { next_ptr.as_ref() };
match next.key.cmp(&key) {
Less => cur_tower = &next.tower as *const _ as *mut _,
Greater => seek_down = true,
Equal => {
let mut ans = value;
unsafe { next_ptr.as_mut() }.key = key;
mem::swap(&mut unsafe { next_ptr.as_mut() }.value, &mut ans);
return Some(ans);
},
}
} else {
seek_down = true;
}
if seek_down {
if h < new_height {
to_modify.push((cur_tower, h));
}
if h > 0 {
h -= 1;
} else {
break;
}
}
}
let mut new_node = Box::new(Node::new(key, value, new_height));
for (prev_tower_ptr, h) in to_modify.iter() {
if let Some(place_mut) = new_node.tower.get_mut(*h) {
*place_mut = unsafe { &**prev_tower_ptr }[*h];
}
}
let new_node_ptr = Box::into_raw_non_null(new_node);
for (prev_tower_ptr, h) in to_modify.iter() {
unsafe { (**prev_tower_ptr)[*h] = Some(new_node_ptr) };
}
self.len += 1;
None
}
pub fn clear(&mut self) {
let mut cur_opt = self.tower[0];
self.tower = vec![None; MAX_HEIGHT];
while let Some(next_ptr) = cur_opt {
let next_node = unsafe { Box::from_raw(next_ptr.as_ptr()) };
cur_opt = next_node.tower[0];
drop(next_node);
self.len -= 1;
}
}
pub fn get<Q: ?Sized>(&self, key: &Q) -> Option<&V>
where K: Borrow<Q>, Q: Ord {
let mut cur_tower = &self.tower as *const _ as *mut Tower<K, V>;
let mut h = self.tower.len() - 1;
loop {
let mut seek_down = false;
if let Some(ref mut next_ptr) = unsafe { &mut *cur_tower }[h] {
let next = unsafe { next_ptr.as_ref() };
match next.key.borrow().cmp(&key) {
Less => cur_tower = &next.tower as *const _ as *mut _,
Greater => seek_down = true,
Equal => return Some(&next.value)
}
} else {
seek_down = true;
}
if seek_down {
if h > 0 {
h -= 1;
} else {
return None;
}
}
}
}
pub fn remove<Q: ?Sized>(&mut self, key: &Q) -> Option<V>
where K: Borrow<Q>, Q: Ord {
let mut cur_tower = &self.tower as *const _ as *mut Tower<K, V>;
let mut h = self.tower.len() - 1;
let cur_height = MAX_HEIGHT;
let mut to_modify: Vec<(*mut Tower<K, V>, usize)> = Vec::with_capacity(cur_height);
loop {
let mut seek_down = false;
if let Some(ref mut next_ptr) = unsafe { &mut *cur_tower }[h] {
let next = unsafe { next_ptr.as_ref() };
match next.key.borrow().cmp(&key) {
Less => cur_tower = &next.tower as *const _ as *mut _,
Greater => seek_down = true,
Equal => {
let value = unsafe { ptr::read(next_ptr.as_ptr()) }.value;
for (prev_tower_ptr, h) in to_modify.iter() {
if *h < unsafe { next_ptr.as_ref() }.tower.len() {
unsafe { (**prev_tower_ptr)[*h] = next_ptr.as_ref().tower[*h] };
}
}
self.len -= 1;
return Some(value);
},
}
} else {
seek_down = true;
}
if seek_down {
if h < cur_height {
to_modify.push((cur_tower, h));
}
if h > 0 {
h -= 1;
} else {
return None;
}
}
}
}
}
impl<K: Ord + fmt::Debug, V: fmt::Debug> fmt::Debug for SkipMap<K, V> {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
f.debug_list().entries(self.iter()).finish()
}
}
impl<K: Ord, V> Drop for SkipMap<K, V> {
fn drop(&mut self) {
self.clear();
}
}
pub struct Iter<'a, K: 'a + Ord, V: 'a> {
head: Option<NonNull<Node<K, V>>>,
len: usize,
_marker: PhantomData<(&'a K, &'a V)>,
}
impl<'a, K: Ord, V> Iterator for Iter<'a, K, V> {
type Item = (&'a K, &'a V);
#[inline]
fn next(&mut self) -> Option<Self::Item> {
if let Some(next_ptr) = self.head {
let next_node = unsafe { &*next_ptr.as_ptr() };
self.head = next_node.tower[0];
self.len -= 1;
Some((&next_node.key, &next_node.value))
} else {
None
}
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
(self.len, Some(self.len))
}
}
impl<'a, K: Ord + fmt::Debug, V> fmt::Debug for Iter<'a, K, V> {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
f.debug_tuple("Iter").field(&self.len).finish()
}
}
impl<'a, K: Ord, V> ExactSizeIterator for Iter<'a, K, V> {}
impl<'a, K: Ord, V> FusedIterator for Iter<'a, K, V> {}
#[derive(Debug)]
struct Node<K: Ord, V> {
key: K,
value: V,
tower: Tower<K, V>,
_marker: PhantomData<Box<Node<K, V>>>,
}
impl<K: Ord, V> Node<K, V> {
fn new(key: K, value: V, height: usize) -> Self {
Self {
key,
value,
tower: vec![None; height],
_marker: PhantomData,
}
}
}
fn next_height() -> usize {
use rand::distributions::{Distribution, Bernoulli};
let mut ans = 1;
let mut rng = rand::thread_rng();
let half_half = Bernoulli::from_ratio(1, 2);
while half_half.sample(&mut rng) {
ans += 1;
if ans == MAX_HEIGHT {
break;
}
}
ans
}
#[test]
fn skip_map_insert_test() {
let generate_value = |key| key + 10000;
let test = |times, numbers| {
for _ in 0..times {
let mut map = SkipMap::new();
assert!(map.is_empty());
for i in (0..numbers).rev() {
map.insert(i, generate_value(i));
}
assert_eq!(map.len(), numbers);
let mut i = 0;
for (key, value) in map.iter() {
assert_eq!(*key, i);
assert_eq!(*value, generate_value(i));
i += 1;
}
map.clear();
assert_eq!(map.len(), 0);
}
};
test(1000, 10);
test(10, 1000);
}
#[test]
fn skip_map_remove_test() {
for _ in 0..1000 {
let mut map = SkipMap::new();
assert!(map.is_empty());
for i in (0..10).rev() {
map.insert(i, i + 1000);
}
println!("{:?}", map);
for i in (3..8).rev() {
map.remove(&i);
}
println!("{:?}", map);
}
}