use super::{Comparator, Forest, Node, NodeData, NodePool, Path, INNER_SIZE};
use crate::packed_option::PackedOption;
#[cfg(test)]
use alloc::string::String;
#[cfg(test)]
use core::fmt;
use core::marker::PhantomData;
struct MapTypes<K, V>(PhantomData<(K, V)>);
impl<K, V> Forest for MapTypes<K, V>
where
K: Copy,
V: Copy,
{
type Key = K;
type Value = V;
type LeafKeys = [K; INNER_SIZE - 1];
type LeafValues = [V; INNER_SIZE - 1];
fn splat_key(key: Self::Key) -> Self::LeafKeys {
[key; INNER_SIZE - 1]
}
fn splat_value(value: Self::Value) -> Self::LeafValues {
[value; INNER_SIZE - 1]
}
}
pub struct MapForest<K, V>
where
K: Copy,
V: Copy,
{
nodes: NodePool<MapTypes<K, V>>,
}
impl<K, V> MapForest<K, V>
where
K: Copy,
V: Copy,
{
pub fn new() -> Self {
Self {
nodes: NodePool::new(),
}
}
pub fn clear(&mut self) {
self.nodes.clear();
}
}
#[derive(Clone)]
pub struct Map<K, V>
where
K: Copy,
V: Copy,
{
root: PackedOption<Node>,
unused: PhantomData<(K, V)>,
}
impl<K, V> Map<K, V>
where
K: Copy,
V: Copy,
{
pub fn new() -> Self {
Self {
root: None.into(),
unused: PhantomData,
}
}
pub fn is_empty(&self) -> bool {
self.root.is_none()
}
pub fn get<C: Comparator<K>>(&self, key: K, forest: &MapForest<K, V>, comp: &C) -> Option<V> {
self.root
.expand()
.and_then(|root| Path::default().find(key, root, &forest.nodes, comp))
}
pub fn get_or_less<C: Comparator<K>>(
&self,
key: K,
forest: &MapForest<K, V>,
comp: &C,
) -> Option<(K, V)> {
self.root.expand().and_then(|root| {
let mut path = Path::default();
match path.find(key, root, &forest.nodes, comp) {
Some(v) => Some((key, v)),
None => path.prev(root, &forest.nodes),
}
})
}
pub fn insert<C: Comparator<K>>(
&mut self,
key: K,
value: V,
forest: &mut MapForest<K, V>,
comp: &C,
) -> Option<V> {
self.cursor(forest, comp).insert(key, value)
}
pub fn remove<C: Comparator<K>>(
&mut self,
key: K,
forest: &mut MapForest<K, V>,
comp: &C,
) -> Option<V> {
let mut c = self.cursor(forest, comp);
if c.goto(key).is_some() {
c.remove()
} else {
None
}
}
pub fn clear(&mut self, forest: &mut MapForest<K, V>) {
if let Some(root) = self.root.take() {
forest.nodes.free_tree(root);
}
}
pub fn retain<F>(&mut self, forest: &mut MapForest<K, V>, mut predicate: F)
where
F: FnMut(K, &mut V) -> bool,
{
let mut path = Path::default();
if let Some(root) = self.root.expand() {
path.first(root, &forest.nodes);
}
while let Some((node, entry)) = path.leaf_pos() {
let keep = {
let (ks, vs) = forest.nodes[node].unwrap_leaf_mut();
predicate(ks[entry], &mut vs[entry])
};
if keep {
path.next(&forest.nodes);
} else {
self.root = path.remove(&mut forest.nodes).into();
}
}
}
pub fn cursor<'a, C: Comparator<K>>(
&'a mut self,
forest: &'a mut MapForest<K, V>,
comp: &'a C,
) -> MapCursor<'a, K, V, C> {
MapCursor::new(self, forest, comp)
}
pub fn iter<'a>(&'a self, forest: &'a MapForest<K, V>) -> MapIter<'a, K, V> {
MapIter {
root: self.root,
pool: &forest.nodes,
path: Path::default(),
}
}
}
impl<K, V> Default for Map<K, V>
where
K: Copy,
V: Copy,
{
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
impl<K, V> Map<K, V>
where
K: Copy + fmt::Display,
V: Copy,
{
fn verify<C: Comparator<K>>(&self, forest: &MapForest<K, V>, comp: &C)
where
NodeData<MapTypes<K, V>>: fmt::Display,
{
if let Some(root) = self.root.expand() {
forest.nodes.verify_tree(root, comp);
}
}
fn tpath<C: Comparator<K>>(&self, key: K, forest: &MapForest<K, V>, comp: &C) -> String {
use alloc::string::ToString;
match self.root.expand() {
None => "map(empty)".to_string(),
Some(root) => {
let mut path = Path::default();
path.find(key, root, &forest.nodes, comp);
path.to_string()
}
}
}
}
pub struct MapCursor<'a, K, V, C>
where
K: 'a + Copy,
V: 'a + Copy,
C: 'a + Comparator<K>,
{
root: &'a mut PackedOption<Node>,
pool: &'a mut NodePool<MapTypes<K, V>>,
comp: &'a C,
path: Path<MapTypes<K, V>>,
}
impl<'a, K, V, C> MapCursor<'a, K, V, C>
where
K: Copy,
V: Copy,
C: Comparator<K>,
{
fn new(container: &'a mut Map<K, V>, forest: &'a mut MapForest<K, V>, comp: &'a C) -> Self {
Self {
root: &mut container.root,
pool: &mut forest.nodes,
comp,
path: Path::default(),
}
}
pub fn is_empty(&self) -> bool {
self.root.is_none()
}
pub fn next(&mut self) -> Option<(K, V)> {
self.path.next(self.pool)
}
pub fn prev(&mut self) -> Option<(K, V)> {
self.root
.expand()
.and_then(|root| self.path.prev(root, self.pool))
}
pub fn key(&self) -> Option<K> {
self.path
.leaf_pos()
.and_then(|(node, entry)| self.pool[node].unwrap_leaf().0.get(entry).cloned())
}
pub fn value(&self) -> Option<V> {
self.path
.leaf_pos()
.and_then(|(node, entry)| self.pool[node].unwrap_leaf().1.get(entry).cloned())
}
pub fn value_mut(&mut self) -> Option<&mut V> {
self.path
.leaf_pos()
.and_then(move |(node, entry)| self.pool[node].unwrap_leaf_mut().1.get_mut(entry))
}
pub fn goto(&mut self, elem: K) -> Option<V> {
self.root.expand().and_then(|root| {
let v = self.path.find(elem, root, self.pool, self.comp);
if v.is_none() {
self.path.normalize(self.pool);
}
v
})
}
pub fn goto_first(&mut self) -> Option<V> {
self.root.map(|root| self.path.first(root, self.pool).1)
}
pub fn insert(&mut self, key: K, value: V) -> Option<V> {
match self.root.expand() {
None => {
let root = self.pool.alloc_node(NodeData::leaf(key, value));
*self.root = root.into();
self.path.set_root_node(root);
None
}
Some(root) => {
let old = self.path.find(key, root, self.pool, self.comp);
if old.is_some() {
*self.path.value_mut(self.pool) = value;
} else {
*self.root = self.path.insert(key, value, self.pool).into();
}
old
}
}
}
pub fn remove(&mut self) -> Option<V> {
let value = self.value();
if value.is_some() {
*self.root = self.path.remove(self.pool).into();
}
value
}
}
pub struct MapIter<'a, K, V>
where
K: 'a + Copy,
V: 'a + Copy,
{
root: PackedOption<Node>,
pool: &'a NodePool<MapTypes<K, V>>,
path: Path<MapTypes<K, V>>,
}
impl<'a, K, V> Iterator for MapIter<'a, K, V>
where
K: 'a + Copy,
V: 'a + Copy,
{
type Item = (K, V);
fn next(&mut self) -> Option<Self::Item> {
match self.root.take() {
Some(root) => Some(self.path.first(root, self.pool)),
None => self.path.next(self.pool),
}
}
}
#[cfg(test)]
impl<'a, K, V, C> MapCursor<'a, K, V, C>
where
K: Copy + fmt::Display,
V: Copy + fmt::Display,
C: Comparator<K>,
{
fn verify(&self) {
self.path.verify(self.pool);
self.root.map(|root| self.pool.verify_tree(root, self.comp));
}
fn tpath(&self) -> String {
use alloc::string::ToString;
self.path.to_string()
}
}
#[cfg(test)]
mod tests {
use super::*;
use alloc::vec::Vec;
use core::mem;
#[test]
fn node_size() {
type F = MapTypes<u32, u32>;
assert_eq!(mem::size_of::<NodeData<F>>(), 64);
}
#[test]
fn empty() {
let mut f = MapForest::<u32, f32>::new();
f.clear();
let mut m = Map::<u32, f32>::new();
assert!(m.is_empty());
m.clear(&mut f);
assert_eq!(m.get(7, &f, &()), None);
assert_eq!(m.iter(&f).next(), None);
assert_eq!(m.get_or_less(7, &f, &()), None);
m.retain(&mut f, |_, _| unreachable!());
let mut c = m.cursor(&mut f, &());
assert!(c.is_empty());
assert_eq!(c.key(), None);
assert_eq!(c.value(), None);
assert_eq!(c.next(), None);
assert_eq!(c.prev(), None);
c.verify();
assert_eq!(c.tpath(), "<empty path>");
assert_eq!(c.goto_first(), None);
assert_eq!(c.tpath(), "<empty path>");
}
#[test]
fn inserting() {
let f = &mut MapForest::<u32, f32>::new();
let mut m = Map::<u32, f32>::new();
assert_eq!(m.insert(50, 5.0, f, &()), None);
assert_eq!(m.insert(50, 5.5, f, &()), Some(5.0));
assert_eq!(m.insert(20, 2.0, f, &()), None);
assert_eq!(m.insert(80, 8.0, f, &()), None);
assert_eq!(m.insert(40, 4.0, f, &()), None);
assert_eq!(m.insert(60, 6.0, f, &()), None);
assert_eq!(m.insert(90, 9.0, f, &()), None);
assert_eq!(m.insert(200, 20.0, f, &()), None);
m.verify(f, &());
assert_eq!(
m.iter(f).collect::<Vec<_>>(),
[
(20, 2.0),
(40, 4.0),
(50, 5.5),
(60, 6.0),
(80, 8.0),
(90, 9.0),
(200, 20.0),
]
);
assert_eq!(m.get(0, f, &()), None);
assert_eq!(m.get(20, f, &()), Some(2.0));
assert_eq!(m.get(30, f, &()), None);
assert_eq!(m.get(40, f, &()), Some(4.0));
assert_eq!(m.get(50, f, &()), Some(5.5));
assert_eq!(m.get(60, f, &()), Some(6.0));
assert_eq!(m.get(70, f, &()), None);
assert_eq!(m.get(80, f, &()), Some(8.0));
assert_eq!(m.get(100, f, &()), None);
assert_eq!(m.get_or_less(0, f, &()), None);
assert_eq!(m.get_or_less(20, f, &()), Some((20, 2.0)));
assert_eq!(m.get_or_less(30, f, &()), Some((20, 2.0)));
assert_eq!(m.get_or_less(40, f, &()), Some((40, 4.0)));
assert_eq!(m.get_or_less(200, f, &()), Some((200, 20.0)));
assert_eq!(m.get_or_less(201, f, &()), Some((200, 20.0)));
{
let mut c = m.cursor(f, &());
assert_eq!(c.prev(), Some((200, 20.0)));
assert_eq!(c.prev(), Some((90, 9.0)));
assert_eq!(c.prev(), Some((80, 8.0)));
assert_eq!(c.prev(), Some((60, 6.0)));
assert_eq!(c.prev(), Some((50, 5.5)));
assert_eq!(c.prev(), Some((40, 4.0)));
assert_eq!(c.prev(), Some((20, 2.0)));
assert_eq!(c.prev(), None);
}
assert_eq!(m.tpath(50, f, &()), "node0[2]");
assert_eq!(m.tpath(80, f, &()), "node0[4]");
assert_eq!(m.tpath(200, f, &()), "node0[6]");
assert_eq!(m.remove(80, f, &()), Some(8.0));
assert_eq!(m.tpath(50, f, &()), "node0[2]");
assert_eq!(m.tpath(80, f, &()), "node0[4]");
assert_eq!(m.tpath(200, f, &()), "node0[5]");
assert_eq!(m.remove(80, f, &()), None);
m.verify(f, &());
assert_eq!(m.remove(20, f, &()), Some(2.0));
assert_eq!(m.tpath(50, f, &()), "node0[1]");
assert_eq!(m.tpath(80, f, &()), "node0[3]");
assert_eq!(m.tpath(200, f, &()), "node0[4]");
assert_eq!(m.remove(20, f, &()), None);
m.verify(f, &());
{
let mut c = m.cursor(f, &());
assert_eq!(c.goto_first(), Some(4.0));
assert_eq!(c.key(), Some(40));
assert_eq!(c.value(), Some(4.0));
assert_eq!(c.next(), Some((50, 5.5)));
assert_eq!(c.next(), Some((60, 6.0)));
assert_eq!(c.next(), Some((90, 9.0)));
assert_eq!(c.next(), Some((200, 20.0)));
c.verify();
assert_eq!(c.next(), None);
c.verify();
}
assert_eq!(m.remove(200, f, &()), Some(20.0));
assert_eq!(m.remove(40, f, &()), Some(4.0));
assert_eq!(m.remove(60, f, &()), Some(6.0));
m.verify(f, &());
assert_eq!(m.remove(50, f, &()), Some(5.5));
m.verify(f, &());
assert_eq!(m.remove(90, f, &()), Some(9.0));
m.verify(f, &());
assert!(m.is_empty());
}
#[test]
fn split_level0_leaf() {
let f = &mut MapForest::<u32, f32>::new();
fn full_leaf(f: &mut MapForest<u32, f32>) -> Map<u32, f32> {
let mut m = Map::new();
for n in 1..8 {
m.insert(n * 10, n as f32 * 1.1, f, &());
}
m
}
let mut m = full_leaf(f);
m.insert(5, 4.2, f, &());
m.verify(f, &());
assert_eq!(m.get(5, f, &()), Some(4.2));
m.retain(f, |k, v| {
*v = (k / 10) as f32;
(k % 20) == 0
});
assert_eq!(
m.iter(f).collect::<Vec<_>>(),
[(20, 2.0), (40, 4.0), (60, 6.0)]
);
let mut m = full_leaf(f);
m.insert(80, 4.2, f, &());
m.verify(f, &());
assert_eq!(m.get(80, f, &()), Some(4.2));
let mut m = full_leaf(f);
m.insert(35, 4.2, f, &());
m.verify(f, &());
assert_eq!(m.get(35, f, &()), Some(4.2));
let mut m = full_leaf(f);
m.insert(45, 4.2, f, &());
m.verify(f, &());
assert_eq!(m.get(45, f, &()), Some(4.2));
m.clear(f);
assert!(m.is_empty());
}
#[test]
fn split_level1_leaf() {
let f = &mut MapForest::<u32, f32>::new();
fn full(f: &mut MapForest<u32, f32>) -> Map<u32, f32> {
let mut m = Map::new();
for row in 1..9 {
for col in 1..5 {
m.insert(row * 100 + col * 10, row as f32 + col as f32 * 0.1, f, &());
}
}
for row in 1..9 {
for col in 5..8 {
m.insert(row * 100 + col * 10, row as f32 + col as f32 * 0.1, f, &());
}
}
m
}
let mut m = full(f);
m.verify(f, &());
assert_eq!(m.tpath(110, f, &()), "node2[0]--node0[0]");
assert_eq!(m.tpath(140, f, &()), "node2[0]--node0[3]");
assert_eq!(m.tpath(210, f, &()), "node2[1]--node1[0]");
assert_eq!(m.tpath(270, f, &()), "node2[1]--node1[6]");
assert_eq!(m.tpath(310, f, &()), "node2[2]--node3[0]");
assert_eq!(m.tpath(810, f, &()), "node2[7]--node8[0]");
assert_eq!(m.tpath(870, f, &()), "node2[7]--node8[6]");
{
let mut c = m.cursor(f, &());
assert_eq!(c.goto_first(), Some(1.1));
assert_eq!(c.key(), Some(110));
}
m.insert(0, 4.2, f, &());
m.verify(f, &());
assert_eq!(m.get(0, f, &()), Some(4.2));
f.clear();
m = full(f);
m.insert(135, 4.2, f, &());
m.verify(f, &());
assert_eq!(m.get(135, f, &()), Some(4.2));
f.clear();
m = full(f);
m.insert(145, 4.2, f, &());
m.verify(f, &());
assert_eq!(m.get(145, f, &()), Some(4.2));
f.clear();
m = full(f);
m.insert(175, 4.2, f, &());
m.verify(f, &());
assert_eq!(m.get(175, f, &()), Some(4.2));
f.clear();
m = full(f);
m.insert(435, 4.2, f, &());
m.verify(f, &());
assert_eq!(m.get(435, f, &()), Some(4.2));
f.clear();
m = full(f);
m.insert(445, 4.2, f, &());
m.verify(f, &());
assert_eq!(m.get(445, f, &()), Some(4.2));
f.clear();
m = full(f);
m.insert(535, 4.2, f, &());
m.verify(f, &());
assert_eq!(m.get(535, f, &()), Some(4.2));
f.clear();
m = full(f);
m.insert(545, 4.2, f, &());
m.verify(f, &());
assert_eq!(m.get(545, f, &()), Some(4.2));
f.clear();
m = full(f);
m.insert(835, 4.2, f, &());
m.verify(f, &());
assert_eq!(m.get(835, f, &()), Some(4.2));
f.clear();
m = full(f);
m.insert(845, 4.2, f, &());
m.verify(f, &());
assert_eq!(m.get(845, f, &()), Some(4.2));
f.clear();
m = full(f);
m.insert(805, 4.2, f, &());
m.verify(f, &());
assert_eq!(m.get(805, f, &()), Some(4.2));
m.clear(f);
m.verify(f, &());
}
fn two_leaf(f: &mut MapForest<u32, f32>) -> Map<u32, f32> {
f.clear();
let mut m = Map::new();
for n in 1..9 {
m.insert(n * 10, n as f32, f, &());
}
m
}
#[test]
fn remove_level1() {
let f = &mut MapForest::<u32, f32>::new();
let mut m = two_leaf(f);
m.verify(f, &());
assert_eq!(m.tpath(10, f, &()), "node2[0]--node0[0]");
assert_eq!(m.tpath(40, f, &()), "node2[0]--node0[3]");
assert_eq!(m.tpath(49, f, &()), "node2[0]--node0[4]");
assert_eq!(m.tpath(50, f, &()), "node2[1]--node1[0]");
assert_eq!(m.tpath(80, f, &()), "node2[1]--node1[3]");
assert_eq!(m.insert(55, 5.5, f, &()), None);
assert_eq!(m.remove(50, f, &()), Some(5.0));
m.verify(f, &());
assert_eq!(m.tpath(49, f, &()), "node2[0]--node0[4]");
assert_eq!(m.tpath(50, f, &()), "node2[0]--node0[4]");
assert_eq!(m.tpath(55, f, &()), "node2[1]--node1[0]");
assert_eq!(m.insert(15, 1.5, f, &()), None);
assert_eq!(m.remove(10, f, &()), Some(1.0));
m.verify(f, &());
assert_eq!(m.remove(55, f, &()), Some(5.5));
m.verify(f, &());
assert_eq!(m.tpath(55, f, &()), "node2[0]--node0[4]");
assert_eq!(m.tpath(60, f, &()), "node2[1]--node1[0]");
assert_eq!(m.insert(90, 9.0, f, &()), None);
assert_eq!(m.insert(100, 10.0, f, &()), None);
m.verify(f, &());
assert_eq!(m.tpath(55, f, &()), "node2[0]--node0[4]");
assert_eq!(m.tpath(60, f, &()), "node2[1]--node1[0]");
assert_eq!(m.remove(20, f, &()), Some(2.0));
m.verify(f, &());
assert_eq!(m.tpath(50, f, &()), "node2[0]--node0[3]");
assert_eq!(m.tpath(60, f, &()), "node2[0]--node0[3]");
assert_eq!(m.tpath(70, f, &()), "node2[1]--node1[0]");
assert_eq!(m.remove(15, f, &()), Some(1.5));
m.verify(f, &());
}
#[test]
fn remove_level1_rightmost() {
let f = &mut MapForest::<u32, f32>::new();
let mut m = two_leaf(f);
assert_eq!(m.remove(60, f, &()), Some(6.0));
assert_eq!(m.remove(80, f, &()), Some(8.0));
assert_eq!(m.remove(50, f, &()), Some(5.0));
m.verify(f, &());
assert_eq!(m.tpath(50, f, &()), "node2[0]--node0[4]");
assert_eq!(m.tpath(70, f, &()), "node2[1]--node1[0]");
assert_eq!(m.remove(70, f, &()), Some(7.0));
m.verify(f, &());
}
fn level3_sparse(f: &mut MapForest<u32, f32>) -> Map<u32, f32> {
f.clear();
let mut m = Map::new();
for n in 1..133 {
m.insert(n * 10, n as f32, f, &());
}
m
}
#[test]
fn level3_removes() {
let f = &mut MapForest::<u32, f32>::new();
let mut m = level3_sparse(f);
m.verify(f, &());
assert_eq!(m.tpath(0, f, &()), "node11[0]--node2[0]--node0[0]");
assert_eq!(m.tpath(10000, f, &()), "node11[7]--node41[4]--node40[4]");
assert_eq!(m.tpath(640, f, &()), "node11[3]--node21[3]--node19[3]");
assert_eq!(m.tpath(650, f, &()), "node11[4]--node26[0]--node20[0]");
assert_eq!(m.remove(640, f, &()), Some(64.0));
m.verify(f, &());
assert_eq!(m.tpath(650, f, &()), "node11[3]--node26[3]--node20[3]");
assert_eq!(m.tpath(1130, f, &()), "node11[6]--node41[0]--node35[0]");
assert_eq!(m.tpath(1140, f, &()), "node11[6]--node41[0]--node35[1]");
assert_eq!(m.remove(1130, f, &()), Some(113.0));
m.verify(f, &());
assert_eq!(m.tpath(1140, f, &()), "node11[6]--node41[0]--node37[0]");
}
#[test]
fn insert_many() {
let f = &mut MapForest::<u32, f32>::new();
let mut m = Map::<u32, f32>::new();
let mm = 4096;
let mut x = 0;
for n in 0..mm {
assert_eq!(m.insert(x, n as f32, f, &()), None);
m.verify(f, &());
x = (x + n + 1) % mm;
}
x = 0;
for n in 0..mm {
assert_eq!(m.get(x, f, &()), Some(n as f32));
x = (x + n + 1) % mm;
}
x = 0;
for n in 0..mm {
assert_eq!(m.remove(x, f, &()), Some(n as f32));
m.verify(f, &());
x = (x + n + 1) % mm;
}
assert!(m.is_empty());
}
}