augmented_rbtree/
node.rs

1use core::ptr::NonNull;
2/// The color of a node in the red-black tree.
3#[derive(Debug, Clone, Copy, PartialEq, Eq)]
4pub enum Color {
5    /// A red node.
6    Red,
7    /// A black node.
8    Black,
9}
10
11/// Tracks which side of the parent a nil node is on.
12/// This is necessary for `delete_fixup` when x is None.
13#[derive(Debug, Clone, Copy, PartialEq, Eq)]
14pub(crate) enum NilSide {
15    Left,
16    Right,
17}
18
19#[derive(Clone)]
20pub(crate) struct Node<K, V, S> {
21    pub(crate) left: Option<NonNull<Node<K, V, S>>>,
22    pub(crate) right: Option<NonNull<Node<K, V, S>>>,
23    pub(crate) parent: Option<NonNull<Node<K, V, S>>>,
24    pub(crate) key: K,
25    pub(crate) value: V,
26    pub(crate) stats: S,
27    pub(crate) color: Color,
28    pub(crate) _marker: core::marker::PhantomData<(K, V, S)>,
29}
30
31pub mod internal_details {
32    use core::ptr::NonNull;
33
34    use crate::{Color, node::Node};
35
36    #[derive(Debug)]
37    pub struct NodeRef<K, V, S> {
38        pub(crate) ptr: NonNull<Node<K, V, S>>,
39    }
40
41    impl<K, V, S> Clone for NodeRef<K, V, S> {
42        #[inline]
43        fn clone(&self) -> Self {
44            *self
45        }
46    }
47
48    impl<K, V, S> Copy for NodeRef<K, V, S> {}
49
50    impl<K, V, S> PartialEq for NodeRef<K, V, S> {
51        #[inline]
52        fn eq(&self, other: &Self) -> bool {
53            // Two handles are equal if they point to the exact same heap node block
54            self.ptr == other.ptr
55        }
56    }
57
58    impl<K, V, S> Eq for NodeRef<K, V, S> {}
59
60    impl<K, V, S> NodeRef<K, V, S> {
61        #[inline]
62        pub(crate) unsafe fn from_raw(ptr: NonNull<Node<K, V, S>>) -> Self {
63            Self { ptr }
64        }
65
66        #[inline]
67        pub(crate) fn parent(self) -> Option<Self> {
68            unsafe { (*self.ptr.as_ptr()).parent.map(|p| NodeRef { ptr: p }) }
69        }
70
71        #[inline]
72        pub(crate) fn set_parent(self, parent: Option<Self>) {
73            unsafe { (*self.ptr.as_ptr()).parent = parent.map(|p| p.ptr) }
74        }
75
76        #[inline]
77        pub(crate) fn left(self) -> Option<Self> {
78            unsafe { (*self.ptr.as_ptr()).left.map(|p| NodeRef { ptr: p }) }
79        }
80
81        #[inline]
82        pub(crate) fn set_left(self, left: Option<Self>) {
83            unsafe { (*self.ptr.as_ptr()).left = left.map(|l| l.ptr) }
84        }
85
86        #[inline]
87        pub(crate) fn right(self) -> Option<Self> {
88            unsafe { (*self.ptr.as_ptr()).right.map(|p| NodeRef { ptr: p }) }
89        }
90
91        #[inline]
92        pub(crate) fn set_right(self, right: Option<Self>) {
93            unsafe { (*self.ptr.as_ptr()).right = right.map(|r| r.ptr) }
94        }
95
96        #[inline]
97        pub(crate) unsafe fn key<'a>(self) -> &'a K {
98            unsafe { &(*self.ptr.as_ptr()).key }
99        }
100
101        #[inline]
102        pub(crate) unsafe fn value<'a>(self) -> &'a V {
103            unsafe { &(*self.ptr.as_ptr()).value }
104        }
105
106        #[inline]
107        pub(crate) unsafe fn stats<'a>(self) -> &'a S {
108            unsafe { &(*self.ptr.as_ptr()).stats }
109        }
110
111        #[inline]
112        pub(crate) fn color(self) -> Color {
113            unsafe { (*self.ptr.as_ptr()).color }
114        }
115
116        #[inline]
117        pub(crate) fn set_color(self, color: Color) {
118            unsafe { (*self.ptr.as_ptr()).color = color }
119        }
120
121        #[inline]
122        pub(crate) unsafe fn value_mut<'a>(self) -> &'a mut V {
123            unsafe { &mut (*self.ptr.as_ptr()).value }
124        }
125
126        #[inline]
127        pub(crate) fn is_black(node: Option<Self>) -> bool {
128            match node {
129                Some(ptr) => ptr.color() == Color::Black,
130                None => true,
131            }
132        }
133
134        #[inline]
135        #[allow(dead_code)]
136        pub(crate) fn next_node(self) -> Option<Self> {
137            if let Some(right) = self.right() {
138                let mut current = right;
139                while let Some(left) = current.left() {
140                    current = left;
141                }
142                return Some(current);
143            }
144
145            let mut current = self;
146            while let Some(parent) = current.parent() {
147                if parent.left() == Some(current) {
148                    return Some(parent);
149                }
150                current = parent;
151            }
152
153            None
154        }
155
156        #[allow(dead_code)]
157        #[inline]
158        pub(crate) fn prev_node(self) -> Option<Self> {
159            if let Some(left) = self.left() {
160                let mut current = left;
161                while let Some(right) = current.right() {
162                    current = right;
163                }
164                return Some(current);
165            }
166
167            let mut current = self;
168            while let Some(parent) = current.parent() {
169                if parent.right() == Some(current) {
170                    return Some(parent);
171                }
172                current = parent;
173            }
174
175            None
176        }
177
178        #[inline]
179        pub(crate) fn leftmost(self) -> Self {
180            let mut node = self;
181            while let Some(left) = node.left() {
182                node = left;
183            }
184            node
185        }
186
187        #[allow(dead_code)]
188        #[inline]
189        pub(crate) fn rightmost(self) -> Self {
190            let mut node = self;
191            while let Some(right) = node.right() {
192                node = right;
193            }
194            node
195        }
196    }
197}
augmented_rbtree/node.rs

augmented_rbtree/
node.rs
