Skip to main content

embed_avl/
lib.rs

1#![allow(rustdoc::redundant_explicit_links)]
2#![cfg_attr(docsrs, feature(doc_cfg))]
3#![cfg_attr(docsrs, allow(unused_attributes))]
4#![cfg_attr(not(feature = "std"), no_std)]
5//! An intrusive AVL tree implementation.
6//!
7//! The algorithm origin from open-zfs
8//!
9//! benchmarks: <https://github.com/NaturalIO/embed-collections-rs/wiki>
10//!
11//! The ordering of avl tree is defined by [AvlItem] trait. There're two usage:
12//! - borrow field as key from the node struct
13//! - impl Ord for struct itself
14//!
15//! We provide Entry API to achieve `CursorMut` like ability, you can use the following combination:
16//! - [AvlTree::entry], [AvlTree::first_entry], [AvlTree::last_entry]
17//! - [OccupiedEntry::peek_prev], [OccupiedEntry::peek_next]
18//! - [OccupiedEntry::move_prev], [OccupiedEntry::move_next]
19//! - [OccupiedEntry::insert_prev_unchecked], [OccupiedEntry::insert_next_unchecked]  ( does not change current `OccupiedEntry`)
20//! - [OccupiedEntry::insert_move_prev_unchecked], OccupiedEntry::insert_move_next_unchecked] (return the newlly inserted OccupiedEntry, reduce additional seek cost)
21//! - [OccupiedEntry::remove_prev] , [OccupiedEntry::remove_next] (does not change current `OccupiedEntry`)
22//!
23//!
24//! Refer to the example of [AvlItem] document.
25//!
26//! # Examples
27//!
28//! ## Using `Box` to search and remove by key
29//!
30//! ```rust
31//! use embed_avl::{AvlTree, AvlItem, AvlNode, Entry, impl_ord_with_borrow};
32//! use core::cell::UnsafeCell;
33//! use core::cmp::{Ordering, Ord};
34//! use core::borrow::Borrow;
35//! extern crate format_attr;
36//! extern crate alloc;
37//! use alloc::sync::Arc;
38//! use core::fmt;
39//!
40//! #[derive(format_attr::DebugAttr)]
41//! #[fmt("MyNode {value}")]
42//! struct MyNode {
43//!     value: i32,
44//!     avl_node: UnsafeCell<AvlNode<MyNode, ()>>,
45//! }
46//!
47//! impl Borrow<i32> for MyNode {
48//!     fn borrow(&self) -> &i32 {
49//!         &self.value
50//!     }
51//! }
52//!
53//! unsafe impl AvlItem<()> for MyNode {
54//!     fn get_node(&self) -> &mut AvlNode<MyNode, ()> {
55//!         unsafe { &mut *self.avl_node.get() }
56//!     }
57//! }
58//!
59//! impl_ord_with_borrow!(MyNode, i32);
60//!
61//! // A boxed example
62//!
63//! let mut tree = AvlTree::<Box<MyNode>, ()>::new();
64//! tree.add(Box::new(MyNode { value: 10, avl_node: UnsafeCell::new(Default::default()) })).expect("add ok");
65//!
66//! // Search and remove
67//! if let Some(node) = tree.remove_by_key::<i32>(&10i32) {
68//!     assert_eq!(node.value, 10);
69//! }
70//! assert!(tree.is_empty());
71//!
72//! // Using `Arc` for multiple ownership
73//! // remove_ref only available to `Arc` and `Rc`
74//!
75//! let mut tree = AvlTree::<Arc<MyNode>, ()>::new();
76//! let node = Arc::new(MyNode { value: 42, avl_node: UnsafeCell::new(Default::default()) });
77//!
78//! // Alternative way of insert
79//! if let Entry::Vacant(ent) = tree.entry(node.value) {
80//!     ent.insert_node(node.clone());
81//! } else {
82//!     unreachable!();
83//! }
84//! assert_eq!(tree.len(), 1);
85//!
86//! // Remove by reference (detach from avl tree)
87//! tree.remove(&node);
88//! assert_eq!(tree.len(), 0);
89//! ```
90
91extern crate alloc;
92#[cfg(any(feature = "std", test))]
93extern crate std;
94
95mod entry;
96mod iter;
97pub use entry::*;
98pub use iter::*;
99#[cfg(test)]
100mod tests;
101
102use alloc::vec::Vec;
103use core::marker::PhantomData;
104use core::{
105    borrow::Borrow,
106    cmp::{Ordering, PartialEq},
107    fmt, mem,
108    ops::Deref,
109    ptr::{NonNull, null_mut},
110};
111use pointers::{CloneablePointer, Pointer};
112
113/// A trait to return internal mutable AvlNode for specified list.
114///
115/// The tag is used to distinguish different AvlNodes within the same item,
116/// allowing an item to belong to multiple lists simultaneously.
117/// For only one ownership, you can use `()`.
118///
119/// # Safety
120///
121/// Implementors must ensure `get_node` returns a valid reference to the `AvlNode`
122/// embedded within `Self`. Users must use `UnsafeCell` to hold `AvlNode` to support
123/// interior mutability required by list operations.
124///
125/// # Example (struct field as key)
126///
127/// ```
128/// use embed_avl::{AvlItem, AvlNode, impl_ord_with_borrow};
129/// extern crate format_attr;
130/// use core::borrow::Borrow;
131/// use core::cell::UnsafeCell;
132///
133/// #[derive(format_attr::DebugAttr)]
134/// #[fmt("IntAvlNode {value}")]
135/// pub struct IntAvlNode {
136///     pub value: i64,
137///     pub node: UnsafeCell<AvlNode<Self, ()>>,
138/// }
139///
140/// unsafe impl AvlItem<()> for IntAvlNode {
141///     fn get_node(&self) -> &mut AvlNode<Self, ()> {
142///         unsafe { &mut *self.node.get() }
143///     }
144/// }
145///
146/// impl Borrow<i64> for IntAvlNode {
147///     fn borrow(&self) -> &i64 {
148///         &self.value
149///     }
150/// }
151///
152/// impl_ord_with_borrow!(IntAvlNode, i64);
153/// ```
154///
155/// # Example (impl Ord for struct)
156///
157/// ```
158/// use embed_avl::{AvlItem, AvlNode};
159/// use core::cell::{Cell, UnsafeCell};
160/// use core::cmp::Ordering;
161///
162/// pub struct RangeSeg {
163///     node: UnsafeCell<AvlNode<Self, ()>>,
164///     pub start: Cell<u64>,
165///     pub end: Cell<u64>,
166/// }
167///
168/// impl PartialOrd for RangeSeg {
169///     fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
170///         Some(Ord::cmp(self, other))
171///     }
172/// }
173/// impl PartialEq for RangeSeg {
174///     fn eq(&self, other: &Self) -> bool {
175///         Ord::cmp(self, other) == Ordering::Equal
176///     }
177/// }
178/// impl Eq for RangeSeg {}
179///
180/// impl Ord for RangeSeg {
181///     fn cmp(&self, other: &Self) -> Ordering {
182///         if self.end.get() <= other.start.get() {
183///             Ordering::Less
184///         } else if self.start.get() >= other.end.get() {
185///             Ordering::Greater
186///         } else {
187///             // intersection
188///             Ordering::Equal
189///         }
190///     }
191/// }
192/// unsafe impl Send for RangeSeg {}
193/// unsafe impl AvlItem<()> for RangeSeg {
194///     fn get_node(&self) -> &mut AvlNode<Self, ()> {
195///         unsafe { &mut *self.node.get() }
196///     }
197/// }
198/// ```
199pub unsafe trait AvlItem<Tag>: Sized {
200    fn get_node(&self) -> &mut AvlNode<Self, Tag>;
201}
202
203/// Implement Ord with `Borrow<K>`
204///
205/// # Arguments
206/// - cls: struct name
207/// - key_type: `cls: Borrow<$key_type>`
208#[macro_export]
209macro_rules! impl_ord_with_borrow {
210    ($cls: ty, $key_type: ty) => {
211        impl PartialOrd for $cls {
212            fn partial_cmp(&self, other: &Self) -> Option<core::cmp::Ordering> {
213                Some(Ord::cmp(self, other))
214            }
215        }
216        impl PartialEq for $cls {
217            fn eq(&self, other: &Self) -> bool {
218                Borrow::<$key_type>::borrow(self) == Borrow::<$key_type>::borrow(other)
219            }
220        }
221        impl Eq for $cls {}
222
223        impl Ord for $cls {
224            fn cmp(&self, other: &Self) -> core::cmp::Ordering {
225                Borrow::<$key_type>::borrow(self).cmp(Borrow::<$key_type>::borrow(other))
226            }
227        }
228    };
229}
230
231/// logging macro for development
232#[macro_export(local_inner_macros)]
233macro_rules! trace_log {
234    ($($arg:tt)+)=>{
235        #[cfg(feature="trace_log")]
236        {
237            log::debug!($($arg)+);
238        }
239    };
240}
241
242pub(crate) const DIR_LEFT: Dir = Dir(false);
243pub(crate) const DIR_RIGHT: Dir = Dir(true);
244
245#[derive(PartialEq, Debug, Copy, Clone)]
246pub(crate) struct Dir(bool);
247
248impl Dir {
249    #[inline(always)]
250    const fn to_usize(&self) -> usize {
251        self.0 as usize
252    }
253
254    #[inline(always)]
255    const fn reverse(&self) -> Self {
256        Self(!self.0)
257    }
258
259    #[inline(always)]
260    fn to_balance(&self) -> i8 {
261        if *self == DIR_LEFT { -1 } else { 1 }
262    }
263}
264
265impl From<i8> for Dir {
266    #[inline(always)]
267    fn from(i: i8) -> Self {
268        Self(i != 0)
269    }
270}
271
272impl From<bool> for Dir {
273    #[inline(always)]
274    fn from(i: bool) -> Self {
275        Self(i)
276    }
277}
278
279macro_rules! balance_to_child {
280    ($balance: expr) => {{ ($balance > 1).into() }};
281}
282
283pub struct AvlNode<T: AvlItem<Tag>, Tag> {
284    pub parent: *mut T,
285    pub children: [*mut T; 2],
286    pub balance: i8,
287    _phan: PhantomData<fn(&Tag)>,
288}
289
290unsafe impl<T: AvlItem<Tag> + Send, Tag> Send for AvlNode<T, Tag> {}
291
292impl<T: AvlItem<Tag>, Tag> AvlNode<T, Tag> {
293    #[inline(always)]
294    pub fn detach(&mut self) {
295        self.children = [null_mut(), null_mut()];
296        self.parent = null_mut();
297        self.balance = 0;
298    }
299
300    #[inline(always)]
301    fn get_child(&self, dir: Dir) -> *mut T {
302        unsafe { *self.children.get_unchecked(dir.to_usize()) }
303    }
304
305    #[inline(always)]
306    fn set_child(&mut self, dir: Dir, child: *mut T) {
307        self.children[dir.to_usize()] = child;
308    }
309
310    #[inline(always)]
311    fn left(&self) -> *mut T {
312        self.children[0]
313    }
314
315    #[inline(always)]
316    fn right(&self) -> *mut T {
317        self.children[1]
318    }
319
320    // set child and both child's parent
321    #[inline(always)]
322    fn set_child2<P: Pointer<Target = T>>(
323        &mut self, dir: Dir, child: *mut P::Target, parent: *mut P::Target,
324    ) {
325        if !child.is_null() {
326            unsafe { (*child).get_node().parent = parent };
327        }
328        self.set_child(dir, child);
329    }
330
331    #[inline(always)]
332    fn get_parent(&self) -> *mut T {
333        self.parent
334    }
335
336    #[inline(always)]
337    fn _child_direction(&self, child_p: *mut T) -> Dir {
338        // this is AvlTree::parent_direction in previous version
339        let is_right = self.children[1] == child_p;
340        #[cfg(debug_assertions)]
341        {
342            if !is_right {
343                debug_assert!(self.children[0] == child_p);
344            }
345        }
346        is_right.into()
347    }
348
349    #[inline(always)]
350    fn child_direction(parent: *mut T, child_p: *mut T) -> Dir {
351        if !parent.is_null() {
352            unsafe { (*parent).get_node() }._child_direction(child_p)
353        } else {
354            DIR_LEFT
355        }
356    }
357
358    #[inline(always)]
359    fn parent_direction(node: *mut T) -> Dir {
360        // this is AvlTree::parent_direction2 in previous version
361        let parent = unsafe { (*node).get_node() }.get_parent();
362        Self::child_direction(parent, node)
363    }
364
365    // Swap two node but not there value
366    #[inline(always)]
367    pub fn swap(&mut self, other: &mut AvlNode<T, Tag>) {
368        mem::swap(self, other);
369    }
370}
371
372impl<T: AvlItem<Tag>, Tag> Default for AvlNode<T, Tag> {
373    fn default() -> Self {
374        Self {
375            children: [null_mut(), null_mut()],
376            parent: null_mut(),
377            balance: 0,
378            _phan: Default::default(),
379        }
380    }
381}
382
383#[allow(unused_must_use)]
384impl<T: AvlItem<Tag>, Tag> fmt::Debug for AvlNode<T, Tag> {
385    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
386        write!(f, "(")?;
387        if !self.left().is_null() {
388            write!(f, "left: {:p}", self.left())?;
389        } else {
390            write!(f, "left: none ")?;
391        }
392
393        if !self.right().is_null() {
394            write!(f, "right: {:p}", self.right())?;
395        } else {
396            write!(f, "right: none ")?;
397        }
398        write!(f, ")")
399    }
400}
401
402/// An intrusive AVL tree (balanced binary search tree).
403///
404/// Elements in the tree must implement the [`AvlItem`] trait.
405/// The tree supports various pointer types (`Box`, `Arc`, `Rc`, etc.) through the [`Pointer`] trait.
406pub struct AvlTree<P, Tag>
407where
408    P: Pointer,
409    P::Target: AvlItem<Tag>,
410{
411    root: *mut P::Target,
412    count: usize,
413    _phan: PhantomData<fn(P, &Tag)>,
414}
415
416unsafe impl<P, Tag> Send for AvlTree<P, Tag>
417where
418    P: Pointer + Send,
419    P::Target: AvlItem<Tag>,
420{
421}
422
423unsafe impl<P, Tag> Sync for AvlTree<P, Tag>
424where
425    P: Pointer + Send,
426    P::Target: AvlItem<Tag>,
427{
428}
429
430impl<P, Tag> fmt::Debug for AvlTree<P, Tag>
431where
432    P: Pointer,
433    P::Target: AvlItem<Tag> + fmt::Debug,
434{
435    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
436        let _ = write!(f, "AvlTree{{");
437        if self.count > 0 {
438            let mut iter = self.iter();
439            loop {
440                let node = iter.next().unwrap();
441                let _ = write!(f, "{:?}", node.as_ref().as_ref());
442                if iter.len() > 0 {
443                    let _ = write!(f, ", ");
444                } else {
445                    break;
446                }
447            }
448        }
449        write!(f, "}}")
450    }
451}
452
453impl<P, Tag> AvlTree<P, Tag>
454where
455    P: Pointer,
456    P::Target: AvlItem<Tag>,
457{
458    /// Creates a new, empty `AvlTree`.
459    #[inline]
460    pub fn new() -> Self {
461        AvlTree { count: 0, root: null_mut(), _phan: Default::default() }
462    }
463
464    /// Returns an iterator that removes all elements from the tree in post-order.
465    ///
466    /// This is an optimized, non-recursive, and stack-less traversal that preserves
467    /// tree invariants during destruction.
468    #[inline]
469    pub fn drain(&mut self) -> AvlDrain<'_, P, Tag> {
470        AvlDrain::new(self)
471    }
472
473    #[inline]
474    pub fn is_empty(&self) -> bool {
475        self.count == 0
476    }
477
478    #[inline]
479    pub fn len(&self) -> usize {
480        self.count
481    }
482
483    /// Adds a new node to the tree, requires P::Target has `Ord`.
484    ///
485    /// (If you want more freedom of Ordering, see also [Self::entry_with])
486    ///
487    /// - Returns `Ok(())` if the node was added, takes the ownership of P.
488    /// - Err(P) if an equivalent node already exists
489    #[inline]
490    pub fn add(&mut self, node: P) -> Result<(), P>
491    where
492        P::Target: Ord,
493    {
494        if !self.root.is_null() {
495            let mut parent = self.root;
496            unsafe {
497                loop {
498                    // match is actually better optimized than if-else or table mapping
499                    let dir = match node.as_ref().cmp((&*parent).as_ref()) {
500                        Ordering::Greater => DIR_RIGHT,
501                        Ordering::Less => DIR_LEFT,
502                        Ordering::Equal => return Err(node),
503                    };
504                    let child = (*parent).get_node().get_child(dir);
505                    if !child.is_null() {
506                        parent = child;
507                    } else {
508                        self._insert(node.into_raw() as *mut P::Target, parent, dir);
509                        return Ok(());
510                    }
511                }
512            }
513        } else {
514            debug_assert_eq!(self.count, 0);
515            self.root = node.into_raw() as *mut P::Target;
516            self.count = 1;
517            Ok(())
518        }
519    }
520
521    /// remove a node from the tree (need P::Target to impl Borrow)
522    #[inline]
523    pub fn remove_by_key<K>(&mut self, key: &K) -> Option<P>
524    where
525        P::Target: Borrow<K>,
526        K: Ord,
527    {
528        let _cmp = |_key: &&K, node: &P::Target| -> Ordering { _key.cmp(&node.borrow()) };
529        if let Entry::Occupied(ent) = self.entry_with::<&K, _>(key, _cmp) {
530            Some(ent.remove())
531        } else {
532            None
533        }
534    }
535
536    /// remove a node from the tree (useful if P is multiple ownership)
537    #[inline]
538    pub fn remove(&mut self, node: &P)
539    where
540        P::Target: Ord,
541    {
542        let _cmp = |_cur: &&P::Target, node: &P::Target| -> Ordering { _cur.cmp(&node.borrow()) };
543        if let Entry::Occupied(ent) = self.entry_with::<&P::Target, _>(node.as_ref(), _cmp) {
544            ent.remove();
545        }
546    }
547
548    /// Searches for an element in the tree with `key`, requires P::Target can borrow to `K` type,
549    /// and have a total ordering by `K`.
550    #[inline]
551    pub fn entry<'a, K>(&'a mut self, key: K) -> Entry<'a, P, Tag, K>
552    where
553        P::Target: Borrow<K>,
554        K: Ord,
555    {
556        self.entry_with(key, |_key: &K, node: &P::Target| _key.cmp(node.borrow()))
557    }
558
559    /// Searches for an element in the tree with `key` and custom compare function.
560    #[inline]
561    pub fn entry_with<'a, K, F>(&'a mut self, key: K, cmp_f: F) -> Entry<'a, P, Tag, K>
562    where
563        F: Fn(&K, &P::Target) -> Ordering,
564    {
565        if !self.root.is_null() {
566            let mut parent = self.root;
567            loop {
568                // match is actually better optimized than if-else or table mapping
569                let dir = match cmp_f(&key, unsafe { &*parent }.as_ref()) {
570                    Ordering::Greater => DIR_RIGHT,
571                    Ordering::Less => DIR_LEFT,
572                    Ordering::Equal => {
573                        return Entry::Occupied(OccupiedEntry {
574                            node: unsafe { NonNull::new_unchecked(parent) },
575                            tree: self,
576                        });
577                    }
578                };
579                let child = unsafe { (*parent).get_node() }.get_child(dir);
580                if !child.is_null() {
581                    parent = child;
582                } else {
583                    return Entry::Vacant(VacantEntry {
584                        node: parent,
585                        direction: dir,
586                        tree: self,
587                        key,
588                    });
589                }
590            }
591        } else {
592            Entry::Vacant(VacantEntry { node: null_mut(), direction: DIR_RIGHT, tree: self, key })
593        }
594    }
595
596    #[inline]
597    pub fn get<'a, K>(&'a self, key: &K) -> Option<NodeRef<'a, P>>
598    where
599        P::Target: Borrow<K>,
600        K: Ord,
601    {
602        self.get_with(|node| key.cmp(node.borrow()))
603    }
604
605    pub fn get_with<'a, F>(&'a self, cmp_f: F) -> Option<NodeRef<'a, P>>
606    where
607        F: Fn(&P::Target) -> Ordering,
608    {
609        if !self.root.is_null() {
610            let mut node = self.root;
611            loop {
612                // match is actually better optimized than if-else or table mapping
613                let dir = match cmp_f(unsafe { &*node }.as_ref()) {
614                    Ordering::Greater => DIR_RIGHT,
615                    Ordering::Less => DIR_LEFT,
616                    Ordering::Equal => return Some(NodeRef::new(node)),
617                };
618                let child = unsafe { (*node).get_node() }.get_child(dir);
619                if !child.is_null() {
620                    node = child;
621                } else {
622                    return None;
623                }
624            }
625        } else {
626            None
627        }
628    }
629
630    #[allow(clippy::not_unsafe_ptr_arg_deref)]
631    #[inline]
632    pub(crate) fn _insert(
633        &mut self,
634        new_ptr: *mut P::Target,
635        parent: *mut P::Target, // parent
636        mut which_child: Dir,
637    ) {
638        if !parent.is_null() {
639            let mut new_balance: i8;
640            unsafe { (*new_ptr).get_node() }.parent = parent;
641            let parent_node = unsafe { (*parent).get_node() };
642            parent_node.set_child(which_child, new_ptr);
643            self.count += 1;
644            /*
645             * Now, back up the tree modifying the balance of all nodes above the
646             * insertion point. If we get to a highly unbalanced ancestor, we
647             * need to do a rotation.  If we back out of the tree we are done.
648             * If we brought any subtree into perfect balance (0), we are also done.
649             */
650            let mut data: *mut P::Target = parent;
651            loop {
652                let node = unsafe { (*data).get_node() };
653                let old_balance = node.balance;
654                new_balance = old_balance + which_child.to_balance();
655                if new_balance == 0 {
656                    node.balance = 0;
657                    return;
658                }
659                if old_balance != 0 {
660                    break;
661                }
662                node.balance = new_balance;
663                let parent_ptr = node.get_parent();
664                if parent_ptr.is_null() {
665                    return;
666                }
667                which_child = unsafe { (*parent_ptr).get_node() }._child_direction(data);
668                data = parent_ptr;
669            }
670            self.rotate(data, new_balance);
671        } else {
672            debug_assert_eq!(self.count, 0);
673            self.root = new_ptr;
674            self.count += 1;
675        }
676    }
677
678    /// Insert `new_ptr` in `tree` in the given `direction` against the node `here`.
679    ///
680    /// For range-tree scenario
681    ///
682    /// Insertions can only be done at empty leaf points in the tree, therefore
683    /// if the given child of the node is already present we move to either
684    /// the AVL_PREV or AVL_NEXT and reverse the insertion direction. Since
685    /// every other node in the tree is a leaf, this always works.
686    unsafe fn _insert_here(
687        &mut self, new_ptr: *mut P::Target, here: NonNull<P::Target>, direction: Dir,
688    ) {
689        crate::trace_log!("_insert_here {new_ptr:p} here {here:p} dir {direction:?}");
690        let mut dir_child = direction;
691        unsafe {
692            let child = here.as_ref().get_node().get_child(dir_child);
693            let parent;
694            if child.is_null() {
695                crate::trace_log!("direct at here {here:p}");
696                parent = here;
697            } else {
698                dir_child = dir_child.reverse();
699                parent = self.bottom_child(NonNull::new_unchecked(child), dir_child);
700                crate::trace_log!("find bottom {parent:p}");
701            }
702            self._insert(new_ptr, parent.as_ptr(), dir_child);
703        }
704    }
705
706    #[inline]
707    fn rotate(&mut self, data: *mut P::Target, balance: i8) -> bool {
708        /* when balance = -2, left will be 0 */
709        //let dir = if balance < 0 { DIR_LEFT } else { DIR_RIGHT };
710        let dir: Dir = (balance >= 0).into();
711        let node = unsafe { (*data).get_node() };
712
713        let parent = node.get_parent();
714        let dir_inverse = dir.reverse();
715        let left_heavy = balance >> 1;
716        let right_heavy = -left_heavy;
717
718        let child = node.get_child(dir);
719        let child_node = unsafe { (*child).get_node() };
720        let mut child_balance = child_node.balance;
721
722        let which_child = AvlNode::<P::Target, Tag>::child_direction(parent, data);
723
724        // node is overly left heavy, the left child is balanced or also left heavy.
725        if child_balance != right_heavy {
726            child_balance += right_heavy;
727
728            let c_right = child_node.get_child(dir_inverse);
729            node.set_child2::<P>(dir, c_right, data);
730            // move node to be child's right child
731            node.balance = -child_balance;
732
733            node.parent = child;
734            child_node.set_child(dir_inverse, data);
735            // update the pointer into this subtree
736
737            child_node.balance = child_balance;
738            if !parent.is_null() {
739                child_node.parent = parent;
740                unsafe { (*parent).get_node() }.set_child(which_child, child);
741            } else {
742                child_node.parent = null_mut();
743                self.root = child;
744            }
745            return child_balance == 0;
746        }
747        // When node is left heavy, but child is right heavy we use
748        // a different rotation.
749
750        let g_child = child_node.get_child(dir_inverse);
751        let g_child_node = unsafe { (*g_child).get_node() };
752        let g_left = g_child_node.get_child(dir);
753        let g_right = g_child_node.get_child(dir_inverse);
754
755        node.set_child2::<P>(dir, g_right, data);
756        child_node.set_child2::<P>(dir_inverse, g_left, child);
757
758        /*
759         * move child to left child of gchild and
760         * move node to right child of gchild and
761         * fixup parent of all this to point to gchild
762         */
763
764        let g_child_balance = g_child_node.balance;
765        if g_child_balance == right_heavy {
766            child_node.balance = left_heavy;
767        } else {
768            child_node.balance = 0;
769        }
770        child_node.parent = g_child;
771        g_child_node.set_child(dir, child);
772
773        if g_child_balance == left_heavy {
774            node.balance = right_heavy;
775        } else {
776            node.balance = 0;
777        }
778        g_child_node.balance = 0;
779
780        node.parent = g_child;
781        g_child_node.set_child(dir_inverse, data);
782
783        if !parent.is_null() {
784            g_child_node.parent = parent;
785            unsafe { (*parent).get_node() }.set_child(which_child, g_child);
786        } else {
787            g_child_node.parent = null_mut();
788            self.root = g_child;
789        }
790        true
791    }
792
793    /// Requires `del` to be a valid pointer to a node in this tree.
794    ///
795    /// # Safety
796    ///
797    /// It does not drop the node data, only unlinks it.
798    /// Caller is responsible for re-taking ownership (e.g. via from_raw) and dropping if needed.
799    ///
800    /// For Arc/Rc, use [Self::remove_ref()] instead.
801    ///
802    pub(crate) unsafe fn _remove(&mut self, del: *mut P::Target) {
803        /*
804         * Deletion is easiest with a node that has at most 1 child.
805         * We swap a node with 2 children with a sequentially valued
806         * neighbor node. That node will have at most 1 child. Note this
807         * has no effect on the ordering of the remaining nodes.
808         *
809         * As an optimization, we choose the greater neighbor if the tree
810         * is right heavy, otherwise the left neighbor. This reduces the
811         * number of rotations needed.
812         */
813        if self.count == 0 {
814            return;
815        }
816        if self.count == 1 && self.root == del {
817            self.root = null_mut();
818            self.count = 0;
819            unsafe { (*del).get_node().detach() };
820            return;
821        }
822        let mut which_child: Dir;
823
824        // Use reference directly to get node, avoiding unsafe dereference of raw pointer
825        let del_node = unsafe { (*del).get_node() };
826
827        let node_swap_flag = !del_node.left().is_null() && !del_node.right().is_null();
828
829        if node_swap_flag {
830            let dir: Dir = balance_to_child!(del_node.balance + 1);
831            let child_temp = del_node.get_child(dir);
832
833            let dir_inverse: Dir = dir.reverse();
834            let child = self
835                .bottom_child(unsafe { NonNull::new_unchecked(child_temp) }, dir_inverse)
836                .as_ptr();
837
838            // Fix Miri UB: Avoid calling parent_direction(child) if child's parent is del,
839            // because that would create a aliasing &mut ref to del while we hold del_node.
840            let dir_child_temp = if child == child_temp {
841                dir
842            } else {
843                AvlNode::<P::Target, Tag>::parent_direction(child)
844            };
845
846            // Fix Miri UB: Do not call parent_direction(del) as it creates a new &mut AvlNode
847            // alias while we hold del_node. Use del_node to find parent direction.
848            let parent = del_node.get_parent();
849            let dir_child_del = AvlNode::<P::Target, Tag>::child_direction(parent, del);
850
851            let child_node = unsafe { (*child).get_node() };
852            child_node.swap(del_node);
853
854            // move 'node' to delete's spot in the tree
855            if child_node.get_child(dir) == child {
856                // if node(d) left child is node(c)
857                child_node.set_child(dir, del);
858            }
859
860            let c_dir = child_node.get_child(dir);
861            if c_dir == del {
862                del_node.parent = child;
863            } else if !c_dir.is_null() {
864                unsafe { (*c_dir).get_node() }.parent = child;
865            }
866
867            let c_inv = child_node.get_child(dir_inverse);
868            if c_inv == del {
869                del_node.parent = child;
870            } else if !c_inv.is_null() {
871                unsafe { (*c_inv).get_node() }.parent = child;
872            }
873
874            let parent = child_node.get_parent();
875            if !parent.is_null() {
876                unsafe { (*parent).get_node() }.set_child(dir_child_del, child);
877            } else {
878                self.root = child;
879            }
880
881            // Put tmp where node used to be (just temporary).
882            // It always has a parent and at most 1 child.
883            let parent = del_node.get_parent();
884            unsafe { (*parent).get_node() }.set_child(dir_child_temp, del);
885            // if right exist, go to right
886            which_child = Dir(!del_node.right().is_null());
887            let child = del_node.get_child(which_child);
888            if !child.is_null() {
889                unsafe { (*child).get_node() }.parent = del;
890            }
891            which_child = dir_child_temp;
892        } else {
893            let parent = del_node.get_parent();
894            which_child = AvlNode::<P::Target, Tag>::child_direction(parent, del);
895        }
896
897        // Here we know "delete" is at least partially a leaf node. It can
898        // be easily removed from the tree.
899        let parent: *mut P::Target = del_node.get_parent();
900
901        let mut imm_data: *mut P::Target = del_node.left();
902        if !imm_data.is_null() {
903            let imm_node = unsafe { (*imm_data).get_node() };
904            imm_node.parent = parent;
905        } else {
906            imm_data = del_node.right();
907            if !imm_data.is_null() {
908                let imm_node = unsafe { (*imm_data).get_node() };
909                imm_node.parent = parent;
910            }
911        }
912        // Connect parent directly to node (leaving out delete).
913
914        if !parent.is_null() {
915            assert!(self.count > 0);
916            self.count -= 1;
917
918            let parent_node = unsafe { (*parent).get_node() };
919            parent_node.set_child(which_child, imm_data);
920
921            //Since the subtree is now shorter, begin adjusting parent balances
922            //and performing any needed rotations.
923            let mut node_data: *mut P::Target = parent;
924            loop {
925                // Move up the tree and adjust the balance.
926                // Capture the parent and which_child values for the next
927                // iteration before any rotations occur.
928                let node = unsafe { (*node_data).get_node() };
929                let balance_delta = -which_child.to_balance();
930                let old_balance = node.balance;
931
932                //If a node was in perfect balance but isn't anymore then
933                //we can stop, since the height didn't change above this point
934                //due to a deletion.
935                if old_balance == 0 {
936                    node.balance = balance_delta;
937                    break;
938                }
939                let new_balance = old_balance + balance_delta;
940
941                let parent = node.get_parent();
942                which_child = AvlNode::<P::Target, Tag>::child_direction(parent, node_data);
943
944                //If the new balance is zero, we don't need to rotate
945                //else
946                //need a rotation to fix the balance.
947                //If the rotation doesn't change the height
948                //of the sub-tree we have finished adjusting.
949                if new_balance == 0 {
950                    node.balance = 0;
951                } else if !self.rotate(node_data, new_balance) {
952                    break;
953                }
954
955                if !parent.is_null() {
956                    node_data = parent;
957                    continue;
958                }
959                break;
960            }
961        } else if !imm_data.is_null() {
962            debug_assert!(self.count > 0);
963            self.count -= 1;
964            self.root = imm_data;
965        }
966        if self.root.is_null() && self.count > 0 {
967            panic!("AvlTree {} nodes left after remove but tree.root == nil", self.count);
968        }
969        del_node.detach();
970    }
971
972    #[inline(always)]
973    fn bottom_child(&self, mut data: NonNull<P::Target>, dir: Dir) -> NonNull<P::Target> {
974        loop {
975            unsafe {
976                let child = data.as_ref().get_node().get_child(dir);
977                if !child.is_null() {
978                    data = NonNull::new_unchecked(child);
979                } else {
980                    return data;
981                }
982            }
983        }
984    }
985
986    /// return a iterator to get the reference
987    ///
988    /// NOTE: The iterator yield [NodeRef], so that you can clone the reference if it's `Arc`
989    #[inline]
990    pub fn iter(&self) -> AvlIter<'_, P, Tag> {
991        let first_p = if !self.root.is_null() {
992            Some(self.bottom_child(unsafe { NonNull::new_unchecked(self.root) }, DIR_LEFT))
993        } else {
994            None
995        };
996        AvlIter::new(self, first_p, DIR_RIGHT)
997    }
998
999    /// return a reversed iterator to get the reference.
1000    ///
1001    /// NOTE: The iterator yield [NodeRef], so that you can clone the reference if it's `Arc`
1002    #[inline]
1003    pub fn iter_rev(&self) -> AvlIter<'_, P, Tag> {
1004        let last_p = if !self.root.is_null() {
1005            Some(self.bottom_child(unsafe { NonNull::new_unchecked(self.root) }, DIR_RIGHT))
1006        } else {
1007            None
1008        };
1009        AvlIter::new(self, last_p, DIR_LEFT)
1010    }
1011
1012    #[inline]
1013    pub fn prev<'a>(&'a self, data: &'a P::Target) -> Option<NodeRef<'a, P>> {
1014        self._walk_dir(data.into(), DIR_LEFT).map(|p| NodeRef::new(p.as_ptr()))
1015    }
1016
1017    #[inline]
1018    pub fn next<'a>(&'a self, data: &'a P::Target) -> Option<NodeRef<'a, P>> {
1019        self._walk_dir(data.into(), DIR_RIGHT).map(|p| NodeRef::new(p.as_ptr()))
1020    }
1021
1022    fn _walk_dir(&self, mut data_ptr: NonNull<P::Target>, dir: Dir) -> Option<NonNull<P::Target>> {
1023        let dir_inverse = dir.reverse();
1024        let node = unsafe { data_ptr.as_ref().get_node() };
1025        let temp = node.get_child(dir);
1026        if !temp.is_null() {
1027            unsafe { Some(self.bottom_child(NonNull::new_unchecked(temp), dir_inverse)) }
1028        } else {
1029            let mut parent = node.parent;
1030            if parent.is_null() {
1031                return None;
1032            }
1033            loop {
1034                let pdir = AvlNode::<P::Target, Tag>::child_direction(parent, data_ptr.as_ptr());
1035                if pdir == dir_inverse {
1036                    return Some(unsafe { NonNull::new_unchecked(parent) });
1037                }
1038                let data_ptr_raw = parent;
1039                parent = unsafe { (*parent).get_node() }.parent;
1040                if parent.is_null() {
1041                    return None;
1042                }
1043                unsafe {
1044                    data_ptr = NonNull::new_unchecked(data_ptr_raw);
1045                }
1046            }
1047        }
1048    }
1049
1050    #[inline]
1051    pub fn first<'a>(&'a self) -> Option<NodeRef<'a, P>> {
1052        if !self.root.is_null() {
1053            let p = self.bottom_child(unsafe { NonNull::new_unchecked(self.root) }, DIR_LEFT);
1054            Some(NodeRef::<'a, P>::new(p.as_ptr()))
1055        } else {
1056            None
1057        }
1058    }
1059
1060    #[inline]
1061    pub fn last<'a>(&'a self) -> Option<NodeRef<'a, P>> {
1062        if !self.root.is_null() {
1063            let p = self.bottom_child(unsafe { NonNull::new_unchecked(self.root) }, DIR_RIGHT);
1064            Some(NodeRef::<'a, P>::new(p.as_ptr()))
1065        } else {
1066            None
1067        }
1068    }
1069
1070    #[inline]
1071    pub fn first_entry<'a>(&'a mut self) -> Option<OccupiedEntry<'a, P, Tag>> {
1072        if !self.root.is_null() {
1073            let p = self.bottom_child(unsafe { NonNull::new_unchecked(self.root) }, DIR_LEFT);
1074            Some(OccupiedEntry { node: p, tree: self })
1075        } else {
1076            None
1077        }
1078    }
1079
1080    #[inline]
1081    pub fn last_entry<'a>(&'a mut self) -> Option<OccupiedEntry<'a, P, Tag>> {
1082        if !self.root.is_null() {
1083            let p = self.bottom_child(unsafe { NonNull::new_unchecked(self.root) }, DIR_RIGHT);
1084            Some(OccupiedEntry { node: p, tree: self })
1085        } else {
1086            None
1087        }
1088    }
1089
1090    #[inline]
1091    pub fn validate(&self)
1092    where
1093        P::Target: Ord,
1094    {
1095        self.validate_with(P::Target::cmp)
1096    }
1097
1098    pub fn validate_with<F>(&self, cmp_f: F)
1099    where
1100        F: Fn(&P::Target, &P::Target) -> Ordering,
1101    {
1102        let validate_node = |data: *mut P::Target| unsafe {
1103            let node = (*data).get_node();
1104            let left = node.left();
1105            if !left.is_null() {
1106                assert!(cmp_f(&*left, &*data) != Ordering::Greater);
1107                assert_eq!((*left).get_node().get_parent(), data);
1108            }
1109            let right = node.right();
1110            if !right.is_null() {
1111                assert!(cmp_f(&*right, &*data) != Ordering::Less);
1112                assert_eq!((*right).get_node().get_parent(), data);
1113            }
1114        };
1115        let c = {
1116            #[cfg(feature = "std")]
1117            {
1118                ((self.len() + 10) as f32).log2() as usize
1119            }
1120            #[cfg(not(feature = "std"))]
1121            {
1122                100
1123            }
1124        };
1125        let mut stack: Vec<*mut P::Target> = Vec::with_capacity(c);
1126        if self.root.is_null() {
1127            assert_eq!(self.count, 0);
1128            return;
1129        }
1130        let mut data = self.root;
1131        let mut visited = 0;
1132        loop {
1133            if !data.is_null() {
1134                let left = unsafe { (*data).get_node() }.get_child(DIR_LEFT);
1135                if !left.is_null() {
1136                    crate::trace_log!("validate: push {data:p}");
1137                    stack.push(data);
1138                    data = left;
1139                    continue;
1140                }
1141            } else if !stack.is_empty() {
1142                data = stack.pop().unwrap();
1143            } else {
1144                break;
1145            }
1146            crate::trace_log!("validate: {data:p}");
1147            validate_node(data);
1148            visited += 1;
1149            data = unsafe { (*data).get_node() }.get_child(DIR_RIGHT);
1150        }
1151        assert_eq!(visited, self.count);
1152    }
1153}
1154
1155impl<P, Tag> Drop for AvlTree<P, Tag>
1156where
1157    P: Pointer,
1158    P::Target: AvlItem<Tag>,
1159{
1160    fn drop(&mut self) {
1161        if mem::needs_drop::<P>() {
1162            for _ in self.drain() {}
1163        }
1164    }
1165}
1166
1167/// A Reference to `P`
1168pub struct NodeRef<'a, P: Pointer> {
1169    node: NonNull<P::Target>,
1170    _phan: PhantomData<&'a P>,
1171}
1172
1173impl<'a, P: Pointer> AsRef<P::Target> for NodeRef<'a, P> {
1174    #[inline(always)]
1175    fn as_ref(&self) -> &P::Target {
1176        unsafe { self.node.as_ref() }
1177    }
1178}
1179
1180impl<'a, P: Pointer> Deref for NodeRef<'a, P> {
1181    type Target = P::Target;
1182
1183    #[inline(always)]
1184    fn deref(&self) -> &P::Target {
1185        unsafe { self.node.as_ref() }
1186    }
1187}
1188
1189impl<'a, P: Pointer> NodeRef<'a, P> {
1190    #[inline]
1191    fn new(p: *mut P::Target) -> Self {
1192        Self { node: unsafe { NonNull::new_unchecked(p) }, _phan: Default::default() }
1193    }
1194}
1195
1196impl<'a, P: CloneablePointer> NodeRef<'a, P> {
1197    /// If P is `Rc`/`Arc`, you can clone the reference out
1198    #[inline]
1199    pub fn clone_ref(&self) -> P {
1200        unsafe { CloneablePointer::from_raw_and_clone(self.node.as_ptr()) }
1201    }
1202}