use std::ops::{Index, IndexMut};
use std::fmt::{Debug, Formatter};
use thin_vec::ThinVec;
use bitvec::prelude::*;
pub type Key = BitVec<usize, Lsb0>;
pub type KeyRef<'a> = &'a BitSlice<usize, Lsb0>;
#[derive(Default)]
struct Bitmap {
bitmap: BitmapType,
}
type BitmapType = u64;
const RESULTS_BITS_END_NODE: usize = 5;
const RESULTS_BITS: usize = RESULTS_BITS_END_NODE - 1;
const CHILDREN_START_END_NODE: usize = 2_usize.pow(RESULTS_BITS_END_NODE as u32 + 1);
const CHILDREN_START: usize = 2_usize.pow(RESULTS_BITS as u32 + 1);
impl Debug for Bitmap {
fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
let field2_name = if self.is_end_node() { "internal2" } else { "external" };
f.debug_struct("Bitmap")
.field("is_end_node", &self.is_end_node())
.field("internal", &format!("{}", &self.bitmap.view_bits::<Lsb0>()[..CHILDREN_START]))
.field(field2_name, &format!("{}", &self.bitmap.view_bits::<Lsb0>()[CHILDREN_START..]))
.finish()
}
}
impl Bitmap {
#[inline]
fn is_end_node(&self) -> bool {
self.bitmap.view_bits::<Lsb0>()[0]
}
fn set_is_end_node(&mut self, is_end_node: bool) {
self.bitmap.view_bits_mut::<Lsb0>().set(0, is_end_node);
}
#[inline]
fn children_start_at(&self) -> usize {
if self.is_end_node() { CHILDREN_START_END_NODE } else { CHILDREN_START }
}
#[inline]
fn results_capacity(&self) -> usize {
if self.is_end_node() { RESULTS_BITS_END_NODE } else { RESULTS_BITS }
}
fn children_bits_mut(&mut self) -> &mut BitSlice<BitmapType, Lsb0> {
let start = self.children_start_at();
self.bitmap.view_bits_mut::<Lsb0>().index_mut(start..)
}
fn children_bits(&self) -> &BitSlice<BitmapType, Lsb0> {
let start = self.children_start_at();
self.bitmap.view_bits::<Lsb0>().index(start..)
}
fn results_bits_mut(&mut self) -> &mut BitSlice<BitmapType, Lsb0> {
let end = self.children_start_at();
self.bitmap.view_bits_mut::<Lsb0>().index_mut(1..end)
}
fn results_bits(&self) -> &BitSlice<BitmapType, Lsb0> {
let end = self.children_start_at();
self.bitmap.view_bits::<Lsb0>().index(1..end)
}
fn results_keys_with_prefix(&self, prefix: Key) -> impl Iterator<Item = Key> + '_ {
self.results_bits()
.iter_ones()
.map(from_index)
.map(move |result_key| {
let mut key = prefix.clone();
key.extend(result_key);
key
})
}
}
#[derive(Debug)]
pub struct Node<T> {
results: Option<ThinVec<T>>,
children: Option<ThinVec<Node<T>>>,
bitmap: Bitmap,
}
impl<T> Default for Node<T> {
fn default() -> Self {
let mut bitmap: Bitmap = Default::default();
bitmap.set_is_end_node(true);
Node {
results: None,
children: None,
bitmap,
}
}
}
fn children_mut<'a, T>(bitmap: &'a Bitmap, children: &'a mut Option<ThinVec<Node<T>>>) -> impl Iterator<Item = (Key, &'a mut Node<T>)> {
let children_iter = children.iter_mut().flat_map(|children| children.iter_mut());
bitmap.children_bits().iter_ones().map(|x| x.view_bits::<Lsb0>().iter().take(RESULTS_BITS_END_NODE).collect()).zip(children_iter)
}
fn results_mut<'a, T>(bitmap: &'a Bitmap, results: &'a mut Option<ThinVec<T>>) -> impl Iterator<Item = (Key, &'a mut T)> {
let results_iter = results.iter_mut().flat_map(|results| results.iter_mut());
bitmap.results_bits().iter_ones().map(from_index).zip(results_iter)
}
fn to_index(key: KeyRef) -> usize {
let leading_one = 2usize.pow(key.len() as u32);
let net_bits: usize = if key.is_empty() { 0 } else { key.load_le() };
(leading_one + net_bits) - 1
}
fn from_index(mut index: usize) -> Key {
index += 1;
let prefix_len = (std::mem::size_of::<usize>() as u32 * 8) - index.leading_zeros() - 1;
let mut key = Key::new();
key.extend(index.view_bits::<Lsb0>().iter().take(prefix_len as usize));
key
}
impl<T: Debug + Send + Sync> Node<T> {
fn children(&self) -> impl Iterator<Item = (Key, &Node<T>)> {
let children_iter = self.children.iter().flat_map(|children| children.iter());
self.bitmap.children_bits().iter_ones().map(|x| x.view_bits::<Lsb0>().iter().take(RESULTS_BITS_END_NODE).collect()).zip(children_iter)
}
fn results(&self) -> impl Iterator<Item = (Key, &T)> {
let results_iter = self.results.iter().flat_map(|results| results.iter());
self.bitmap.results_bits().iter_ones().map(from_index).zip(results_iter)
}
fn get_child(&self, key: KeyRef) -> Option<&Node<T>> {
if self.bitmap.is_end_node() { return None; }
let nibble: usize = key.load_le();
self.bitmap.children_bits()[nibble].then(|| {
let vec_index = self.bitmap.children_bits()[..nibble].count_ones();
&self.children.as_ref().unwrap()[vec_index]
})
}
fn get_child_mut(&mut self, key: KeyRef) -> Option<&mut Node<T>> {
if self.bitmap.is_end_node() { return None; }
let nibble: usize = key.load_le();
self.bitmap.children_bits()[nibble].then(|| {
let vec_index = self.bitmap.children_bits()[..nibble].count_ones();
&mut self.children.as_mut().unwrap()[vec_index]
})
}
fn convert_to_normal(&mut self) {
if !self.bitmap.is_end_node() { return; }
let results_iter = self.results.take().into_iter().flat_map(|results| results.into_iter());
let results = self.bitmap.results_bits().iter_ones().map(from_index).zip(results_iter).collect::<Vec<_>>();
self.bitmap = Default::default();
self.bitmap.set_is_end_node(false);
for (key, value) in results {
self.insert(&key, value);
}
}
fn get_or_insert_child(&mut self, key: KeyRef) -> &mut Node<T> {
self.convert_to_normal();
{
let nibble: usize = key.load_le();
if !self.bitmap.children_bits()[nibble] {
self.bitmap.children_bits_mut().set(nibble, true);
let children = self.children.get_or_insert(Default::default());
let vec_index = self.bitmap.children_bits()[..nibble].count_ones();
children.insert(vec_index, Node::default());
}
}
self.get_child_mut(key).unwrap()
}
pub fn insert(&mut self, key: KeyRef, value: T) -> Option<T> {
if key.len() <= self.bitmap.results_capacity() {
let index = to_index(&key);
let results = self.results.get_or_insert(Default::default());
let vec_index = self.bitmap.results_bits()[..index].count_ones();
if self.bitmap.results_bits()[index] {
Some(std::mem::replace(&mut results[vec_index], value))
} else {
self.bitmap.results_bits_mut().set(index, true);
results.insert(vec_index, value);
None
}
} else {
let (key, remaining) = key.split_at(RESULTS_BITS_END_NODE);
let child = self.get_or_insert_child(key);
child.insert(remaining, value)
}
}
pub fn remove(&mut self, key: KeyRef) -> Option<T> {
if key.len() <= self.bitmap.results_capacity() {
let index = to_index(key);
self.bitmap.results_bits()[index].then(|| {
self.bitmap.results_bits_mut().set(index, false);
let results = self.results.get_or_insert(Default::default());
let vec_index = self.bitmap.results_bits()[..index].count_ones();
results.remove(vec_index)
})
} else {
let (key, remaining) = key.split_at(RESULTS_BITS_END_NODE);
self.get_child_mut(key).and_then(|child| child.remove(remaining))
}
}
fn iter_with_prefix(&self, prefix: Key) -> impl Iterator<Item = (Key, &T)> + Send + Sync + '_ {
let results_iter = {
let prefix = prefix.clone();
self.results().map(move |(result_key, val)| {
let mut key = prefix.clone();
key.extend(result_key);
(key, val)
})
};
let children_iter = self.children()
.flat_map(move |(child_key, child)| {
let mut key = prefix.clone();
key.extend(child_key);
child.iter_with_prefix(key)
});
let children_iter: Box<dyn Iterator<Item = (Key, &T)> + Send + Sync + '_> = Box::new(children_iter);
results_iter.chain(children_iter)
}
fn iter_mut_with_prefix(&mut self, prefix: Key) -> impl Iterator<Item = (Key, &mut T)> + Send + Sync + '_ {
let results_iter = {
let prefix = prefix.clone();
results_mut(&self.bitmap, &mut self.results).map(move |(result_key, val)| {
let mut key = prefix.clone();
key.extend(result_key);
(key, val)
})
};
let children_iter = children_mut(&self.bitmap, &mut self.children)
.flat_map(move |(child_key, child)| {
let mut key = prefix.clone();
key.extend(child_key);
child.iter_mut_with_prefix(key)
});
let children_iter: Box<dyn Iterator<Item = (Key, &mut T)> + Send + Sync + '_> = Box::new(children_iter);
results_iter.chain(children_iter)
}
pub fn iter(&self) -> impl Iterator<Item = (Key, &T)> + '_ {
self.iter_with_prefix(Key::new())
}
pub fn iter_mut(&mut self) -> impl Iterator<Item = (Key, &mut T)> + '_ {
self.iter_mut_with_prefix(Key::new())
}
pub fn values(&self) -> impl Iterator<Item = &T> + Send + Sync + '_ {
let results_iter = self.results.iter().flat_map(|values| values.iter());
let children_iter = self.children.iter()
.flat_map(|children| children.iter())
.flat_map(|child| child.values());
let children_iter: Box<dyn Iterator<Item = &T> + Send + Sync + '_> = Box::new(children_iter);
results_iter.chain(children_iter)
}
pub fn values_mut(&mut self) -> impl Iterator<Item = &mut T> + Send + Sync + '_ {
let results_iter = self.results.iter_mut().flat_map(|values| values.iter_mut());
let children_iter = self.children.iter_mut()
.flat_map(|children| children.iter_mut())
.flat_map(|child| child.values_mut());
let children_iter: Box<dyn Iterator<Item = &mut T> + Send + Sync + '_> = Box::new(children_iter);
results_iter.chain(children_iter)
}
fn keys_with_prefix<'a>(&'a self, prefix: Key) -> impl Iterator<Item = Key> + Send + Sync + '_ {
let results_keys_iter = self.bitmap.results_keys_with_prefix(prefix.clone());
let children_keys_iter = self.children()
.flat_map(move |(child_key, child)| {
let mut key = prefix.clone();
key.extend(child_key);
child.keys_with_prefix(key)
});
let children_keys_iter: Box<dyn Iterator<Item = Key> + Send + Sync + '_> = Box::new(children_keys_iter);
results_keys_iter.chain(children_keys_iter)
}
pub fn keys(&self) -> impl Iterator<Item = Key> + '_ {
self.keys_with_prefix(Key::new())
}
pub fn exact(&self, key: KeyRef) -> Option<&T> {
if key.len() <= self.bitmap.results_capacity() {
let index = to_index(key);
self.bitmap.results_bits()[index].then(|| {
let vec_index = self.bitmap.results_bits()[..index].count_ones();
&self.results.as_ref().unwrap()[vec_index]
})
} else {
let (key, remaining) = key.split_at(RESULTS_BITS_END_NODE);
self.get_child(key).and_then(|child| child.exact(remaining))
}
}
pub fn exact_mut(&mut self, key: KeyRef) -> Option<&mut T> {
if key.len() <= self.bitmap.results_capacity() {
let index = to_index(key);
self.bitmap.results_bits()[index].then(|| {
let vec_index = self.bitmap.results_bits()[..index].count_ones();
&mut self.results.as_mut().unwrap()[vec_index]
})
} else {
let (key, remaining) = key.split_at(RESULTS_BITS_END_NODE);
self.get_child_mut(key).and_then(|child| child.exact_mut(remaining))
}
}
fn longest_match_with_prefix(&self, mut prefix: Key, mut key: KeyRef) -> Option<(Key, &T)> {
(key.len() > self.bitmap.results_capacity()).then(|| {
let mut prefix = prefix.clone();
let (key, remaining) = key.split_at(RESULTS_BITS_END_NODE);
prefix.extend(key);
self.get_child(key).and_then(|child| child.longest_match_with_prefix(prefix, remaining))
})
.flatten()
.or_else(|| {
loop {
if let Some(result) = self.exact(key) {
prefix.extend(key);
return Some((prefix, result));
}
if !key.is_empty() {
key = &key[..key.len() - 1];
} else {
break;
}
}
None
})
}
pub fn longest_match(&self, key: KeyRef) -> Option<(Key, &T)> {
self.longest_match_with_prefix(Key::new(), key)
}
fn or_longer_with_prefix(&self, prefix: Key, mut key: Key) -> Box<dyn Iterator<Item = (Key, &T)> + Send + Sync + '_> {
if key.len() > self.bitmap.results_capacity() {
let mut prefix = prefix.clone();
let remaining = key.split_off(RESULTS_BITS_END_NODE);
prefix.extend(&key);
if let Some(child) = self.get_child(&key) {
child.or_longer_with_prefix(prefix, remaining)
} else {
Box::new(std::iter::empty())
}
} else {
let results_iter = {
let prefix = prefix.clone();
let key = key.clone();
self.results()
.filter(move |(result_key, _)| {
result_key.starts_with(&key)
})
.map(move |(result_key, val)| {
let mut key = prefix.clone();
key.extend(result_key);
(key, val)
})
};
let children_iter = self.children()
.filter(move |(child_key, _)| {
child_key.starts_with(&key)
})
.flat_map(move |(child_key, child)| {
let mut key = prefix.clone();
key.extend(child_key);
child.iter_with_prefix(key)
});
let children_iter: Box<dyn Iterator<Item = (Key, &T)> + Send + Sync + '_> = Box::new(children_iter);
Box::new(results_iter.chain(children_iter))
}
}
pub fn or_longer(&self, key: Key) -> impl Iterator<Item = (Key, &T)> + '_ {
self.or_longer_with_prefix(Key::new(), key)
}
fn matches_with_prefix(&self, prefix: Key, mut key: Key) -> impl Iterator<Item = (Key, &T)> + Send + Sync + '_ {
let results_iter = {
let prefix = prefix.clone();
let key = key.clone();
self.results()
.filter(move |(result_key, _)| {
key.starts_with(result_key)
})
.map(move |(result_key, val)| {
let mut key = prefix.clone();
key.extend(result_key);
(key, val)
})
};
let children_iter = {
let key = key.clone();
self.children()
.filter(move |(child_key, _)| {
key.starts_with(child_key)
})
};
let children_iter = children_iter
.flat_map(move |(child_key, child)| {
let remaining = key.split_off(RESULTS_BITS_END_NODE);
let mut key = prefix.clone();
key.extend(child_key);
child.matches_with_prefix(key, remaining)
});
let children_iter: Box<dyn Iterator<Item = (Key, &T)> + Send + Sync + '_> = Box::new(children_iter);
results_iter.chain(children_iter)
}
pub fn matches(&self, key: Key) -> impl Iterator<Item = (Key, &T)> + '_ {
self.matches_with_prefix(Key::new(), key)
}
}