ngx/collections/
rbtree.rs

1//! Types and utilities for working with [ngx_rbtree_t].
2//!
3//! This module provides both the tools for interaction with the existing `ngx_rbtree_t` objects in
4//! the NGINX, and useful high-level types built on top of the `ngx_rbtree_t`.
5//!
6//! See <https://nginx.org/en/docs/dev/development_guide.html#red_black_tree>.
7
8use core::alloc::Layout;
9use core::cmp::Ordering;
10use core::hash::{self, BuildHasher, Hash};
11use core::marker::PhantomData;
12use core::ptr::{self, NonNull};
13use core::{borrow, mem};
14
15use nginx_sys::{
16    ngx_rbt_red, ngx_rbtree_data, ngx_rbtree_delete, ngx_rbtree_init, ngx_rbtree_insert,
17    ngx_rbtree_key_t, ngx_rbtree_min, ngx_rbtree_next, ngx_rbtree_node_t, ngx_rbtree_t,
18};
19
20use crate::allocator::{self, AllocError, Allocator};
21
22/// Raw iterator over the `ngx_rbtree_t` nodes.
23///
24/// This iterator type can be used to access elements of any correctly initialized `ngx_rbtree_t`
25/// instance, including those already embedded in the nginx structures.  The iterator stores pointer
26/// to the next node and thus remains valid and usable even if the last returned item is removed
27/// from the tree.
28pub struct NgxRbTreeIter<'a> {
29    tree: NonNull<ngx_rbtree_t>,
30    node: *mut ngx_rbtree_node_t,
31    _lifetime: PhantomData<&'a ()>,
32}
33
34impl<'a> NgxRbTreeIter<'a> {
35    /// Creates an iterator for the `ngx_rbtree_t`.
36    ///
37    /// # Safety
38    ///
39    /// The tree must outlive the iterator.
40    pub unsafe fn new(tree: NonNull<ngx_rbtree_t>) -> Self {
41        let t = unsafe { tree.as_ref() };
42        let node = if ptr::addr_eq(t.root, t.sentinel) {
43            // empty tree
44            ptr::null_mut()
45        } else {
46            unsafe { ngx_rbtree_min(t.root, t.sentinel) }
47        };
48
49        Self {
50            tree,
51            node,
52            _lifetime: PhantomData,
53        }
54    }
55}
56
57impl<'a> Iterator for NgxRbTreeIter<'a> {
58    type Item = NonNull<ngx_rbtree_node_t>;
59
60    fn next(&mut self) -> Option<Self::Item> {
61        let item = NonNull::new(self.node)?;
62        // ngx_rbtree_next does not mutate the tree
63        self.node = unsafe { ngx_rbtree_next(self.tree.as_mut(), self.node) };
64        Some(item)
65    }
66}
67
68#[allow(deprecated)]
69type BuildMapHasher = core::hash::BuildHasherDefault<hash::SipHasher>;
70
71/// A map type based on the `ngx_rbtree_t`.
72///
73/// This map implementation owns the stored keys and values and ensures that the data is dropped.
74/// The order of the elements is an undocumented implementation detail.
75///
76/// This is a `ngx`-specific high-level type with no direct counterpart in the NGINX code.
77#[derive(Debug)]
78pub struct RbTreeMap<K, V, A>
79where
80    A: Allocator,
81{
82    tree: ngx_rbtree_t,
83    sentinel: NonNull<ngx_rbtree_node_t>,
84    alloc: A,
85    _kv_type: PhantomData<(K, V)>,
86}
87
88/// Entry type for the [RbTreeMap].
89///
90/// The struct is used from the Rust code only and thus does not need to be compatible with C.
91struct MapEntry<K, V> {
92    node: ngx_rbtree_node_t,
93    key: K,
94    value: V,
95}
96
97impl<K, V> MapEntry<K, V>
98where
99    K: Hash,
100{
101    fn new(key: K, value: V) -> Self {
102        let mut node: ngx_rbtree_node_t = unsafe { mem::zeroed() };
103        node.key = BuildMapHasher::default().hash_one(&key) as ngx_rbtree_key_t;
104
105        Self { node, key, value }
106    }
107
108    fn into_kv(self) -> (K, V) {
109        (self.key, self.value)
110    }
111}
112
113/// An iterator for the [RbTreeMap].
114pub struct Iter<'a, K: 'a, V: 'a>(NgxRbTreeIter<'a>, PhantomData<(K, V)>);
115
116impl<'a, K: 'a, V: 'a> Iter<'a, K, V> {
117    /// Creates an iterator for the [RbTreeMap].
118    pub fn new<A: Allocator>(tree: &'a RbTreeMap<K, V, A>) -> Self {
119        // msrv(1.89.0): NonNull::from_ref()
120        let rbtree = NonNull::from(&tree.tree);
121        // SAFETY: Iter borrows from the tree, ensuring that the tree would outlive it.
122        Self(unsafe { NgxRbTreeIter::new(rbtree) }, Default::default())
123    }
124}
125
126impl<'a, K: 'a, V: 'a> Iterator for Iter<'a, K, V> {
127    type Item = (&'a K, &'a V);
128
129    fn next(&mut self) -> Option<Self::Item> {
130        let item = self.0.next()?;
131        let item = unsafe { ngx_rbtree_data!(item, MapEntry<K, V>, node).as_ref() };
132        Some((&item.key, &item.value))
133    }
134}
135
136/// A mutable iterator for the [RbTreeMap].
137pub struct IterMut<'a, K: 'a, V: 'a>(NgxRbTreeIter<'a>, PhantomData<(K, V)>);
138
139impl<'a, K: 'a, V: 'a> IterMut<'a, K, V> {
140    /// Creates an iterator for the [RbTreeMap].
141    pub fn new<A: Allocator>(tree: &'a mut RbTreeMap<K, V, A>) -> Self {
142        // msrv(1.89.0): NonNull::from_mut()
143        let rbtree = NonNull::from(&mut tree.tree);
144        // SAFETY: IterMut borrows from the tree, ensuring that the tree would outlive it.
145        Self(unsafe { NgxRbTreeIter::new(rbtree) }, Default::default())
146    }
147}
148
149impl<'a, K: 'a, V: 'a> Iterator for IterMut<'a, K, V> {
150    type Item = (&'a K, &'a mut V);
151
152    fn next(&mut self) -> Option<Self::Item> {
153        let item = self.0.next()?;
154        let item = unsafe { ngx_rbtree_data!(item, MapEntry<K, V>, node).as_mut() };
155        Some((&item.key, &mut item.value))
156    }
157}
158
159impl<K, V, A> RbTreeMap<K, V, A>
160where
161    A: Allocator,
162{
163    /// Returns a reference to the underlying allocator.
164    pub fn allocator(&self) -> &A {
165        &self.alloc
166    }
167
168    /// Clears the tree, removing all elements.
169    pub fn clear(&mut self) {
170        // SAFETY: the iter lives until the end of the scope
171        let iter = unsafe { NgxRbTreeIter::new(NonNull::from(&self.tree)) };
172        let layout = Layout::new::<MapEntry<K, V>>();
173
174        for node in iter {
175            unsafe {
176                let mut data = ngx_rbtree_data!(node, MapEntry<K, V>, node);
177
178                ngx_rbtree_delete(&mut self.tree, &mut data.as_mut().node);
179                ptr::drop_in_place(data.as_mut());
180                self.allocator().deallocate(data.cast(), layout)
181            }
182        }
183    }
184
185    /// Returns true if the tree contains no entries.
186    pub fn is_empty(&self) -> bool {
187        ptr::addr_eq(self.tree.root, self.tree.sentinel)
188    }
189
190    /// Returns an iterator over the entries of the tree.
191    #[inline]
192    pub fn iter(&self) -> Iter<'_, K, V> {
193        Iter::new(self)
194    }
195
196    /// Returns a mutable iterator over the entries of the tree.
197    #[inline]
198    pub fn iter_mut(&mut self) -> IterMut<'_, K, V> {
199        IterMut::new(self)
200    }
201}
202
203impl<K, V, A> RbTreeMap<K, V, A>
204where
205    A: Allocator,
206    K: Hash + Ord,
207{
208    /// Attempts to create and initialize a new RbTreeMap with specified allocator.
209    pub fn try_new_in(alloc: A) -> Result<Self, AllocError> {
210        let layout = Layout::new::<ngx_rbtree_node_t>();
211        let sentinel: NonNull<ngx_rbtree_node_t> = alloc.allocate_zeroed(layout)?.cast();
212
213        let mut this = RbTreeMap {
214            tree: unsafe { mem::zeroed() },
215            sentinel,
216            alloc,
217            _kv_type: PhantomData,
218        };
219
220        unsafe { ngx_rbtree_init(&mut this.tree, this.sentinel.as_ptr(), Some(Self::insert)) };
221
222        Ok(this)
223    }
224
225    /// Returns a reference to the value corresponding to the key.
226    pub fn get<Q>(&self, key: &Q) -> Option<&V>
227    where
228        K: borrow::Borrow<Q>,
229        Q: Hash + Ord + ?Sized,
230    {
231        self.lookup(key).map(|x| unsafe { &x.as_ref().value })
232    }
233
234    /// Returns a mutable reference to the value corresponding to the key.
235    pub fn get_mut<Q>(&mut self, key: &Q) -> Option<&mut V>
236    where
237        K: borrow::Borrow<Q>,
238        Q: Hash + Ord + ?Sized,
239    {
240        self.lookup(key)
241            .map(|mut x| unsafe { &mut x.as_mut().value })
242    }
243
244    /// Removes a key from the tree, returning the value at the key if the key was previously in the
245    /// tree.
246    pub fn remove<Q>(&mut self, key: &Q) -> Option<V>
247    where
248        K: borrow::Borrow<Q>,
249        Q: Hash + Ord + ?Sized,
250    {
251        self.remove_entry(key).map(|(_, v)| v)
252    }
253
254    /// Removes a key from the tree, returning the stored key and value if the key was previously in
255    /// the tree.
256    pub fn remove_entry<Q>(&mut self, key: &Q) -> Option<(K, V)>
257    where
258        K: borrow::Borrow<Q>,
259        Q: Hash + Ord + ?Sized,
260    {
261        let mut node = self.lookup(key)?;
262        unsafe {
263            ngx_rbtree_delete(&mut self.tree, &mut node.as_mut().node);
264            let layout = Layout::for_value(node.as_ref());
265            // SAFETY: we make a bitwise copy of the node and dispose of the original value without
266            // dropping it.
267            let copy = node.as_ptr().read();
268            self.allocator().deallocate(node.cast(), layout);
269            Some(copy.into_kv())
270        }
271    }
272
273    /// Attempts to insert a new element into the tree.
274    pub fn try_insert(&mut self, key: K, value: V) -> Result<&mut V, AllocError> {
275        let mut node = if let Some(mut node) = self.lookup(&key) {
276            unsafe { node.as_mut().value = value };
277            node
278        } else {
279            let node = MapEntry::new(key, value);
280            let mut node = allocator::allocate(node, self.allocator())?;
281            unsafe { ngx_rbtree_insert(&mut self.tree, &mut node.as_mut().node) };
282            node
283        };
284
285        Ok(unsafe { &mut node.as_mut().value })
286    }
287
288    extern "C" fn insert(
289        mut temp: *mut ngx_rbtree_node_t,
290        node: *mut ngx_rbtree_node_t,
291        sentinel: *mut ngx_rbtree_node_t,
292    ) {
293        let n = unsafe { &mut *ngx_rbtree_data!(node, MapEntry<K, V>, node) };
294
295        loop {
296            let t = unsafe { &mut *ngx_rbtree_data!(temp, MapEntry<K, V>, node) };
297            let p = match Ord::cmp(&n.node.key, &t.node.key) {
298                Ordering::Less => &mut t.node.left,
299                Ordering::Greater => &mut t.node.right,
300                Ordering::Equal => match Ord::cmp(&n.key, &t.key) {
301                    Ordering::Less => &mut t.node.left,
302                    Ordering::Greater => &mut t.node.right,
303                    // should be handled in try_insert
304                    Ordering::Equal => &mut t.node.right,
305                },
306            };
307
308            if ptr::addr_eq(*p, sentinel) {
309                *p = node;
310                break;
311            }
312
313            temp = *p;
314        }
315
316        n.node.parent = temp;
317        n.node.left = sentinel;
318        n.node.right = sentinel;
319        unsafe { ngx_rbt_red(node) };
320    }
321
322    fn lookup<Q>(&self, key: &Q) -> Option<NonNull<MapEntry<K, V>>>
323    where
324        K: borrow::Borrow<Q>,
325        Q: Hash + Ord + ?Sized,
326    {
327        let mut node = self.tree.root;
328        let hash = BuildMapHasher::default().hash_one(key) as ngx_rbtree_key_t;
329
330        while !ptr::addr_eq(node, self.tree.sentinel) {
331            let n = unsafe { NonNull::new_unchecked(ngx_rbtree_data!(node, MapEntry<K, V>, node)) };
332            let nr = unsafe { n.as_ref() };
333
334            node = match Ord::cmp(&hash, &nr.node.key) {
335                Ordering::Less => nr.node.left,
336                Ordering::Greater => nr.node.right,
337                Ordering::Equal => match Ord::cmp(key, nr.key.borrow()) {
338                    Ordering::Less => nr.node.left,
339                    Ordering::Greater => nr.node.right,
340                    Ordering::Equal => return Some(n),
341                },
342            }
343        }
344
345        None
346    }
347}
348
349impl<K, V, A> Drop for RbTreeMap<K, V, A>
350where
351    A: Allocator,
352{
353    fn drop(&mut self) {
354        self.clear();
355
356        unsafe {
357            self.allocator().deallocate(
358                self.sentinel.cast(),
359                Layout::for_value(self.sentinel.as_ref()),
360            )
361        };
362    }
363}
364
365unsafe impl<K, V, A> Send for RbTreeMap<K, V, A>
366where
367    A: Send + Allocator,
368    K: Send,
369    V: Send,
370{
371}
372
373unsafe impl<K, V, A> Sync for RbTreeMap<K, V, A>
374where
375    A: Sync + Allocator,
376    K: Sync,
377    V: Sync,
378{
379}
ngx/collections/rbtree.rs

ngx/collections/
rbtree.rs
