Skip to main content

orengine_utils/
treap.rs

1//! A randomized binary search tree with subtree-augmented filtering.
2//!
3//! This module provides [`Treap`], an augmented treap that combines
4//! standard BST ordering (via [`TreapEntry::sorting_key`]) with
5//! efficient filtered min/max queries (via [`TreapEntry::filtering_key`]).
6//!
7//! Each node stores the maximum filtering key in its subtree, allowing
8//! [`pop_max_with_filter`](Treap::pop_max_with_filter) and related
9//! methods to prune entire branches.
10use crate::cheap_random::cheap_random_with_current_u32;
11use crate::hints::{cold_path, unlikely, unreachable_hint, unwrap_or_bug_hint};
12use crate::ArrayBuffer;
13use alloc::boxed::Box;
14use alloc::format;
15use core::cmp::Ordering;
16use core::fmt;
17use core::fmt::Display;
18use core::mem;
19use core::num::NonZeroU32;
20use core::ops::Deref;
21use core::ptr::NonNull;
22
23/// Trait for entries that can be stored in a [`Treap`].
24///
25/// Implementors provide two types of keys:
26/// - **Sorting Key**: Used for BST (binary search tree) ordering
27/// - **Filtering Key**: Used for subtree pruning during range queries
28///
29/// # Examples
30///
31/// ```rust
32/// use core::cmp::Ordering;
33/// use orengine_utils::treap::TreapEntry;
34///
35/// #[derive(Clone, Copy)]
36/// struct Point {
37///     x: i32,
38///     y: i32,
39///     payload: i32,
40/// }
41///
42/// impl TreapEntry for Point {
43///     type SortingKey = i32;
44///     type FilteringKey = i32;
45///     type Value = i32;
46///
47///     fn sorting_key(&self) -> &Self::SortingKey {
48///         &self.x
49///     }
50///
51///     fn filtering_key(&self) -> &Self::FilteringKey {
52///         &self.y
53///     }
54///
55///     fn value(&self) -> &Self::Value {
56///         &self.payload
57///     }
58///
59///     fn value_mut(&mut self) -> &mut Self::Value {
60///         &mut self.payload
61///     }
62/// }
63/// ```
64pub trait TreapEntry {
65    /// Key type for BST ordering. Can be a tuple for compound keys.
66    type SortingKey: Ord;
67
68    /// Key type for filtering/pruning.
69    type FilteringKey: Ord + Clone;
70
71    /// Value type stored in the entry.
72    type Value;
73
74    /// Returns the sorting key for BST ordering.
75    fn sorting_key(&self) -> &Self::SortingKey;
76
77    /// Returns the filtering key for pruning decisions.
78    fn filtering_key(&self) -> &Self::FilteringKey;
79
80    /// Returns a shared reference to the underlying value.
81    fn value(&self) -> &Self::Value;
82
83    /// Returns an exclusive reference to the underlying value.
84    fn value_mut(&mut self) -> &mut Self::Value;
85}
86
87/// Node in the [`Treap`].
88///
89/// It can be dereferenced into a shared reference to the [`TreapEntry`].
90/// You can also use [`Node::neighbors`] to get an iterator over the neighbors of this node.
91pub struct Node<E: TreapEntry> {
92    entry: E,
93
94    /// Maximum filter key in this subtree
95    max_filter: E::FilteringKey,
96
97    priority: u32,
98    left: Option<NonNull<Node<E>>>,
99    right: Option<NonNull<Node<E>>>,
100    parent: Option<NonNull<Node<E>>>,
101}
102
103impl<E: TreapEntry> Node<E> {
104    /// Returns a shared reference to the entry.
105    #[inline]
106    fn entry(&self) -> &E {
107        &self.entry
108    }
109
110    /// Returns an iterator over nodes reachable from `self` whose
111    /// `filtering_key() >= filter`.
112    ///
113    /// Traversal visits both subtree descendants and ancestors,
114    /// pruning branches where the subtree's `max_filter < filter`.
115    ///
116    /// `skip_right` skips the right subtree on the first step, useful when
117    /// the caller has already consumed the greatest node (e.g., after
118    /// [`Treap::peek_max_with_filter`]).
119    ///
120    /// # Example
121    ///
122    /// ```rust
123    /// use orengine_utils::treap::{BaseTreapEntry, Treap};
124    ///
125    /// let mut treap = Treap::<BaseTreapEntry<usize, usize, ()>>::new();
126    ///
127    /// for i in 1..=5 {
128    ///     treap.set(BaseTreapEntry::new(i, i, ()));
129    /// }
130    ///
131    /// let node = treap.peek_max_with_filter(&3).unwrap();
132    ///
133    /// let keys: Vec<_> = node.neighbors(&3, true) // `true` because we already know the greatest entry with `FilteringKey` >= 3
134    ///     .map(|n| n.sorting_key)
135    ///     .collect();
136    ///
137    /// assert_eq!(keys, vec![5, 4, 3]);
138    /// ```
139    #[allow(
140        clippy::too_many_lines,
141        reason = "There is only a slight excess here,\
142         but moving the iterator out will ruin readability."
143    )]
144    pub fn neighbors<'treap>(
145        &'treap self,
146        filter: &'treap E::FilteringKey,
147        skip_right: bool,
148    ) -> impl Iterator<Item = &'treap Self> {
149        /// An internal state machine for the neighbor iterator,
150        /// controlling traversal direction through the treap.
151        #[derive(Clone, Copy, PartialEq)]
152        enum Phase {
153            Start,              // Initial state for the node
154            StartAndNextGoLeft, // Initial state for the node
155            GoRight,            // Go to right child
156            GoLeft,             // Go to left child
157            GoUp,               // Return to parent
158        }
159
160        struct NeighborsIterator<'treap, E: TreapEntry> {
161            current: NonNull<Node<E>>,
162            filter: &'treap E::FilteringKey,
163            // Depth relative to the start node.
164            // 0: Ancestors or Start Node. > 0: Descendants.
165            // -1 is used transiently to detect moving into an ancestor.
166            depth: i32,
167            phase: Phase,
168        }
169
170        impl<E: TreapEntry> NeighborsIterator<'_, E> {
171            /// Checks if a subtree contains nodes meeting the filter threshold.
172            fn check_subtree(
173                node: Option<NonNull<Node<E>>>,
174                filter: &E::FilteringKey,
175            ) -> Option<NonNull<Node<E>>> {
176                node.filter(|n| unsafe { n.as_ref().max_filter >= *filter })
177            }
178        }
179
180        impl<'a, E: TreapEntry + 'a> Iterator for NeighborsIterator<'a, E> {
181            type Item = &'a Node<E>;
182
183            fn next(&mut self) -> Option<Self::Item> {
184                loop {
185                    unsafe {
186                        let node = self.current.as_ref();
187
188                        match self.phase {
189                            Phase::Start => {
190                                // 1. Yield Current Node
191                                self.phase = Phase::GoRight;
192
193                                if node.entry.filtering_key() >= self.filter {
194                                    return Some(node);
195                                }
196                            }
197
198                            Phase::StartAndNextGoLeft => {
199                                self.phase = Phase::GoLeft;
200
201                                if node.entry.filtering_key() >= self.filter {
202                                    return Some(node);
203                                }
204                            }
205
206                            Phase::GoRight => {
207                                // 2. Try Right Child
208                                self.phase = Phase::GoLeft;
209
210                                if let Some(right) = Self::check_subtree(node.right, self.filter) {
211                                    self.current = right;
212                                    self.depth += 1;
213                                    self.phase = Phase::Start;
214                                }
215                            }
216
217                            Phase::GoLeft => {
218                                // 3. Try Left Child
219                                self.phase = Phase::GoUp;
220
221                                if let Some(left) = Self::check_subtree(node.left, self.filter) {
222                                    self.current = left;
223                                    self.depth += 1;
224                                    self.phase = Phase::Start;
225                                }
226                            }
227
228                            Phase::GoUp => {
229                                // 4. Go Up
230                                if let Some(parent_ptr) = node.parent {
231                                    let parent = parent_ptr.as_ref();
232
233                                    // Determine relationship BEFORE moving current pointer logic fully
234                                    let is_right_child = parent.right == Some(self.current);
235
236                                    self.current = parent_ptr;
237                                    self.depth -= 1;
238
239                                    if self.depth < 0 {
240                                        // We just stepped into an Ancestor (relative to start)
241                                        self.depth = 0;
242
243                                        // Yield the ancestor if valid
244                                        if parent.entry.filtering_key() >= self.filter {
245                                            // After yielding, determine next step
246                                            if is_right_child {
247                                                self.phase = Phase::GoLeft;
248                                            } else {
249                                                self.phase = Phase::GoUp;
250                                            }
251                                            return Some(parent);
252                                        }
253
254                                        // The ancestor is invalid, but maybe another subtree is valid
255                                        if is_right_child {
256                                            self.phase = Phase::GoLeft;
257                                        } else {
258                                            self.phase = Phase::GoUp;
259                                        }
260                                    } else {
261                                        // We are bubbling up inside the Start Node's subtree.
262                                        // We visited this node and its Right child.
263                                        // If we came from Right, we must now check Left.
264                                        // If we came from Left, we are done with this node (Go Up).
265                                    }
266
267                                    if is_right_child {
268                                        self.phase = Phase::GoLeft;
269                                    } else {
270                                        self.phase = Phase::GoUp;
271                                    }
272                                } else {
273                                    // Root reached
274                                    return None;
275                                }
276                            }
277                        }
278                    }
279                }
280            }
281        }
282
283        let initial_phase = if skip_right {
284            Phase::StartAndNextGoLeft
285        } else {
286            Phase::Start
287        };
288
289        NeighborsIterator {
290            current: NonNull::from(self),
291            filter,
292            depth: 0,
293            phase: initial_phase,
294        }
295    }
296}
297
298impl<E: TreapEntry> Deref for Node<E> {
299    type Target = E;
300
301    fn deref(&self) -> &Self::Target {
302        self.entry()
303    }
304}
305
306/// A treap (tree and heap) data structure combining a binary search tree with randomized
307/// heap-based balancing.
308///
309/// Supports efficient filtering via subtree augmentation with maximum filtering keys.
310///
311/// # Example
312///
313/// ```rust
314/// use orengine_utils::treap::{Node, Treap, TreapEntry};
315///
316/// #[derive(Clone, Copy)]
317/// struct Point {
318///     x: i32,
319///     y: i32,
320///     payload: i32,
321/// }
322///
323/// impl TreapEntry for Point {
324///     type SortingKey = i32;
325///     type FilteringKey = i32;
326///     type Value = i32;
327///
328///     fn sorting_key(&self) -> &Self::SortingKey {
329///         &self.x
330///     }
331///
332///     fn filtering_key(&self) -> &Self::FilteringKey {
333///         &self.y
334///     }
335///
336///     fn value(&self) -> &Self::Value {
337///         &self.payload
338///     }
339///
340///     fn value_mut(&mut self) -> &mut Self::Value {
341///         &mut self.payload
342///     }
343/// }
344///
345/// let mut treap = Treap::new();
346///
347/// treap.set(Point { x: 1, y: 5, payload: -1 });
348/// treap.set(Point { x: 2, y: 4, payload: -2 });
349/// treap.set(Point { x: 3, y: 3, payload: -3 });
350/// treap.set(Point { x: 4, y: 2, payload: -4 });
351/// treap.set(Point { x: 5, y: 1, payload: -5 });
352///
353/// let the_greatest: &Node<Point> = treap.peek_max_with_filter(&4).unwrap(); // Peeks the greatest by `x` entry with `y` >= 4
354///
355/// assert_eq!(the_greatest.x, 2);
356///
357/// // Iterator over points with `y` >= 4 around `the_greatest` to start from the greatest by `x`
358/// // `true` for `skip_right` because we already know the greatest and want to search for `y <= the_greatest.y`.
359/// let neighbors = the_greatest.neighbors(&4, true);
360///
361/// assert_eq!(neighbors.map(|point| point.payload).collect::<Vec<i32>>(), vec![-2, -1]);
362/// ```
363pub struct Treap<E: TreapEntry> {
364    root: Option<NonNull<Node<E>>>,
365    rng: NonZeroU32,
366}
367
368/// Node in the [`Treap`].
369///
370/// It can be dereferenced into a shared reference to the [`TreapEntry`], and it provides
371/// a mutable reference to the underlying value by [`NodeMut::value_mut`].
372///
373/// You can also use [`Node::neighbors`] to get an iterator over the neighbors of this node.
374///
375/// And it can be used to remove the node from the treap by [`NodeMut::remove_from_treap`].
376pub struct NodeMut<'handle, E: TreapEntry> {
377    treap: &'handle mut Treap<E>,
378    node_ptr: NonNull<Node<E>>,
379}
380
381impl<E: TreapEntry> NodeMut<'_, E> {
382    /// Returns an exclusive reference to the value of this node.
383    pub fn value_mut(&mut self) -> &mut E::Value {
384        let node = unsafe { self.node_ptr.as_mut() };
385
386        node.entry.value_mut()
387    }
388
389    /// Removes this node from the associated [`Treap`].
390    ///
391    /// # Example
392    ///
393    /// ```rust
394    /// use orengine_utils::treap::{BaseTreapEntry, Treap};
395    ///
396    /// let mut treap = Treap::<BaseTreapEntry<u32, u32, ()>>::new();
397    ///
398    /// unsafe { treap.add(BaseTreapEntry::new(1, 1, ())) };
399    ///
400    /// let mut node = treap.peek_max_with_filter_mut(&1).unwrap();
401    /// if node.sorting_key > 0 { // Remove on condition, you can use the node not to use search it again
402    ///     let entry = node.remove_from_treap();
403    ///
404    ///     assert_eq!(entry.sorting_key, 1);
405    ///     assert!(treap.is_empty());
406    /// }
407    /// ```
408    pub fn remove_from_treap(&mut self) -> E {
409        self.treap.remove_node_by_ptr(self.node_ptr)
410    }
411}
412
413impl<E: TreapEntry> Deref for NodeMut<'_, E> {
414    type Target = Node<E>;
415
416    fn deref(&self) -> &Self::Target {
417        unsafe { self.node_ptr.as_ref() }
418    }
419}
420
421impl<E: TreapEntry> Treap<E> {
422    /// Creates a new empty treap.
423    ///
424    /// # Example
425    ///
426    /// ```rust
427    /// use orengine_utils::treap::{BaseTreapEntry, Treap};
428    ///
429    /// let treap = Treap::<BaseTreapEntry<u32, u32, ()>>::new();
430    ///
431    /// assert!(treap.is_empty());
432    /// ```
433    #[inline]
434    pub const fn new() -> Self {
435        Self {
436            root: None,
437            rng: unsafe { NonZeroU32::new_unchecked(1_406_868_647) },
438        }
439    }
440
441    /// Returns `true` if the treap contains no entries.
442    ///
443    /// # Example
444    ///
445    /// ```rust
446    /// use orengine_utils::treap::{BaseTreapEntry, Treap};
447    ///
448    /// let mut treap = Treap::<BaseTreapEntry<u32, u32, ()>>::new();
449    ///
450    /// assert!(treap.is_empty());
451    ///
452    /// treap.set(BaseTreapEntry::new(1, 1, ()));
453    ///
454    /// assert!(!treap.is_empty());
455    /// ```
456    #[inline]
457    pub fn is_empty(&self) -> bool {
458        self.root.is_none()
459    }
460
461    /// Recomputes the `max_filter` for `node` from its entry and children.
462    /// Must be called after any structural change to `node`'s children.
463    #[inline]
464    fn update_augmentation_(node: &mut Node<E>) {
465        let mut max_f = node.entry.filtering_key();
466
467        if let Some(left_ptr) = node.left {
468            let left = unsafe { left_ptr.as_ref() };
469            if left.max_filter > *max_f {
470                max_f = &left.max_filter;
471            }
472        }
473
474        if let Some(right_ptr) = node.right {
475            let right = unsafe { right_ptr.as_ref() };
476            if right.max_filter > *max_f {
477                max_f = &right.max_filter;
478            }
479        }
480
481        node.max_filter = max_f.clone();
482    }
483
484    /// Recomputes the `max_filter` for the node at `node`.
485    /// This is a pointer-based wrapper around [`update_augmentation_`].
486    fn update_augmentation(mut node: NonNull<Node<E>>) {
487        let node = unsafe { node.as_mut() };
488
489        Self::update_augmentation_(node);
490    }
491
492    /// Performs a right rotation at the given node.
493    fn rotate_right(&mut self, mut x: NonNull<Node<E>>) {
494        let x_ref = unsafe { x.as_mut() };
495        let mut y = unwrap_or_bug_hint(x_ref.left);
496        let y_ref = unsafe { y.as_mut() };
497
498        // Update x's left child
499        x_ref.left = y_ref.right;
500        if let Some(right_of_y) = y_ref.right {
501            unsafe { (*right_of_y.as_ptr()).parent = Some(x) };
502        }
503
504        // Update y's parent
505        y_ref.parent = x_ref.parent;
506
507        // Update parent's child pointer
508        match x_ref.parent {
509            Some(p) => {
510                let parent = unsafe { &mut *p.as_ptr() };
511                if parent.left == Some(x) {
512                    parent.left = Some(y);
513                } else {
514                    parent.right = Some(y);
515                }
516            }
517            None => {
518                self.root = Some(y);
519            }
520        }
521
522        // Complete the rotation
523        y_ref.right = Some(x);
524        x_ref.parent = Some(y);
525
526        // Update augmentations: only x and y changed their children
527        Self::update_augmentation_(x_ref);
528        Self::update_augmentation_(y_ref);
529
530        // Update ancestors if needed
531        if let Some(parent_ptr) = y_ref.parent {
532            Self::update_augmentation(parent_ptr);
533        }
534    }
535
536    /// Performs a left rotation at the given node.
537    fn rotate_left(&mut self, mut x: NonNull<Node<E>>) {
538        let x_ref = unsafe { x.as_mut() };
539        let mut y = unwrap_or_bug_hint(x_ref.right);
540        let y_ref = unsafe { y.as_mut() };
541
542        // Update x's right child
543        x_ref.right = y_ref.left;
544        if let Some(left_of_y) = y_ref.left {
545            unsafe { (*left_of_y.as_ptr()).parent = Some(x) };
546        }
547
548        // Update y's parent
549        y_ref.parent = x_ref.parent;
550
551        // Update parent's child pointer
552        match x_ref.parent {
553            Some(p) => {
554                let parent = unsafe { &mut *p.as_ptr() };
555                if parent.left == Some(x) {
556                    parent.left = Some(y);
557                } else {
558                    parent.right = Some(y);
559                }
560            }
561            None => {
562                self.root = Some(y);
563            }
564        }
565
566        // Complete the rotation
567        y_ref.left = Some(x);
568        x_ref.parent = Some(y);
569
570        // Update augmentations: only x and y changed their children
571        Self::update_augmentation_(x_ref);
572        Self::update_augmentation_(y_ref);
573
574        // Update ancestors if needed
575        if let Some(parent_ptr) = y_ref.parent {
576            Self::update_augmentation(parent_ptr);
577        }
578    }
579
580    /// Moves `node` upward via rotations until the heap priority property
581    /// (parent.priority >= child.priority) is restored.
582    fn bubble_up(&mut self, node: NonNull<Node<E>>) {
583        loop {
584            let node_ref = unsafe { node.as_ref() };
585
586            let Some(parent) = node_ref.parent else { break };
587
588            let parent_ref = unsafe { parent.as_ref() };
589            if parent_ref.priority >= node_ref.priority {
590                break;
591            }
592
593            if parent_ref.left == Some(node) {
594                self.rotate_right(parent);
595            } else {
596                self.rotate_left(parent);
597            }
598        }
599    }
600
601    /// Propagates updated `max_filter` values upward from `node` toward the
602    /// root, stopping early once a node is found whose `max_filter` is
603    /// greater than `deleted_filtering_key`.
604    fn update_filter_to_root_after_deleting(
605        mut node: NonNull<Node<E>>,
606        deleted_filtering_key: &E::FilteringKey,
607    ) {
608        loop {
609            let node_ref = unsafe { node.as_ref() };
610            let parent = node_ref.parent;
611
612            Self::update_augmentation(node);
613
614            if let Some(parent_ptr) = parent {
615                let parent_ref = unsafe { parent_ptr.as_ref() };
616                if &parent_ref.max_filter == deleted_filtering_key {
617                    node = parent_ptr;
618                } else {
619                    if cfg!(test) {
620                        assert!(&parent_ref.max_filter > deleted_filtering_key);
621                    }
622
623                    break;
624                }
625            } else {
626                break;
627            }
628        }
629    }
630
631    /// Core insert/update implementation.
632    ///
633    /// When `ONLY_ADD` is `true` (used by [`add`](Treap::add)), panics in
634    /// debug mode if the key already exists.  When `false` (used by
635    /// [`set`](Treap::set)), replaces the existing entry and returns it.
636    ///
637    /// Returns `(node_ptr, old_entry)`.
638    fn set_<const ONLY_ADD: bool>(&mut self, entry: E) -> (NonNull<Node<E>>, Option<E>) {
639        let mut new_node = |entry: E| -> NonNull<Node<E>> {
640            let max_filter = entry.filtering_key().clone();
641            let node = Box::new(Node {
642                entry,
643                priority: cheap_random_with_current_u32(&mut self.rng),
644                left: None,
645                right: None,
646                parent: None,
647                max_filter,
648            });
649
650            unsafe { NonNull::new_unchecked(Box::into_raw(node)) }
651        };
652
653        // Handle empty tree
654        if unlikely(self.root.is_none()) {
655            let node_ptr = new_node(entry);
656
657            self.root = Some(node_ptr);
658
659            return (node_ptr, None);
660        }
661
662        let mut current_ = self.root;
663        let mut prev_node: Option<NonNull<Node<E>>> = None;
664        let mut is_left_child = false;
665
666        unsafe {
667            while let Some(mut current) = current_ {
668                prev_node = Some(current);
669
670                if entry.filtering_key() > &current.as_mut().max_filter {
671                    current.as_mut().max_filter = entry.filtering_key().clone();
672                }
673
674                match (
675                    current
676                        .as_mut()
677                        .entry
678                        .sorting_key()
679                        .cmp(entry.sorting_key()),
680                    ONLY_ADD,
681                ) {
682                    (Ordering::Greater, _) => {
683                        current_ = current.as_mut().left;
684                        is_left_child = true;
685                    }
686                    (Ordering::Less, _) => {
687                        current_ = current.as_mut().right;
688                        is_left_child = false;
689                    }
690                    (Ordering::Equal, true) => {
691                        if cfg!(debug_assertions) {
692                            panic!("The method `add` can only adds entries, not update them");
693                        } else {
694                            unreachable_hint();
695                        }
696                    }
697                    (Ordering::Equal, false) => {
698                        let old_entry = mem::replace(&mut current.as_mut().entry, entry);
699                        if unlikely(current.as_mut().filtering_key() < old_entry.filtering_key()) {
700                            Self::update_augmentation(current);
701
702                            if let Some(parent_ptr) = current.as_mut().parent {
703                                Self::update_filter_to_root_after_deleting(
704                                    parent_ptr,
705                                    old_entry.filtering_key(),
706                                );
707                            }
708                        } // else the treap augmentation is already correct
709
710                        return (current, Some(old_entry));
711                    }
712                }
713            }
714
715            let mut node_ptr = new_node(entry);
716
717            // Link to parent
718            node_ptr.as_mut().parent = prev_node;
719            if is_left_child {
720                unwrap_or_bug_hint(prev_node).as_mut().left = Some(node_ptr);
721            } else {
722                unwrap_or_bug_hint(prev_node).as_mut().right = Some(node_ptr);
723            }
724
725            self.bubble_up(node_ptr);
726
727            (node_ptr, None)
728        }
729    }
730
731    /// Inserts `entry` into the treap, returning [`NodeMut`] with this entry.
732    ///
733    /// # Safety
734    ///
735    /// The treap must not yet contain an entry with the same `SortingKey`.
736    /// In debug builds, violating this panics; in release builds it is UB.
737    ///
738    /// Because this function does not check if it already contains the entry,
739    /// it is faster than [`Treap::set`].
740    ///
741    /// # Example
742    ///
743    /// ```rust
744    /// use orengine_utils::treap::{Treap, BaseTreapEntry, NodeMut};
745    ///
746    /// let mut treap = Treap::<BaseTreapEntry<u32, u32, ()>>::new();
747    /// let node_mut: NodeMut<'_, BaseTreapEntry<u32, u32, ()>> = unsafe {
748    ///     treap.add(BaseTreapEntry::new(42, 7, ()))
749    /// };
750    ///
751    /// assert_eq!(treap.find(&42).unwrap().sorting_key, 42);
752    /// ```
753    pub unsafe fn add(&mut self, entry: E) -> NodeMut<'_, E> {
754        let node_ptr = self.set_::<true>(entry).0;
755
756        NodeMut {
757            treap: self,
758            node_ptr,
759        }
760    }
761
762    /// Inserts `entry`, or replaces the existing entry with the same
763    /// `SortingKey` if one exists.
764    ///
765    /// Returns <code>([NodeMut], Some(old_entry))</code> on a replacement, or
766    /// <code>([NodeMut], None)</code> on an insertion.
767    ///
768    /// # Example
769    ///
770    /// ```rust
771    /// use orengine_utils::treap::{Treap, BaseTreapEntry};
772    ///
773    /// let mut treap = Treap::<BaseTreapEntry<u32, u32, &str>>::new();
774    ///
775    /// let (_node_mut, old) = treap.set(BaseTreapEntry::new(1, 10, "first"));
776    /// assert!(old.is_none());
777    ///
778    /// let (_node_mut, old) = treap.set(BaseTreapEntry::new(1, 20, "second"));
779    /// assert_eq!(old.unwrap().value, "first");
780    /// ```
781    pub fn set(&mut self, entry: E) -> (NodeMut<'_, E>, Option<E>) {
782        let (node_ptr, old_entry) = self.set_::<false>(entry);
783
784        (
785            NodeMut {
786                treap: self,
787                node_ptr,
788            },
789            old_entry,
790        )
791    }
792
793    /// Rotates `node` downward until it becomes a leaf, preserving the heap
794    /// property for all other nodes. Used as the first step of deletion.
795    ///
796    /// # Safety
797    ///
798    /// `node` must be a valid pointer to a node currently in this treap.
799    unsafe fn rotate_down_to_leaf(&mut self, node: NonNull<Node<E>>) {
800        loop {
801            let node_ref = unsafe { node.as_ref() };
802            let left_priority = node_ref
803                .left
804                .map_or(0, |l| unsafe { (*l.as_ptr()).priority });
805            let right_priority = node_ref
806                .right
807                .map_or(0, |r| unsafe { (*r.as_ptr()).priority });
808
809            if node_ref.left.is_none() && node_ref.right.is_none() {
810                break;
811            }
812
813            if left_priority > right_priority {
814                self.rotate_right(node);
815            } else if right_priority > 0 {
816                self.rotate_left(node);
817            } else {
818                break;
819            }
820        }
821    }
822
823    /// Removes a node from the treap by its pointer.
824    fn remove_node_by_ptr(&mut self, node: NonNull<Node<E>>) -> E {
825        let node_ptr = node;
826
827        unsafe {
828            self.rotate_down_to_leaf(node_ptr);
829        }
830
831        let node_ref = unsafe { node.as_ref() };
832
833        if let Some(parent) = node_ref.parent {
834            let parent_ref = unsafe { &mut *parent.as_ptr() };
835            if parent_ref.left == Some(node_ptr) {
836                parent_ref.left = None;
837            } else {
838                parent_ref.right = None;
839            }
840
841            Self::update_filter_to_root_after_deleting(parent, node_ref.filtering_key());
842        } else {
843            self.root = None;
844        }
845
846        unsafe { Box::from_raw(node_ptr.as_ptr()) }.entry
847    }
848
849    /// Removes and returns the entry with `sorting_key`, or `None` if absent.
850    ///
851    /// # Example
852    ///
853    /// ```rust
854    /// use orengine_utils::treap::{BaseTreapEntry, Treap};
855    ///
856    /// let mut treap = Treap::<BaseTreapEntry<u32, u32, ()>>::new();
857    ///
858    /// treap.set(BaseTreapEntry::new(5, 5, ()));
859    ///
860    /// assert!(treap.remove_by_sorting_key(&5).is_some());
861    /// assert!(treap.remove_by_sorting_key(&5).is_none());
862    /// ```
863    pub fn remove_by_sorting_key(&mut self, sorting_key: &E::SortingKey) -> Option<E> {
864        let node = self.find_ptr_mut(sorting_key)?;
865
866        Some(self.remove_node_by_ptr(node))
867    }
868
869    /// Returns a pointer to the node with the greatest sorting key, or `None`.
870    fn find_max_ptr(&self) -> Option<NonNull<Node<E>>> {
871        let mut current = self.root?;
872
873        unsafe {
874            while let Some(right) = (*current.as_ptr()).right {
875                current = right;
876            }
877        }
878
879        Some(current)
880    }
881
882    /// Removes and returns the entry with the greatest sorting key, or `None`.
883    ///
884    /// # Example
885    ///
886    /// ```rust
887    /// use orengine_utils::treap::{BaseTreapEntry, Treap};
888    ///
889    /// let mut treap = Treap::<BaseTreapEntry<u32, u32, ()>>::new();
890    ///
891    /// treap.set(BaseTreapEntry::new(1, 1, ()));
892    /// treap.set(BaseTreapEntry::new(3, 3, ()));
893    ///
894    /// assert_eq!(treap.pop_max().unwrap().sorting_key, 3);
895    /// assert_eq!(treap.pop_max().unwrap().sorting_key, 1);
896    /// ```
897    pub fn pop_max(&mut self) -> Option<E> {
898        let min_node = self.find_max_ptr()?;
899
900        Some(self.remove_node_by_ptr(min_node))
901    }
902
903    /// Returns [`Node`] with the greatest sorting key without removing it.
904    ///
905    /// # Example
906    ///
907    /// ```rust
908    /// use orengine_utils::treap::{BaseTreapEntry, Treap};
909    ///
910    /// let mut treap = Treap::<BaseTreapEntry<u32, u32, ()>>::new();
911    ///
912    /// treap.set(BaseTreapEntry::new(1, 1, ()));
913    /// treap.set(BaseTreapEntry::new(3, 3, ()));
914    ///
915    /// assert_eq!(treap.peek_max().unwrap().sorting_key, 3);
916    /// assert_eq!(treap.peek_max().unwrap().sorting_key, 3);
917    /// ```
918    pub fn peek_max(&self) -> Option<&Node<E>> {
919        self.find_max_ptr().map(|ptr| unsafe { ptr.as_ref() })
920    }
921
922    /// Returns [`NodeMut`] with the greatest sorting key without removing it.
923    ///
924    /// # Example
925    ///
926    /// ```rust
927    /// use orengine_utils::treap::{BaseTreapEntry, Treap};
928    ///
929    /// let mut treap = Treap::<BaseTreapEntry<u32, u32, ()>>::new();
930    ///
931    /// treap.set(BaseTreapEntry::new(1, 1, ()));
932    /// treap.set(BaseTreapEntry::new(3, 3, ()));
933    ///
934    /// assert_eq!(treap.peek_max_mut().unwrap().sorting_key, 3);
935    /// assert_eq!(treap.peek_max_mut().unwrap().sorting_key, 3);
936    ///
937    /// let mut node_mut = treap.peek_max_mut().unwrap();
938    /// if node_mut.sorting_key > 2 { // conditional remove
939    ///     node_mut.remove_from_treap();
940    /// }
941    /// ```
942    pub fn peek_max_mut(&mut self) -> Option<NodeMut<'_, E>> {
943        let mut current = self.root?;
944
945        unsafe {
946            while let Some(right) = (*current.as_ptr()).right {
947                current = right;
948            }
949        }
950
951        Some(NodeMut {
952            node_ptr: current,
953            treap: self,
954        })
955    }
956
957    /// Returns a pointer to the node with the smallest sorting key, or `None`.
958    fn find_min_ptr(&self) -> Option<NonNull<Node<E>>> {
959        let mut current = self.root?;
960
961        unsafe {
962            while let Some(left) = (*current.as_ptr()).left {
963                current = left;
964            }
965        }
966
967        Some(current)
968    }
969
970    /// Removes and returns the entry with the smallest sorting key, or `None`.
971    ///
972    /// # Example
973    ///
974    /// ```rust
975    /// use orengine_utils::treap::{BaseTreapEntry, Treap};
976    ///
977    /// let mut treap = Treap::<BaseTreapEntry<u32, u32, ()>>::new();
978    ///
979    /// treap.set(BaseTreapEntry::new(1, 1, ()));
980    /// treap.set(BaseTreapEntry::new(3, 3, ()));
981    ///
982    /// assert_eq!(treap.pop_min().unwrap().sorting_key, 1);
983    /// assert_eq!(treap.pop_min().unwrap().sorting_key, 3);
984    /// ```
985    pub fn pop_min(&mut self) -> Option<E> {
986        let min_node = self.find_min_ptr()?;
987
988        Some(self.remove_node_by_ptr(min_node))
989    }
990
991    /// Returns [`Node`] with the smallest sorting key without removing it.
992    ///
993    /// ```rust
994    /// use orengine_utils::treap::{BaseTreapEntry, Treap};
995    ///
996    /// let mut treap = Treap::<BaseTreapEntry<u32, u32, ()>>::new();
997    ///
998    /// treap.set(BaseTreapEntry::new(1, 1, ()));
999    /// treap.set(BaseTreapEntry::new(3, 3, ()));
1000    ///
1001    /// assert_eq!(treap.peek_min().unwrap().sorting_key, 1);
1002    /// assert_eq!(treap.peek_min().unwrap().sorting_key, 1);
1003    /// ```
1004    pub fn peek_min(&self) -> Option<&Node<E>> {
1005        self.find_min_ptr().map(|ptr| unsafe { ptr.as_ref() })
1006    }
1007
1008    /// Returns [`NodeMut`] with the smallest sorting key without removing it.
1009    ///
1010    /// # Example
1011    ///
1012    /// ```rust
1013    /// use orengine_utils::treap::{BaseTreapEntry, Treap};
1014    ///
1015    /// let mut treap = Treap::<BaseTreapEntry<u32, u32, ()>>::new();
1016    ///
1017    /// treap.set(BaseTreapEntry::new(1, 1, ()));
1018    /// treap.set(BaseTreapEntry::new(3, 3, ()));
1019    ///
1020    /// assert_eq!(treap.peek_min_mut().unwrap().sorting_key, 1);
1021    /// assert_eq!(treap.peek_min_mut().unwrap().sorting_key, 1);
1022    ///
1023    /// let mut node_mut = treap.peek_min_mut().unwrap();
1024    /// if node_mut.sorting_key < 2 { // conditional remove
1025    ///     node_mut.remove_from_treap();
1026    /// }
1027    /// ```
1028    pub fn peek_min_mut(&mut self) -> Option<NodeMut<'_, E>> {
1029        let mut current = self.root?;
1030
1031        unsafe {
1032            while let Some(left) = (*current.as_ptr()).left {
1033                current = left;
1034            }
1035        }
1036
1037        Some(NodeMut {
1038            node_ptr: current,
1039            treap: self,
1040        })
1041    }
1042
1043    /// Generic filtered search. When `FIND_MAX` is `true`, returns the
1044    /// greatest node whose `filtering_key() >= min_filter`; when `false`,
1045    /// the smallest such node. Prunes branches via the `max_filter` augmentation.
1046    fn find_with_filter<const FIND_MAX: bool>(
1047        &self,
1048        min_filter: &E::FilteringKey,
1049    ) -> Option<NonNull<Node<E>>> {
1050        let root = self.root?;
1051        if unlikely(unsafe { root.as_ref().max_filter < *min_filter }) {
1052            return None;
1053        }
1054
1055        let mut current = root;
1056
1057        loop {
1058            unsafe {
1059                let node_ref = current.as_ref();
1060                let (first_, second_) = if FIND_MAX {
1061                    (node_ref.right, node_ref.left)
1062                } else {
1063                    (node_ref.left, node_ref.right)
1064                };
1065
1066                if let Some(first) = first_ {
1067                    if &(*first.as_ptr()).max_filter >= min_filter {
1068                        current = first;
1069
1070                        continue;
1071                    }
1072                }
1073
1074                if node_ref.entry.filtering_key() >= min_filter {
1075                    return Some(current);
1076                }
1077
1078                if let Some(second) = second_ {
1079                    if cfg!(test) {
1080                        assert!(
1081                            &(*second.as_ptr()).max_filter >= min_filter,
1082                            "the second node max_filter is not >= min_filter"
1083                        );
1084                    }
1085
1086                    current = second;
1087                } else {
1088                    unreachable_hint();
1089                }
1090            }
1091        }
1092    }
1093
1094    /// Finds the maximum node that satisfies `filtering_key() >= min_filter`.
1095    ///
1096    /// Uses subtree augmentation to prune branches.
1097    fn find_max_with_filter(&self, min_filter: &E::FilteringKey) -> Option<NonNull<Node<E>>> {
1098        Self::find_with_filter::<true>(self, min_filter)
1099    }
1100
1101    /// Finds the minimum node that satisfies `filtering_key() >= min_filter`.
1102    ///
1103    /// Uses subtree augmentation to prune branches.
1104    fn find_min_with_filter(&self, min_filter: &E::FilteringKey) -> Option<NonNull<Node<E>>> {
1105        Self::find_with_filter::<false>(self, min_filter)
1106    }
1107
1108    /// Removes and returns the entry with the greatest sorting key that satisfies
1109    /// `filtering_key() >= min_filter`, or `None` if no such entry exists.
1110    ///
1111    /// # Example
1112    ///
1113    /// ```rust
1114    /// use orengine_utils::treap::{Treap, BaseTreapEntry};
1115    ///
1116    /// let mut treap = Treap::<BaseTreapEntry<u32, u32, ()>>::new();
1117    ///
1118    /// treap.set(BaseTreapEntry::new(1, 10, ()));
1119    /// treap.set(BaseTreapEntry::new(2, 5, ()));
1120    /// treap.set(BaseTreapEntry::new(3, 1, ()));
1121    ///
1122    /// assert_eq!(treap.pop_max_with_filter(&5).unwrap().sorting_key, 2);
1123    /// assert_eq!(treap.pop_max_with_filter(&5).unwrap().sorting_key, 1);
1124    /// assert!(treap.pop_max_with_filter(&5).is_none());
1125    /// ```
1126    pub fn pop_max_with_filter(&mut self, min_filter: &E::FilteringKey) -> Option<E> {
1127        let node = self.find_max_with_filter(min_filter)?;
1128
1129        Some(self.remove_node_by_ptr(node))
1130    }
1131
1132    /// Returns [`Node`] with the greatest sorting key satisfying
1133    /// `filtering_key() >= min_filter`, without removing it.
1134    ///
1135    /// # Example
1136    ///
1137    /// ```rust
1138    /// use orengine_utils::treap::{Treap, BaseTreapEntry};
1139    ///
1140    /// let mut treap = Treap::<BaseTreapEntry<u32, u32, ()>>::new();
1141    ///
1142    /// treap.set(BaseTreapEntry::new(1, 10, ()));
1143    /// treap.set(BaseTreapEntry::new(2, 5, ()));
1144    /// treap.set(BaseTreapEntry::new(3, 1, ()));
1145    ///
1146    /// assert_eq!(treap.peek_max_with_filter(&5).unwrap().sorting_key, 2);
1147    /// assert_eq!(treap.peek_max_with_filter(&5).unwrap().sorting_key, 2);
1148    /// ```
1149    pub fn peek_max_with_filter(&self, min_filter: &E::FilteringKey) -> Option<&Node<E>> {
1150        self.find_max_with_filter(min_filter)
1151            .map(|ptr| unsafe { ptr.as_ref() })
1152    }
1153
1154    /// Returns [`NodeMut`] with the greatest sorting key satisfying
1155    /// `filtering_key() >= min_filter`, without removing it.
1156    ///
1157    /// # Example
1158    ///
1159    /// ```rust
1160    /// use orengine_utils::treap::{Treap, BaseTreapEntry};
1161    ///
1162    /// let mut treap = Treap::<BaseTreapEntry<u32, u32, ()>>::new();
1163    ///
1164    /// treap.set(BaseTreapEntry::new(1, 10, ()));
1165    /// treap.set(BaseTreapEntry::new(2, 5, ()));
1166    /// treap.set(BaseTreapEntry::new(3, 1, ()));
1167    ///
1168    /// assert_eq!(treap.peek_max_with_filter_mut(&5).unwrap().sorting_key, 2);
1169    /// assert_eq!(treap.peek_max_with_filter_mut(&5).unwrap().sorting_key, 2);
1170    ///
1171    /// let mut node_mut = treap.peek_max_with_filter_mut(&5).unwrap();
1172    /// if node_mut.sorting_key > 2 { // conditional remove
1173    ///     node_mut.remove_from_treap();
1174    /// }
1175    /// ```
1176    pub fn peek_max_with_filter_mut(
1177        &mut self,
1178        min_filter: &E::FilteringKey,
1179    ) -> Option<NodeMut<'_, E>> {
1180        self.find_max_with_filter(min_filter).map(|ptr| NodeMut {
1181            node_ptr: ptr,
1182            treap: self,
1183        })
1184    }
1185
1186    /// Removes and returns the entry with the smallest sorting key that satisfies
1187    /// `filtering_key() >= min_filter`, or `None` if no such entry exists.
1188    ///
1189    /// # Example
1190    ///
1191    /// ```rust
1192    /// use orengine_utils::treap::{Treap, BaseTreapEntry};
1193    ///
1194    /// let mut treap = Treap::<BaseTreapEntry<u32, u32, ()>>::new();
1195    ///
1196    /// treap.set(BaseTreapEntry::new(1, 10, ()));
1197    /// treap.set(BaseTreapEntry::new(2, 5, ()));
1198    /// treap.set(BaseTreapEntry::new(3, 1, ()));
1199    ///
1200    /// assert_eq!(treap.pop_min_with_filter(&5).unwrap().sorting_key, 1);
1201    /// assert_eq!(treap.pop_min_with_filter(&5).unwrap().sorting_key, 2);
1202    /// assert!(treap.pop_min_with_filter(&5).is_none());
1203    /// ```
1204    pub fn pop_min_with_filter(&mut self, min_filter: &E::FilteringKey) -> Option<E> {
1205        let node = self.find_min_with_filter(min_filter)?;
1206
1207        Some(self.remove_node_by_ptr(node))
1208    }
1209
1210    /// Returns [`Node`] with the smallest sorting key satisfying
1211    /// `filtering_key() >= min_filter`, without removing it.
1212    ///
1213    /// # Example
1214    ///
1215    /// ```rust
1216    /// use orengine_utils::treap::{Treap, BaseTreapEntry};
1217    ///
1218    /// let mut treap = Treap::<BaseTreapEntry<u32, u32, ()>>::new();
1219    ///
1220    /// treap.set(BaseTreapEntry::new(1, 10, ()));
1221    /// treap.set(BaseTreapEntry::new(2, 5, ()));
1222    /// treap.set(BaseTreapEntry::new(3, 1, ()));
1223    ///
1224    /// assert_eq!(treap.peek_min_with_filter(&5).unwrap().sorting_key, 1);
1225    /// assert_eq!(treap.peek_min_with_filter(&5).unwrap().sorting_key, 1);
1226    /// ```
1227    pub fn peek_min_with_filter(&self, min_filter: &E::FilteringKey) -> Option<&Node<E>> {
1228        self.find_min_with_filter(min_filter)
1229            .map(|ptr| unsafe { ptr.as_ref() })
1230    }
1231
1232    /// Returns [`NodeMut`] with the smallest sorting key satisfying
1233    /// `filtering_key() >= min_filter`, without removing it.
1234    ///
1235    /// # Example
1236    ///
1237    /// ```rust
1238    /// use orengine_utils::treap::{Treap, BaseTreapEntry};
1239    ///
1240    /// let mut treap = Treap::<BaseTreapEntry<u32, u32, ()>>::new();
1241    ///
1242    /// treap.set(BaseTreapEntry::new(1, 10, ()));
1243    /// treap.set(BaseTreapEntry::new(2, 5, ()));
1244    /// treap.set(BaseTreapEntry::new(3, 1, ()));
1245    ///
1246    /// assert_eq!(treap.peek_min_with_filter_mut(&5).unwrap().sorting_key, 1);
1247    /// assert_eq!(treap.peek_min_with_filter_mut(&5).unwrap().sorting_key, 1);
1248    ///
1249    /// let mut node_mut = treap.peek_min_with_filter_mut(&5).unwrap();
1250    /// if node_mut.sorting_key < 2 { // conditional remove
1251    ///     node_mut.remove_from_treap();
1252    /// }
1253    /// ```
1254    pub fn peek_min_with_filter_mut(
1255        &mut self,
1256        min_filter: &E::FilteringKey,
1257    ) -> Option<NodeMut<'_, E>> {
1258        self.find_min_with_filter(min_filter).map(|ptr| NodeMut {
1259            node_ptr: ptr,
1260            treap: self,
1261        })
1262    }
1263
1264    /// Returns an in-order iterator over all nodes from the smallest to the greatest sorting key.
1265    ///
1266    /// # Example
1267    ///
1268    /// ```rust
1269    /// use orengine_utils::treap::{Treap, BaseTreapEntry};
1270    ///
1271    /// let mut treap = Treap::<BaseTreapEntry<u32, u32, ()>>::new();
1272    ///
1273    /// treap.set(BaseTreapEntry::new(3, 3, ()));
1274    /// treap.set(BaseTreapEntry::new(1, 1, ()));
1275    /// treap.set(BaseTreapEntry::new(2, 2, ()));
1276    ///
1277    /// let keys: Vec<_> = treap.iter().map(|n| n.sorting_key).collect();
1278    /// assert_eq!(keys, vec![1, 2, 3]);
1279    /// ```
1280    pub fn iter(&self) -> impl Iterator<Item = &Node<E>> {
1281        /// In-order iterator over [`Treap`] nodes, from the smallest to the greatest sorting key.
1282        pub struct Iter<'treap, E: TreapEntry> {
1283            current: Option<NonNull<Node<E>>>,
1284            _marker: core::marker::PhantomData<&'treap Treap<E>>,
1285        }
1286
1287        impl<'treap, E: TreapEntry> Iterator for Iter<'treap, E> {
1288            type Item = &'treap Node<E>;
1289
1290            fn next(&mut self) -> Option<Self::Item> {
1291                let curr = self.current?;
1292
1293                let next_node;
1294
1295                if let Some(node) = unsafe { curr.as_ref() }.right {
1296                    let mut temp = node;
1297
1298                    while let Some(left) = unsafe { temp.as_ref() }.left {
1299                        temp = left;
1300                    }
1301
1302                    next_node = Some(temp);
1303                } else {
1304                    let mut temp = curr;
1305
1306                    loop {
1307                        if let Some(parent) = unsafe { temp.as_ref() }.parent {
1308                            let is_left_child = unsafe { parent.as_ref() }.left == Some(temp);
1309
1310                            if is_left_child {
1311                                next_node = Some(parent);
1312
1313                                break;
1314                            }
1315
1316                            temp = parent;
1317                        } else {
1318                            next_node = None;
1319
1320                            break;
1321                        }
1322                    }
1323                }
1324
1325                let val = unsafe { curr.as_ref() };
1326
1327                self.current = next_node;
1328
1329                Some(val)
1330            }
1331        }
1332
1333        let mut current = self.root;
1334
1335        if let Some(node) = current {
1336            let mut curr = node;
1337
1338            while let Some(left) = unsafe { curr.as_ref() }.left {
1339                curr = left;
1340            }
1341
1342            current = Some(curr);
1343        }
1344
1345        Iter {
1346            current,
1347            _marker: core::marker::PhantomData,
1348        }
1349    }
1350
1351    /// Validate treap structure (for testing).
1352    #[cfg(test)]
1353    pub(crate) fn validate(&self)
1354    where
1355        E::FilteringKey: core::fmt::Debug,
1356    {
1357        if self.root.is_none() {
1358            return;
1359        }
1360
1361        Self::validate_helper(self.root.unwrap(), None, None, None);
1362    }
1363
1364    #[cfg(test)]
1365    fn validate_helper(
1366        node: NonNull<Node<E>>,
1367        min_key: Option<&E::SortingKey>,
1368        max_key: Option<&E::SortingKey>,
1369        parent: Option<NonNull<Node<E>>>,
1370    ) where
1371        E::FilteringKey: core::fmt::Debug,
1372    {
1373        let node_ref = unsafe { node.as_ref() };
1374
1375        assert_eq!(node_ref.parent, parent);
1376
1377        // BST property
1378        if let Some(min) = min_key {
1379            assert!(node_ref.entry.sorting_key() > min);
1380        }
1381        if let Some(max) = max_key {
1382            assert!(node_ref.entry.sorting_key() <= max);
1383        }
1384
1385        // Heap property
1386        if let Some(p) = parent {
1387            assert!(unsafe { p.as_ref() }.priority >= node_ref.priority);
1388        }
1389
1390        // Augmentation
1391        let mut expected_max = node_ref.entry.filtering_key();
1392
1393        if let Some(left) = node_ref.left {
1394            Self::validate_helper(
1395                left,
1396                min_key,
1397                Some(node_ref.entry.sorting_key()),
1398                Some(node),
1399            );
1400
1401            let left_ref = unsafe { left.as_ref() };
1402            if &left_ref.max_filter > expected_max {
1403                expected_max = &left_ref.max_filter;
1404            }
1405        }
1406
1407        if let Some(right) = node_ref.right {
1408            Self::validate_helper(
1409                right,
1410                Some(node_ref.entry.sorting_key()),
1411                max_key,
1412                Some(node),
1413            );
1414
1415            let right_ref = unsafe { right.as_ref() };
1416            if &right_ref.max_filter > expected_max {
1417                expected_max = &right_ref.max_filter;
1418            }
1419        }
1420
1421        assert_eq!(expected_max, &node_ref.max_filter);
1422    }
1423}
1424
1425macro_rules! generate_find_ptr {
1426    (
1427        $name:ident,
1428        $self_type:ty,
1429        $self_name:ident,
1430        $current_ptr_name:ident,
1431        $get_current_block:block
1432    ) => {
1433        fn $name($self_name: $self_type, key: &E::SortingKey) -> Option<NonNull<Node<E>>> {
1434            let mut $current_ptr_name = $self_name.root?;
1435
1436            unsafe {
1437                loop {
1438                    let current_ref = $get_current_block;
1439
1440                    match current_ref.entry.sorting_key().cmp(key) {
1441                        Ordering::Equal => return Some(NonNull::from(current_ref)),
1442                        Ordering::Greater => {
1443                            $current_ptr_name = current_ref.left?;
1444                        }
1445                        Ordering::Less => {
1446                            $current_ptr_name = current_ref.right?;
1447                        }
1448                    }
1449                }
1450            }
1451        }
1452    };
1453}
1454
1455impl<E: TreapEntry> Treap<E> {
1456    generate_find_ptr!(find_ptr_, &Self, self_, current, { current.as_ref() });
1457    generate_find_ptr!(find_ptr_mut_, &mut Self, self_, current, {
1458        current.as_mut()
1459    });
1460
1461    /// Returns a raw shared pointer to the node with the given sorting key,
1462    /// or `None` if absent.
1463    ///
1464    /// The pointer is valid until the node is removed or the treap is dropped.
1465    fn find_ptr(&self, key: &E::SortingKey) -> Option<NonNull<Node<E>>> {
1466        Self::find_ptr_(self, key)
1467    }
1468
1469    /// Returns a raw mutable pointer to the node with the given sorting key,
1470    /// or `None` if absent.
1471    ///
1472    /// See [`find_ptr`](Treap::find_ptr) for a usage example.
1473    fn find_ptr_mut(&mut self, key: &E::SortingKey) -> Option<NonNull<Node<E>>> {
1474        Self::find_ptr_mut_(self, key)
1475    }
1476
1477    /// Returns [`Node`] with the given sorting key, or `None` if absent.
1478    ///
1479    /// # Example
1480    ///
1481    /// ```rust
1482    /// use orengine_utils::treap::{BaseTreapEntry, Treap};
1483    ///
1484    /// let mut treap = Treap::<BaseTreapEntry<u32, u32, &str>>::new();
1485    ///
1486    /// treap.set(BaseTreapEntry::new(1, 1, "hello"));
1487    ///
1488    /// assert_eq!(treap.find(&1).unwrap().value, "hello");
1489    /// ```
1490    pub fn find(&self, key: &E::SortingKey) -> Option<&Node<E>> {
1491        let ptr = self.find_ptr(key)?;
1492
1493        Some(unsafe { ptr.as_ref() })
1494    }
1495
1496    /// Returns [`NodeMut`] with the given sorting key, or `None` if absent.
1497    ///
1498    /// See [`find`](Treap::find) for a usage example.
1499    pub fn find_mut(&mut self, key: &E::SortingKey) -> Option<NodeMut<'_, E>> {
1500        let ptr = self.find_ptr_mut(key)?;
1501
1502        Some(NodeMut {
1503            treap: self,
1504            node_ptr: ptr,
1505        })
1506    }
1507}
1508
1509impl<E: TreapEntry> Default for Treap<E> {
1510    fn default() -> Self {
1511        Self::new()
1512    }
1513}
1514
1515impl<E: TreapEntry + Display> Display for Treap<E> {
1516    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1517        fn print_node<E: TreapEntry + Display>(
1518            f: &mut fmt::Formatter<'_>,
1519            node: NonNull<Node<E>>,
1520            prefix: &str,
1521            is_last: bool,
1522            is_left: bool,
1523        ) -> fmt::Result {
1524            let n = unsafe { node.as_ref() };
1525
1526            // 1. Print the connector
1527            // If it's the last child, use └──, else use ├──
1528            let connector = if is_last { "└── " } else { "├── " };
1529            let order_msg = if is_left { "Left" } else { "Right" };
1530            write!(f, "{prefix}{connector}{order_msg} ")?;
1531
1532            // 2. Print the node value
1533            writeln!(f, "{}", n.entry)?;
1534
1535            // 3. Prepare prefix for children
1536            // If this node was the last child, the extension for children is spaces (don't draw line).
1537            // If it wasn't the last child, the extension is a vertical bar (to connect to the sibling below).
1538            let child_prefix = if is_last {
1539                format!("{prefix}    ") // 4 spaces
1540            } else {
1541                format!("{prefix}│   ") // Pipe + 3 spaces
1542            };
1543
1544            // 4. Recurse for children
1545            let has_right = n.right.is_some();
1546
1547            // Print Left Child
1548            // It is the "last" child visible at this level if there is NO right child.
1549            if let Some(left) = &n.left {
1550                print_node(f, *left, &child_prefix, !has_right, true)?;
1551            }
1552
1553            // Print Right Child
1554            // It is always the "last" child if it exists.
1555            if let Some(right) = &n.right {
1556                print_node(f, *right, &child_prefix, true, false)?;
1557            }
1558
1559            Ok(())
1560        }
1561
1562        match &self.root {
1563            Some(node) => {
1564                // We treat the root specially: it has no prefix and no "is_last" status
1565                // in the traditional sense, but we start the recursion here.
1566                writeln!(f, "{}", unsafe { node.as_ref() }.entry)?;
1567
1568                // Print children. We need to know if a child is the "last" one
1569                // to draw the correct branch shape (└── vs. ├──).
1570                let rb = unsafe { node.as_ref() };
1571
1572                // Process Right Child (visually "top" branch in rotated view, or second in text)
1573                // Typically we print Left then Right.
1574                // In a text diagram, often Right is printed first to keep it "upright"
1575                // or Left first to follow reading order.
1576                // Let's stick to the standard Left-then-Right reading order for the diagram.
1577
1578                let has_right = rb.right.is_some();
1579
1580                if let Some(left) = &rb.left {
1581                    print_node(f, *left, "", !has_right, true)?;
1582                }
1583                if let Some(right) = &rb.right {
1584                    print_node(f, *right, "", true, false)?;
1585                }
1586
1587                Ok(())
1588            }
1589            None => write!(f, "Empty Tree"),
1590        }
1591    }
1592}
1593
1594impl<E: TreapEntry> Drop for Treap<E> {
1595    fn drop(&mut self) {
1596        fn drop_subtree<E: TreapEntry>(mut current: Option<NonNull<Node<E>>>) {
1597            // Stack memory is almost free to allocate, so we can allocate
1598            // 256 * 8 = 2KB of stack memory and do not care about the performance.
1599            // But the treap can become a linked-list if we are extremely unlucky,
1600            // so we need to be careful about the stack overflow.
1601            // We handle it below.
1602            let mut stack = ArrayBuffer::<_, 256>::new();
1603            let mut last_freed: Option<NonNull<Node<E>>> = None;
1604
1605            while current.is_some() || !stack.is_empty() {
1606                // Dive as far left as possible
1607                while let Some(n) = current {
1608                    let res = stack.push(n);
1609                    if let Err(left) = res {
1610                        // Drop the subtree in a new function with a new stack
1611
1612                        cold_path();
1613
1614                        drop_subtree(Some(left));
1615                    } else {
1616                        current = unsafe { n.as_ref() }.left;
1617                    }
1618                }
1619
1620                let &top = stack.last().unwrap();
1621                let node_ref = unsafe { top.as_ref() };
1622
1623                // If there's an unprocessed right child, go there
1624                if node_ref.right.is_some() && node_ref.right != last_freed {
1625                    current = node_ref.right;
1626                } else {
1627                    // Both children are done, free this node
1628
1629                    stack.pop();
1630
1631                    last_freed = Some(top);
1632
1633                    let _ = unsafe { Box::from_raw(top.as_ptr()) };
1634                }
1635            }
1636        }
1637
1638        drop_subtree(self.root);
1639    }
1640}
1641
1642/// A ready-to-use [`TreapEntry`] implementation wrapping a sorting key,
1643/// filtering key, and a value.
1644///
1645/// Use this when you don't need a custom entry type.
1646///
1647/// # Example
1648///
1649/// ```rust
1650/// use orengine_utils::treap::{BaseTreapEntry, Treap};
1651///
1652/// let mut treap = Treap::<BaseTreapEntry<u32, u32, &str>>::new();
1653///
1654/// treap.set(BaseTreapEntry::new(1, 10, "hello"));
1655///
1656/// assert_eq!(treap.find(&1).unwrap().value, "hello");
1657/// ```
1658pub struct BaseTreapEntry<SortingKey: Ord, FilteringKey: Ord + Clone, V> {
1659    pub sorting_key: SortingKey,
1660    pub filtering_key: FilteringKey,
1661    pub value: V,
1662}
1663
1664impl<SortingKey: Ord, FilteringKey: Ord + Clone, V> BaseTreapEntry<SortingKey, FilteringKey, V> {
1665    /// Creates a new [`BaseTreapEntry`] with the given sorting key, filtering key, and value.
1666    ///
1667    /// # Example
1668    ///
1669    /// ```rust
1670    /// use orengine_utils::treap::BaseTreapEntry;
1671    ///
1672    /// let entry = BaseTreapEntry::new(42u32, 7u32, "payload");
1673    ///
1674    /// assert_eq!(entry.sorting_key, 42);
1675    /// assert_eq!(entry.filtering_key, 7);
1676    /// assert_eq!(entry.value, "payload");
1677    /// ```
1678    #[inline]
1679    pub fn new(sorting_key: SortingKey, filtering_key: FilteringKey, value: V) -> Self {
1680        Self {
1681            sorting_key,
1682            filtering_key,
1683            value,
1684        }
1685    }
1686}
1687
1688impl<SortingKey: Ord, FilteringKey: Ord + Clone, V> TreapEntry
1689    for BaseTreapEntry<SortingKey, FilteringKey, V>
1690{
1691    type SortingKey = SortingKey;
1692    type FilteringKey = FilteringKey;
1693    type Value = V;
1694
1695    #[inline]
1696    fn sorting_key(&self) -> &Self::SortingKey {
1697        &self.sorting_key
1698    }
1699
1700    #[inline]
1701    fn filtering_key(&self) -> &Self::FilteringKey {
1702        &self.filtering_key
1703    }
1704
1705    #[inline]
1706    fn value(&self) -> &Self::Value {
1707        &self.value
1708    }
1709
1710    #[inline]
1711    fn value_mut(&mut self) -> &mut Self::Value {
1712        &mut self.value
1713    }
1714}
1715
1716#[cfg(test)]
1717mod tests {
1718    use super::*;
1719    use crate::alloc::string::ToString;
1720    use alloc::string::String;
1721    use alloc::vec::Vec;
1722    use core::iter::from_fn;
1723
1724    #[derive(Debug, Clone, Eq, PartialEq)]
1725    struct TestEntry {
1726        primary: i32,
1727        filter: i32,
1728        value: String,
1729    }
1730
1731    fn generate_filtering_key(primary: i32) -> i32 {
1732        primary % 50
1733    }
1734
1735    impl TestEntry {
1736        fn new(primary: i32, filter: i32, value: &str) -> Self {
1737            Self {
1738                primary,
1739                filter,
1740                value: value.to_string(),
1741            }
1742        }
1743    }
1744
1745    impl TreapEntry for TestEntry {
1746        type SortingKey = i32;
1747        type FilteringKey = i32;
1748        type Value = String;
1749
1750        fn sorting_key(&self) -> &Self::SortingKey {
1751            &self.primary
1752        }
1753
1754        fn filtering_key(&self) -> &Self::FilteringKey {
1755            &self.filter
1756        }
1757
1758        fn value(&self) -> &Self::Value {
1759            &self.value
1760        }
1761
1762        fn value_mut(&mut self) -> &mut Self::Value {
1763            &mut self.value
1764        }
1765    }
1766
1767    impl Display for TestEntry {
1768        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1769            write!(f, "({}, {})", self.primary, self.filter)
1770        }
1771    }
1772
1773    #[test]
1774    #[cfg(not(feature = "no_std"))]
1775    fn test_as_tree() {
1776        let mut treap = Treap::new();
1777
1778        for i in 0..100 {
1779            unsafe {
1780                treap.add(TestEntry::new(
1781                    i,
1782                    generate_filtering_key(i),
1783                    &format!("item_{i}"),
1784                ));
1785            };
1786
1787            let item = treap.find(&i).unwrap();
1788
1789            assert_eq!(item.primary, i);
1790            assert_eq!(item.filter, generate_filtering_key(i));
1791            assert_eq!(&item.value, &format!("item_{i}"));
1792        }
1793
1794        for i in 0..100 {
1795            let item = treap.find(&i).unwrap();
1796
1797            assert_eq!(item.primary, i);
1798            assert_eq!(item.filter, generate_filtering_key(i));
1799            assert_eq!(&item.value, &format!("item_{i}"));
1800        }
1801
1802        for i in 0..100 {
1803            if i % 2 == 0 {
1804                let delta = if i % 4 == 0 { 10 } else { -10 };
1805                treap.set(TestEntry::new(
1806                    i,
1807                    generate_filtering_key(i + delta),
1808                    &format!("updated_item_{i}"),
1809                ));
1810
1811                let item = treap.find(&i).unwrap();
1812
1813                assert_eq!(item.primary, i);
1814                assert_eq!(item.filter, generate_filtering_key(i + delta));
1815                assert_eq!(&item.value, &format!("updated_item_{i}"));
1816            } else {
1817                let entry = treap.remove_by_sorting_key(&i).unwrap();
1818
1819                assert_eq!(entry.primary, i);
1820                assert_eq!(entry.filter, generate_filtering_key(i));
1821                assert_eq!(&entry.value, &format!("item_{i}"));
1822            }
1823
1824            treap.validate();
1825        }
1826
1827        println!("{treap}");
1828    }
1829
1830    #[test]
1831    fn test_pop_best() {
1832        let mut treap: Treap<TestEntry> = Treap::new();
1833
1834        treap.set(TestEntry::new(3, 30, "three"));
1835        treap.set(TestEntry::new(1, 10, "one"));
1836        treap.set(TestEntry::new(2, 20, "two"));
1837
1838        treap.validate();
1839
1840        let e = treap.peek_min().unwrap();
1841        assert_eq!(e.primary, 1);
1842        assert_eq!(e.value, "one");
1843
1844        treap.validate();
1845
1846        let e = treap.pop_min().unwrap();
1847        assert_eq!(e.primary, 1);
1848        assert_eq!(e.value, "one");
1849
1850        treap.validate();
1851
1852        let e = treap.peek_min().unwrap();
1853        assert_eq!(e.primary, 2);
1854        assert_eq!(e.value, "two");
1855
1856        treap.validate();
1857
1858        let e = treap.pop_min().unwrap();
1859        assert_eq!(e.primary, 2);
1860
1861        treap.validate();
1862
1863        let e = treap.peek_min().unwrap();
1864        assert_eq!(e.primary, 3);
1865
1866        treap.validate();
1867
1868        let e = treap.pop_min().unwrap();
1869        assert_eq!(e.primary, 3);
1870
1871        treap.validate();
1872
1873        treap.set(TestEntry::new(3, 30, "three"));
1874        treap.set(TestEntry::new(1, 10, "one"));
1875        treap.set(TestEntry::new(2, 20, "two"));
1876
1877        let e = treap.pop_max().unwrap();
1878        assert_eq!(e.primary, 3);
1879        assert_eq!(e.value, "three");
1880
1881        treap.validate();
1882
1883        let e = treap.pop_max().unwrap();
1884        assert_eq!(e.primary, 2);
1885
1886        treap.validate();
1887
1888        let e = treap.pop_max().unwrap();
1889        assert_eq!(e.primary, 1);
1890
1891        treap.validate();
1892
1893        assert!(treap.pop_max().is_none());
1894    }
1895
1896    #[test]
1897    fn test_remove_by_pointer() {
1898        let mut treap: Treap<TestEntry> = Treap::new();
1899        let mut to_remove = Vec::new();
1900
1901        for i in 0..100 {
1902            if i % 2 == 0 {
1903                to_remove.push(
1904                    *treap
1905                        .set(TestEntry::new(i, i * 10, &format!("even_{i}")))
1906                        .0
1907                        .sorting_key(),
1908                );
1909            } else {
1910                unsafe { treap.add(TestEntry::new(i, i * 10, &format!("odd_{i}"))) };
1911            }
1912        }
1913
1914        assert_eq!(treap.iter().count(), 100);
1915
1916        for sorting_key in to_remove {
1917            let e = if sorting_key % 2 == 0 {
1918                treap.remove_by_sorting_key(&sorting_key).unwrap()
1919            } else {
1920                treap.find_mut(&sorting_key).unwrap().remove_from_treap()
1921            };
1922
1923            assert!(e.value.starts_with("even_"));
1924
1925            treap.validate();
1926        }
1927
1928        assert_eq!(treap.iter().count(), 50);
1929
1930        treap.validate();
1931    }
1932
1933    #[test]
1934    fn test_filtering_max() {
1935        let mut treap: Treap<TestEntry> = Treap::new();
1936
1937        treap.set(TestEntry::new(1, 20, "a"));
1938        treap.set(TestEntry::new(2, 18, "b"));
1939        treap.set(TestEntry::new(3, 15, "c"));
1940        treap.set(TestEntry::new(4, 13, "d"));
1941        treap.set(TestEntry::new(5, 10, "e"));
1942
1943        assert_eq!(treap.peek_max_with_filter(&15).unwrap().primary, 3);
1944        assert_eq!(treap.peek_max_with_filter(&20).unwrap().primary, 1);
1945        assert_eq!(treap.peek_max_with_filter(&16).unwrap().primary, 2);
1946        assert_eq!(treap.peek_max_with_filter(&1).unwrap().primary, 5);
1947        assert_eq!(treap.peek_max_with_filter(&11).unwrap().primary, 4);
1948
1949        assert!(treap.peek_max_with_filter(&21).is_none());
1950
1951        assert_eq!(
1952            &from_fn(|| treap.pop_max_with_filter(&15).map(|e| e.primary)).collect::<Vec<_>>(),
1953            &[3, 2, 1]
1954        );
1955
1956        assert_eq!(treap.iter().count(), 2);
1957
1958        treap.validate();
1959    }
1960
1961    #[test]
1962    fn test_filtering_min() {
1963        let mut treap: Treap<TestEntry> = Treap::new();
1964
1965        treap.set(TestEntry::new(1, 10, "a"));
1966        treap.set(TestEntry::new(2, 13, "b"));
1967        treap.set(TestEntry::new(3, 15, "c"));
1968        treap.set(TestEntry::new(4, 18, "d"));
1969        treap.set(TestEntry::new(5, 20, "e"));
1970
1971        assert_eq!(treap.peek_min_with_filter(&15).unwrap().primary, 3);
1972        assert_eq!(treap.peek_min_with_filter(&20).unwrap().primary, 5);
1973        assert_eq!(treap.peek_min_with_filter(&16).unwrap().primary, 4);
1974        assert_eq!(treap.peek_min_with_filter(&1).unwrap().primary, 1);
1975        assert_eq!(treap.peek_min_with_filter(&12).unwrap().primary, 2);
1976
1977        assert!(treap.peek_min_with_filter(&21).is_none());
1978
1979        assert_eq!(
1980            &from_fn(|| treap.pop_min_with_filter(&15).map(|e| e.primary)).collect::<Vec<_>>(),
1981            &[3, 4, 5]
1982        );
1983
1984        assert_eq!(treap.iter().count(), 2);
1985
1986        treap.validate();
1987    }
1988
1989    #[test]
1990    #[cfg(not(feature = "no_std"))]
1991    fn test_neighbors() {
1992        let mut treap: Treap<TestEntry> = Treap::new();
1993
1994        // region filling
1995
1996        treap.set(TestEntry::new(1, 1, "a"));
1997        treap.set(TestEntry::new(3, 2, "c"));
1998        treap.set(TestEntry::new(2, 10, "b"));
1999        treap.set(TestEntry::new(4, 4, "d"));
2000        treap.set(TestEntry::new(5, 5, "e"));
2001        treap.set(TestEntry::new(6, 6, "e"));
2002        treap.set(TestEntry::new(7, 7, "e"));
2003        treap.set(TestEntry::new(8, 8, "e"));
2004        treap.set(TestEntry::new(9, 9, "e"));
2005        treap.set(TestEntry::new(10, 3, "e"));
2006        treap.set(TestEntry::new(11, 3, "e"));
2007        treap.set(TestEntry::new(12, 0, "e"));
2008        treap.set(TestEntry::new(13, 0, "e"));
2009        treap.set(TestEntry::new(14, 1, "e"));
2010        let node15_key = *treap.set(TestEntry::new(15, 3, "e")).0.sorting_key();
2011        treap.set(TestEntry::new(16, 100, "e"));
2012        treap.set(TestEntry::new(17, 1, "e"));
2013        treap.set(TestEntry::new(18, 1, "e"));
2014        treap.set(TestEntry::new(19, 1, "e"));
2015
2016        // endregion
2017
2018        println!("Treap now: \n{treap}");
2019
2020        // Treap now:
2021        // (17, 1)
2022        // ├── Left (10, 3)
2023        // │   ├── Left (6, 6)
2024        // │   │   ├── Left (3, 2)
2025        // │   │   │   ├── Left (1, 1)
2026        // │   │   │   │   └── Right (2, 10)
2027        // │   │   │   └── Right (5, 5)
2028        // │   │   │       └── Left (4, 4)
2029        // │   │   └── Right (9, 9)
2030        // │   │       └── Left (7, 7)
2031        // │   │           └── Right (8, 8)
2032        // │   └── Right (15, 3)
2033        // │       ├── Left (13, 0)
2034        // │       │   ├── Left (12, 0)
2035        // │       │   │   └── Left (11, 3)
2036        // │       │   └── Right (14, 1)
2037        // │       └── Right (16, 100)
2038        // └── Right (18, 1)
2039        //     └── Right (19, 1)
2040
2041        let node15 = treap.find(&node15_key).unwrap();
2042
2043        let neighbors: Vec<_> = node15
2044            .neighbors(&4, false)
2045            .map(|n| n.entry.primary)
2046            .collect();
2047
2048        assert_eq!(neighbors, alloc::vec![16, 6, 9, 7, 8, 5, 4, 2]);
2049
2050        let neighbors: Vec<_> = node15
2051            .neighbors(&4, true)
2052            .map(|n| n.entry.primary)
2053            .collect();
2054
2055        assert_eq!(neighbors, alloc::vec![6, 9, 7, 8, 5, 4, 2]);
2056    }
2057
2058    #[test]
2059    fn many_items() {
2060        const N: usize = if !cfg!(miri) { 3000 } else { 100 };
2061
2062        let mut state = NonZeroU32::new(1).unwrap();
2063        let mut random = || {
2064            cheap_random_with_current_u32(&mut state);
2065
2066            #[allow(clippy::cast_possible_wrap, reason = "It is fine here")]
2067            {
2068                state.get() as i32
2069            }
2070        };
2071        let mut treap = Treap::new();
2072        let mut inserted_nodes_keys = Vec::with_capacity(N);
2073
2074        for _ in 0..N {
2075            inserted_nodes_keys.push(
2076                *treap
2077                    .set(TestEntry::new(
2078                        random(),
2079                        generate_filtering_key(random()),
2080                        "e",
2081                    ))
2082                    .0
2083                    .sorting_key(),
2084            );
2085
2086            treap.validate();
2087
2088            if random() % 5 == 0 {
2089                #[allow(clippy::cast_sign_loss, reason = "It is fine here")]
2090                treap
2091                    .find_mut(
2092                        &inserted_nodes_keys.remove(random() as usize % inserted_nodes_keys.len()),
2093                    )
2094                    .unwrap()
2095                    .remove_from_treap();
2096
2097                treap.validate();
2098            }
2099        }
2100    }
2101
2102    #[test]
2103    #[cfg(not(feature = "no_std"))]
2104    fn insert_and_remove_rand() {
2105        const N: usize = if !cfg!(miri) { 2000 } else { 20 };
2106
2107        let mut state =
2108            NonZeroU32::new((crate::instant::OrengineInstant::now().into_u64() % 1000) as u32 + 1)
2109                .unwrap();
2110
2111        for _ in 0..10 {
2112            let mut key_value_pairs = std::collections::HashMap::with_capacity(N);
2113            let mut tree = Treap::new();
2114            #[allow(clippy::cast_possible_truncation, reason = "False positive.")]
2115            #[allow(clippy::cast_possible_wrap, reason = "False positive.")]
2116            let mut rand_i32 =
2117                || cheap_random_with_current_u32(&mut state).cast_signed() % (N as i32 * 9 / 10);
2118
2119            for i in 0..N {
2120                let key = rand_i32();
2121                let value_fn = || format!("{i}");
2122
2123                if let Some(old_value) = key_value_pairs.insert(key, value_fn()) {
2124                    assert_eq!(
2125                        tree.find(&key).map(|v: &Node<TestEntry>| &v.value),
2126                        Some(&old_value)
2127                    );
2128
2129                    let old_from_tree = tree.set(TestEntry::new(key, rand_i32(), &value_fn())).1;
2130
2131                    assert!(old_from_tree.is_some());
2132                    assert_eq!(old_from_tree.unwrap().value, old_value);
2133                } else {
2134                    let r = rand_i32();
2135                    if r % 2 == 0 {
2136                        unsafe { tree.add(TestEntry::new(key, rand_i32(), &value_fn())) };
2137                    } else {
2138                        tree.set(TestEntry::new(key, rand_i32(), &value_fn()));
2139                    }
2140                }
2141
2142                tree.validate();
2143
2144                if rand_i32() % 5 == 0 {
2145                    let mut slice = key_value_pairs.iter().take(5).map(|(k, _v)| *k);
2146                    let mut wait = rand_i32() % 5;
2147
2148                    while wait > 0 {
2149                        if slice.next().is_none() {
2150                            break;
2151                        }
2152
2153                        wait -= 1;
2154                    }
2155
2156                    if let Some(key) = slice.next() {
2157                        let value = key_value_pairs.remove(&key).unwrap();
2158                        let value_from_tree = tree.remove_by_sorting_key(&key);
2159
2160                        assert!(value_from_tree.is_some());
2161                        assert_eq!(value_from_tree.unwrap().value, value);
2162                    }
2163
2164                    tree.validate();
2165                }
2166            }
2167        }
2168    }
2169
2170    #[test]
2171    #[cfg(not(feature = "no_std"))]
2172    fn test_drop() {
2173        use core::cell::Cell;
2174
2175        thread_local! {
2176            static SORTING_KEY_DROP_COUNTER: Cell<usize> = const { Cell::new(0) };
2177            static VALUE_DROP_COUNTER: Cell<usize> = const { Cell::new(0) };
2178        }
2179
2180        #[derive(Eq, PartialEq, Ord, PartialOrd)]
2181        struct SortingKeyWrapper(usize);
2182
2183        impl Drop for SortingKeyWrapper {
2184            fn drop(&mut self) {
2185                SORTING_KEY_DROP_COUNTER.replace(SORTING_KEY_DROP_COUNTER.get() + 1);
2186            }
2187        }
2188
2189        #[allow(dead_code, reason = "This value is used for debugging")]
2190        struct ValueWrapper(usize);
2191
2192        impl Drop for ValueWrapper {
2193            fn drop(&mut self) {
2194                VALUE_DROP_COUNTER.replace(VALUE_DROP_COUNTER.get() + 1);
2195            }
2196        }
2197
2198        let mut treap = Treap::<BaseTreapEntry<SortingKeyWrapper, usize, ValueWrapper>>::new();
2199
2200        for i in 0..10 {
2201            let entry = BaseTreapEntry {
2202                sorting_key: SortingKeyWrapper(i),
2203                filtering_key: i,
2204                value: ValueWrapper(i),
2205            };
2206
2207            if i % 2 == 0 {
2208                unsafe { treap.add(entry) };
2209            } else {
2210                treap.set(entry);
2211            }
2212        }
2213
2214        assert_eq!(SORTING_KEY_DROP_COUNTER.get(), 0);
2215        assert_eq!(VALUE_DROP_COUNTER.get(), 0);
2216
2217        for i in 0..10 {
2218            let entry = BaseTreapEntry {
2219                sorting_key: SortingKeyWrapper(i),
2220                filtering_key: i * 2,
2221                value: ValueWrapper(i * 2),
2222            };
2223
2224            treap.set(entry);
2225        }
2226
2227        assert_eq!(SORTING_KEY_DROP_COUNTER.get(), 10);
2228        assert_eq!(VALUE_DROP_COUNTER.get(), 10);
2229
2230        drop(treap);
2231
2232        assert_eq!(SORTING_KEY_DROP_COUNTER.get(), 20);
2233        assert_eq!(VALUE_DROP_COUNTER.get(), 20);
2234    }
2235}