use core::{
borrow::Borrow,
cmp::Ordering,
intrinsics,
mem,
ops::{
Bound::{Excluded, Included, Unbounded},
RangeBounds,
},
ptr,
};
use super::{
node::{marker, ForceResult::*, Handle, NodeRef},
search::{self, SearchResult},
unwrap_unchecked,
};
use crate::alloc::AllocRef;
fn range_search<BorrowType, K, V, Q, R>(
root1: NodeRef<BorrowType, K, V, marker::LeafOrInternal>,
root2: NodeRef<BorrowType, K, V, marker::LeafOrInternal>,
range: R,
) -> (
Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge>,
Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge>,
)
where
Q: ?Sized + Ord,
K: Borrow<Q>,
R: RangeBounds<Q>,
{
match (range.start_bound(), range.end_bound()) {
(Excluded(s), Excluded(e)) if s == e => {
panic!("range start and end are equal and excluded in BTreeMap")
}
(Included(s) | Excluded(s), Included(e) | Excluded(e)) if s > e => {
panic!("range start is greater than range end in BTreeMap")
}
_ => {}
};
let mut min_node = root1;
let mut max_node = root2;
let mut min_found = false;
let mut max_found = false;
loop {
let front = match (min_found, range.start_bound()) {
(false, Included(key)) => match search::search_node(min_node, key) {
SearchResult::Found(kv) => {
min_found = true;
kv.left_edge()
}
SearchResult::GoDown(edge) => edge,
},
(false, Excluded(key)) => match search::search_node(min_node, key) {
SearchResult::Found(kv) => {
min_found = true;
kv.right_edge()
}
SearchResult::GoDown(edge) => edge,
},
(true, Included(_)) => min_node.last_edge(),
(true, Excluded(_)) => min_node.first_edge(),
(_, Unbounded) => min_node.first_edge(),
};
let back = match (max_found, range.end_bound()) {
(false, Included(key)) => match search::search_node(max_node, key) {
SearchResult::Found(kv) => {
max_found = true;
kv.right_edge()
}
SearchResult::GoDown(edge) => edge,
},
(false, Excluded(key)) => match search::search_node(max_node, key) {
SearchResult::Found(kv) => {
max_found = true;
kv.left_edge()
}
SearchResult::GoDown(edge) => edge,
},
(true, Included(_)) => max_node.first_edge(),
(true, Excluded(_)) => max_node.last_edge(),
(_, Unbounded) => max_node.last_edge(),
};
if front.partial_cmp(&back) == Some(Ordering::Greater) {
panic!("Ord is ill-defined in BTreeMap range");
}
match (front.force(), back.force()) {
(Leaf(f), Leaf(b)) => {
return (f, b);
}
(Internal(min_int), Internal(max_int)) => {
min_node = min_int.descend();
max_node = max_int.descend();
}
_ => unreachable!("BTreeMap has different depths"),
};
}
}
fn full_range<BorrowType, K, V>(
root1: NodeRef<BorrowType, K, V, marker::LeafOrInternal>,
root2: NodeRef<BorrowType, K, V, marker::LeafOrInternal>,
) -> (
Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge>,
Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge>,
) {
let mut min_node = root1;
let mut max_node = root2;
loop {
let front = min_node.first_edge();
let back = max_node.last_edge();
match (front.force(), back.force()) {
(Leaf(f), Leaf(b)) => {
return (f, b);
}
(Internal(min_int), Internal(max_int)) => {
min_node = min_int.descend();
max_node = max_int.descend();
}
_ => unreachable!("BTreeMap has different depths"),
};
}
}
impl<'a, K: 'a, V: 'a> NodeRef<marker::Immut<'a>, K, V, marker::LeafOrInternal> {
pub fn range_search<Q, R>(
self,
range: R,
) -> (
Handle<NodeRef<marker::Immut<'a>, K, V, marker::Leaf>, marker::Edge>,
Handle<NodeRef<marker::Immut<'a>, K, V, marker::Leaf>, marker::Edge>,
)
where
Q: ?Sized + Ord,
K: Borrow<Q>,
R: RangeBounds<Q>,
{
range_search(self, self, range)
}
pub fn full_range(
self,
) -> (
Handle<NodeRef<marker::Immut<'a>, K, V, marker::Leaf>, marker::Edge>,
Handle<NodeRef<marker::Immut<'a>, K, V, marker::Leaf>, marker::Edge>,
) {
full_range(self, self)
}
}
impl<'a, K: 'a, V: 'a> NodeRef<marker::ValMut<'a>, K, V, marker::LeafOrInternal> {
pub fn range_search<Q, R>(
self,
range: R,
) -> (
Handle<NodeRef<marker::ValMut<'a>, K, V, marker::Leaf>, marker::Edge>,
Handle<NodeRef<marker::ValMut<'a>, K, V, marker::Leaf>, marker::Edge>,
)
where
Q: ?Sized + Ord,
K: Borrow<Q>,
R: RangeBounds<Q>,
{
let self2 = unsafe { ptr::read(&self) };
range_search(self, self2, range)
}
pub fn full_range(
self,
) -> (
Handle<NodeRef<marker::ValMut<'a>, K, V, marker::Leaf>, marker::Edge>,
Handle<NodeRef<marker::ValMut<'a>, K, V, marker::Leaf>, marker::Edge>,
) {
let self2 = unsafe { ptr::read(&self) };
full_range(self, self2)
}
}
impl<K, V> NodeRef<marker::Owned, K, V, marker::LeafOrInternal> {
pub fn full_range(
self,
) -> (
Handle<NodeRef<marker::Owned, K, V, marker::Leaf>, marker::Edge>,
Handle<NodeRef<marker::Owned, K, V, marker::Leaf>, marker::Edge>,
) {
let self2 = unsafe { ptr::read(&self) };
full_range(self, self2)
}
}
impl<BorrowType, K, V> Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge> {
pub fn next_kv(
self,
) -> Result<
Handle<NodeRef<BorrowType, K, V, marker::LeafOrInternal>, marker::KV>,
NodeRef<BorrowType, K, V, marker::LeafOrInternal>,
> {
let mut edge = self.forget_node_type();
loop {
edge = match edge.right_kv() {
Ok(kv) => return Ok(kv),
Err(last_edge) => match last_edge.into_node().ascend() {
Ok(parent_edge) => parent_edge.forget_node_type(),
Err(root) => return Err(root),
},
}
}
}
pub fn next_back_kv(
self,
) -> Result<
Handle<NodeRef<BorrowType, K, V, marker::LeafOrInternal>, marker::KV>,
NodeRef<BorrowType, K, V, marker::LeafOrInternal>,
> {
let mut edge = self.forget_node_type();
loop {
edge = match edge.left_kv() {
Ok(kv) => return Ok(kv),
Err(last_edge) => match last_edge.into_node().ascend() {
Ok(parent_edge) => parent_edge.forget_node_type(),
Err(root) => return Err(root),
},
}
}
}
}
impl<BorrowType, K, V> Handle<NodeRef<BorrowType, K, V, marker::Internal>, marker::Edge> {
pub fn next_kv(
self,
) -> Result<
Handle<NodeRef<BorrowType, K, V, marker::Internal>, marker::KV>,
NodeRef<BorrowType, K, V, marker::Internal>,
> {
let mut edge = self;
loop {
edge = match edge.right_kv() {
Ok(internal_kv) => return Ok(internal_kv),
Err(last_edge) => match last_edge.into_node().ascend() {
Ok(parent_edge) => parent_edge,
Err(root) => return Err(root),
},
}
}
}
}
macro_rules! def_next_kv_uncheched_dealloc {
{ unsafe fn $name:ident : $adjacent_kv:ident } => {
unsafe fn $name <K, V, A: AllocRef>(
leaf_edge: Handle<NodeRef<marker::Owned, K, V, marker::Leaf>, marker::Edge>,
mut alloc: A
) -> Handle<NodeRef<marker::Owned, K, V, marker::LeafOrInternal>, marker::KV> {
let mut edge = leaf_edge.forget_node_type();
loop {
edge = match edge.$adjacent_kv() {
Ok(internal_kv) => return internal_kv,
Err(last_edge) => {
unsafe {
let parent_edge = last_edge.into_node().deallocate_and_ascend(&mut alloc);
unwrap_unchecked(parent_edge).forget_node_type()
}
}
}
}
}
};
}
def_next_kv_uncheched_dealloc! {unsafe fn next_kv_unchecked_dealloc: right_kv}
def_next_kv_uncheched_dealloc! {unsafe fn next_back_kv_unchecked_dealloc: left_kv}
#[inline]
fn take_mut<T>(v: &mut T, change: impl FnOnce(T) -> T) {
replace(v, |value| (change(value), ()))
}
#[inline]
fn replace<T, R>(v: &mut T, change: impl FnOnce(T) -> (T, R)) -> R {
struct PanicGuard;
impl Drop for PanicGuard {
fn drop(&mut self) {
intrinsics::abort()
}
}
let guard = PanicGuard;
let value = unsafe { ptr::read(v) };
let (new_value, ret) = change(value);
unsafe {
ptr::write(v, new_value);
}
mem::forget(guard);
ret
}
impl<'a, K, V> Handle<NodeRef<marker::Immut<'a>, K, V, marker::Leaf>, marker::Edge> {
pub unsafe fn next_unchecked(&mut self) -> (&'a K, &'a V) {
replace(self, |leaf_edge| {
let kv = leaf_edge.next_kv();
let kv = unsafe { unwrap_unchecked(kv.ok()) };
(kv.next_leaf_edge(), kv.into_kv())
})
}
pub unsafe fn next_back_unchecked(&mut self) -> (&'a K, &'a V) {
replace(self, |leaf_edge| {
let kv = leaf_edge.next_back_kv();
let kv = unsafe { unwrap_unchecked(kv.ok()) };
(kv.next_back_leaf_edge(), kv.into_kv())
})
}
}
impl<'a, K, V> Handle<NodeRef<marker::ValMut<'a>, K, V, marker::Leaf>, marker::Edge> {
pub unsafe fn next_unchecked(&mut self) -> (&'a K, &'a mut V) {
let kv = replace(self, |leaf_edge| {
let kv = leaf_edge.next_kv();
let kv = unsafe { unwrap_unchecked(kv.ok()) };
(unsafe { ptr::read(&kv) }.next_leaf_edge(), kv)
});
kv.into_kv_valmut()
}
pub unsafe fn next_back_unchecked(&mut self) -> (&'a K, &'a mut V) {
let kv = replace(self, |leaf_edge| {
let kv = leaf_edge.next_back_kv();
let kv = unsafe { unwrap_unchecked(kv.ok()) };
(unsafe { ptr::read(&kv) }.next_back_leaf_edge(), kv)
});
kv.into_kv_valmut()
}
}
impl<'a, K, V> Handle<NodeRef<marker::Mut<'a>, K, V, marker::Leaf>, marker::Edge> {
pub unsafe fn move_next_unchecked(&mut self) {
take_mut(self, |leaf_edge| {
let kv = leaf_edge.next_kv();
let kv = unsafe { unwrap_unchecked(kv.ok()) };
kv.next_leaf_edge()
})
}
}
impl<K, V> Handle<NodeRef<marker::Owned, K, V, marker::Leaf>, marker::Edge> {
pub unsafe fn next_unchecked<A: AllocRef>(&mut self, alloc: A) -> (K, V) {
replace(self, |leaf_edge| {
let kv = unsafe { next_kv_unchecked_dealloc(leaf_edge, alloc) };
let k = unsafe { ptr::read(kv.reborrow().into_kv().0) };
let v = unsafe { ptr::read(kv.reborrow().into_kv().1) };
(kv.next_leaf_edge(), (k, v))
})
}
pub unsafe fn next_back_unchecked<A: AllocRef>(&mut self, alloc: A) -> (K, V) {
replace(self, |leaf_edge| {
let kv = unsafe { next_back_kv_unchecked_dealloc(leaf_edge, alloc) };
let k = unsafe { ptr::read(kv.reborrow().into_kv().0) };
let v = unsafe { ptr::read(kv.reborrow().into_kv().1) };
(kv.next_back_leaf_edge(), (k, v))
})
}
}
impl<BorrowType, K, V> NodeRef<BorrowType, K, V, marker::LeafOrInternal> {
#[inline]
pub fn first_leaf_edge(self) -> Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge> {
let mut node = self;
loop {
match node.force() {
Leaf(leaf) => return leaf.first_edge(),
Internal(internal) => node = internal.first_edge().descend(),
}
}
}
#[inline]
pub fn last_leaf_edge(self) -> Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge> {
let mut node = self;
loop {
match node.force() {
Leaf(leaf) => return leaf.last_edge(),
Internal(internal) => node = internal.last_edge().descend(),
}
}
}
}
pub enum Position<BorrowType, K, V> {
Leaf(NodeRef<BorrowType, K, V, marker::Leaf>),
Internal(NodeRef<BorrowType, K, V, marker::Internal>),
InternalKV(Handle<NodeRef<BorrowType, K, V, marker::Internal>, marker::KV>),
}
impl<'a, K: 'a, V: 'a> NodeRef<marker::Immut<'a>, K, V, marker::LeafOrInternal> {
pub fn visit_nodes_in_order<F>(self, mut visit: F)
where
F: FnMut(Position<marker::Immut<'a>, K, V>),
{
match self.force() {
Leaf(leaf) => visit(Position::Leaf(leaf)),
Internal(internal) => {
visit(Position::Internal(internal));
let mut edge = internal.first_edge();
loop {
edge = match edge.descend().force() {
Leaf(leaf) => {
visit(Position::Leaf(leaf));
match edge.next_kv() {
Ok(kv) => {
visit(Position::InternalKV(kv));
kv.right_edge()
}
Err(_) => return,
}
}
Internal(internal) => {
visit(Position::Internal(internal));
internal.first_edge()
}
}
}
}
}
}
pub fn calc_length(self) -> usize {
let mut result = 0;
self.visit_nodes_in_order(|pos| match pos {
Position::Leaf(node) => result += node.len(),
Position::Internal(node) => result += node.len(),
Position::InternalKV(_) => (),
});
result
}
}
impl<BorrowType, K, V> Handle<NodeRef<BorrowType, K, V, marker::LeafOrInternal>, marker::KV> {
pub fn next_leaf_edge(self) -> Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge> {
match self.force() {
Leaf(leaf_kv) => leaf_kv.right_edge(),
Internal(internal_kv) => {
let next_internal_edge = internal_kv.right_edge();
next_internal_edge.descend().first_leaf_edge()
}
}
}
pub fn next_back_leaf_edge(
self,
) -> Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge> {
match self.force() {
Leaf(leaf_kv) => leaf_kv.left_edge(),
Internal(internal_kv) => {
let next_internal_edge = internal_kv.left_edge();
next_internal_edge.descend().last_leaf_edge()
}
}
}
}