use core::fmt::Debug;
use std::cmp::Ordering;
use std::ops::Range;
use crate::rle::{CanRemove, HasLength, Mergeable, Sliceable, TryInsert};
use crate::{BTree, BTreeTrait, FindResult, Query, SplitInfo};
#[derive(Debug)]
#[repr(transparent)]
struct OrdTrait<Key, Value> {
_phantom: core::marker::PhantomData<(Key, Value)>,
}
#[derive(Debug)]
pub struct OrdTreeMap<Key: Clone + Ord + Debug + 'static, Value: Clone + Debug> {
tree: BTree<OrdTrait<Key, Value>>,
len: usize,
}
#[derive(Debug)]
pub struct OrdTreeSet<Key: Clone + Ord + Debug + 'static>(OrdTreeMap<Key, ()>);
impl<Key: Clone + Ord + Debug + 'static, Value: Clone + Debug + 'static> OrdTreeMap<Key, Value> {
#[inline(always)]
pub fn new() -> Self {
Self {
tree: BTree::new(),
len: 0,
}
}
#[inline(always)]
pub fn insert(&mut self, key: Key, value: Value) {
let Some(result) = self.tree.query::<OrdTrait<Key, Value>>(&key) else {
self.len += 1;
self.tree.push(Unmergeable((key, value)));
return;
};
if !result.found {
self.len += 1;
let tree = &mut self.tree;
let data = Unmergeable((key, value));
let index = result.leaf();
let leaf = tree.leaf_nodes.get_mut(index.0).unwrap();
let parent = leaf.parent();
let mut is_full = false;
if result.cursor.offset == 0 && data.can_merge(&leaf.elem) {
leaf.elem.merge_left(&data);
} else if result.cursor.offset == leaf.elem.rle_len() && leaf.elem.can_merge(&data) {
leaf.elem.merge_right(&data);
} else {
let child = tree.alloc_leaf_child(data, parent.unwrap_internal());
let SplitInfo {
parent_idx: parent_index,
insert_slot: insert_index,
..
} = tree.split_leaf_if_needed(result.cursor);
let parent = tree.in_nodes.get_mut(parent_index).unwrap();
parent.children.insert(insert_index, child).unwrap();
is_full = parent.is_full();
}
tree.recursive_update_cache(parent, false, None);
if is_full {
tree.split(parent);
}
} else {
let leaf = self.tree.get_elem_mut(result.leaf()).unwrap();
leaf.0 .1 = value;
}
}
#[inline(always)]
pub fn delete(&mut self, key: &Key) -> Option<(Key, Value)> {
let q = self.tree.query::<OrdTrait<Key, Value>>(key)?;
match self.tree.remove_leaf(q.cursor) {
Some(v) => {
self.len -= 1;
Some(v.0)
}
None => None,
}
}
#[inline(always)]
pub fn iter(&self) -> impl Iterator<Item = &(Key, Value)> {
self.tree.iter().map(|x| &x.0)
}
#[inline(always)]
pub fn iter_key(&self) -> impl Iterator<Item = &Key> {
self.tree.iter().map(|x| &x.0 .0)
}
#[inline(always)]
pub fn len(&self) -> usize {
self.len
}
#[inline(always)]
pub fn is_empty(&self) -> bool {
self.len == 0
}
#[allow(unused)]
pub(crate) fn check(&self) {
self.tree.check()
}
}
impl<Key: Clone + Ord + Debug + 'static> OrdTreeSet<Key> {
#[inline(always)]
pub fn new() -> Self {
Self(OrdTreeMap::new())
}
#[inline(always)]
pub fn insert(&mut self, key: Key) {
self.0.insert(key, ());
}
#[inline(always)]
pub fn delete(&mut self, key: &Key) -> bool {
self.0.delete(key).is_some()
}
#[inline(always)]
pub fn iter(&self) -> impl Iterator<Item = &Key> {
self.0.iter_key()
}
pub fn len(&self) -> usize {
self.0.len
}
pub fn is_empty(&self) -> bool {
self.0.len == 0
}
#[allow(unused)]
fn check(&self) {
self.0.check()
}
}
impl<Key: Clone + Ord + Debug + 'static> Default for OrdTreeSet<Key> {
#[inline(always)]
fn default() -> Self {
Self::new()
}
}
impl<Key: Clone + Ord + Debug + 'static, Value: Clone + Debug + 'static> Default
for OrdTreeMap<Key, Value>
{
#[inline(always)]
fn default() -> Self {
Self::new()
}
}
impl<Key, Value> Default for OrdTrait<Key, Value> {
#[inline(always)]
fn default() -> Self {
Self {
_phantom: Default::default(),
}
}
}
#[repr(transparent)]
#[derive(Debug, Clone)]
pub struct Unmergeable<T>(T);
impl<T> HasLength for Unmergeable<T> {
fn rle_len(&self) -> usize {
1
}
}
impl<T: Clone> Sliceable for Unmergeable<T> {
fn _slice(&self, range: Range<usize>) -> Self {
if range.end - range.start != 1 {
panic!("Invalid range");
}
self.clone()
}
}
impl<T> Mergeable for Unmergeable<T> {
fn can_merge(&self, _rhs: &Self) -> bool {
false
}
fn merge_right(&mut self, _rhs: &Self) {
unreachable!()
}
fn merge_left(&mut self, _left: &Self) {
unreachable!()
}
}
impl<T> TryInsert for Unmergeable<T> {
fn try_insert(&mut self, _pos: usize, elem: Self) -> Result<(), Self> {
Err(elem)
}
}
impl<T> CanRemove for Unmergeable<T> {
fn can_remove(&self) -> bool {
false
}
}
impl<Key: Clone + Ord + Debug + 'static, Value: Clone + Debug> BTreeTrait for OrdTrait<Key, Value> {
type Elem = Unmergeable<(Key, Value)>;
type Cache = Option<(Key, Key)>;
type CacheDiff = ();
const USE_DIFF: bool = false;
#[inline(always)]
fn calc_cache_internal(cache: &mut Self::Cache, caches: &[crate::Child<Self>]) {
if caches.is_empty() {
return;
}
*cache = Some((
caches[0].cache.as_ref().unwrap().0.clone(),
caches[caches.len() - 1].cache.as_ref().unwrap().1.clone(),
));
}
#[inline(always)]
fn apply_cache_diff(_: &mut Self::Cache, _: &Self::CacheDiff) {
unreachable!()
}
#[inline(always)]
fn merge_cache_diff(_: &mut Self::CacheDiff, _: &Self::CacheDiff) {}
#[inline(always)]
fn get_elem_cache(elem: &Self::Elem) -> Self::Cache {
Some((elem.0 .0.clone(), elem.0 .0.clone()))
}
#[inline(always)]
fn new_cache_to_diff(_: &Self::Cache) -> Self::CacheDiff {}
fn sub_cache(_: &Self::Cache, _: &Self::Cache) -> Self::CacheDiff {}
}
impl<Key: Ord + Clone + Debug + 'static, Value: Clone + Debug + 'static> Query<OrdTrait<Key, Value>>
for OrdTrait<Key, Value>
{
type QueryArg = Key;
#[inline(always)]
fn init(_target: &Self::QueryArg) -> Self {
Self::default()
}
#[inline]
fn find_node(
&mut self,
target: &Self::QueryArg,
child_caches: &[crate::Child<OrdTrait<Key, Value>>],
) -> crate::FindResult {
let result = child_caches.binary_search_by(|x| {
let (min, max) = x.cache.as_ref().unwrap();
if target < min {
core::cmp::Ordering::Greater
} else if target > max {
core::cmp::Ordering::Less
} else {
core::cmp::Ordering::Equal
}
});
match result {
Ok(i) => FindResult::new_found(i, 0),
Err(i) => FindResult::new_missing(
i.min(child_caches.len() - 1),
if i == child_caches.len() { 1 } else { 0 },
),
}
}
#[inline(always)]
fn confirm_elem(
&mut self,
q: &Self::QueryArg,
elem: &<OrdTrait<Key, Value> as BTreeTrait>::Elem,
) -> (usize, bool) {
match q.cmp(&elem.0 .0) {
Ordering::Less => (0, false),
Ordering::Equal => (0, true),
Ordering::Greater => (1, false),
}
}
}
#[cfg(test)]
mod test {
use std::cmp::Ordering;
use rand::{Rng, SeedableRng};
use crate::HeapVec;
use super::*;
#[test]
fn test() {
let mut tree: OrdTreeSet<u64> = OrdTreeSet::new();
let mut rng = rand::rngs::StdRng::seed_from_u64(123);
let mut data: HeapVec<u64> = (0..1000).map(|_| rng.gen()).collect();
for &value in data.iter() {
tree.insert(value);
}
data.sort_unstable();
assert_eq!(tree.iter().copied().collect::<HeapVec<_>>(), data);
tree.check();
}
#[test]
fn test_delete() {
let mut tree: OrdTreeSet<u64> = OrdTreeSet::new();
tree.insert(12);
tree.delete(&12);
assert_eq!(tree.len(), 0);
}
#[test]
fn test_compare_pos() {
let mut tree: OrdTreeSet<u64> = OrdTreeSet::new();
for i in 0..100 {
tree.insert(i);
}
for i in 0..99 {
let a = tree.0.tree.query::<OrdTrait<u64, ()>>(&i).unwrap();
assert_eq!(
tree.0.tree.compare_pos(a.cursor(), a.cursor()),
Ordering::Equal
);
for j in i + 1..100 {
let b = tree.0.tree.query::<OrdTrait<u64, ()>>(&j).unwrap();
assert_eq!(
tree.0.tree.compare_pos(a.cursor(), b.cursor()),
Ordering::Less
);
assert_eq!(
tree.0.tree.compare_pos(b.cursor(), a.cursor()),
Ordering::Greater
);
}
}
}
mod move_event_test {
use super::*;
#[test]
fn test() {
let mut tree: OrdTreeMap<u64, usize> = OrdTreeMap::new();
let mut rng = rand::rngs::StdRng::seed_from_u64(123);
let mut data: HeapVec<u64> = (0..1000).map(|_| rng.gen()).collect();
for &value in data.iter() {
tree.insert(value, 0);
}
for value in data.drain(0..100) {
tree.delete(&value);
}
for value in data.drain(0..800) {
tree.delete(&value);
}
tree.tree.check();
for _ in (0..100).rev() {
tree.delete(&data.pop().unwrap());
}
}
}
}