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() > ¤t.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}