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/// Trait for pointer conversions between the tree entry and its container.
23///
24/// # Safety
25///
26/// This trait must only be implemented on types that contain a tree node or wrappers with
27/// compatible layout. The type then can be used to access elements of a raw rbtree type
28/// [NgxRbTree] linked via specified field.
29///
30/// If the struct can belong to several trees through multiple embedded `ngx_rbtree_node_t` fields,
31/// a separate [NgxRbTreeEntry] implementation via wrapper type should be used for each tree.
32pub unsafe trait NgxRbTreeEntry {
33    /// Gets a container pointer from tree node.
34    fn from_rbtree_node(node: NonNull<ngx_rbtree_node_t>) -> NonNull<Self>;
35    /// Gets an rbtree node from a container reference.
36    fn to_rbtree_node(&mut self) -> &mut ngx_rbtree_node_t;
37}
38
39unsafe impl NgxRbTreeEntry for ngx_rbtree_node_t {
40    fn from_rbtree_node(node: NonNull<ngx_rbtree_node_t>) -> NonNull<Self> {
41        node
42    }
43
44    fn to_rbtree_node(&mut self) -> &mut ngx_rbtree_node_t {
45        self
46    }
47}
48
49/// A wrapper over a raw `ngx_rbtree_t`, a red-black tree implementation.
50///
51/// This wrapper is defined in terms of type `T` that embeds and can be converted from or to the
52/// tree nodes.
53///
54/// See <https://nginx.org/en/docs/dev/development_guide.html#red_black_tree>.
55#[derive(Debug)]
56#[repr(transparent)]
57pub struct NgxRbTree<T> {
58    inner: ngx_rbtree_t,
59    _type: PhantomData<T>,
60}
61
62impl<T> NgxRbTree<T>
63where
64    T: NgxRbTreeEntry,
65{
66    /// Creates a tree reference from a pointer to [ngx_rbtree_t].
67    ///
68    /// # Safety
69    ///
70    /// `tree` is a valid pointer to [ngx_rbtree_t], and `T::from_rbtree_node` on the tree nodes
71    /// results in valid pointers to `T`.
72    pub unsafe fn from_ptr<'a>(tree: *const ngx_rbtree_t) -> &'a Self {
73        &*tree.cast()
74    }
75
76    /// Creates a mutable tree reference from a pointer to [ngx_rbtree_t].
77    ///
78    /// # Safety
79    ///
80    /// `tree` is a valid pointer to [ngx_rbtree_t], and `T::from_rbtree_node` on the tree nodes
81    /// results in valid pointers to `T`.
82    pub unsafe fn from_ptr_mut<'a>(tree: *mut ngx_rbtree_t) -> &'a mut Self {
83        &mut *tree.cast()
84    }
85
86    /// Returns `true` if the tree contains no elements.
87    pub fn is_empty(&self) -> bool {
88        ptr::addr_eq(self.inner.root, self.inner.sentinel)
89    }
90
91    /// Appends a node to the tree.
92    pub fn insert(&mut self, node: &mut T) {
93        unsafe { ngx_rbtree_insert(&mut self.inner, node.to_rbtree_node()) };
94    }
95
96    /// Removes the specified node from the tree.
97    pub fn remove(&mut self, node: &mut T) {
98        unsafe { ngx_rbtree_delete(&mut self.inner, node.to_rbtree_node()) };
99    }
100
101    /// Returns an iterator over the nodes of the tree.
102    pub fn iter(&self) -> NgxRbTreeIter<'_> {
103        unsafe { NgxRbTreeIter::new(NonNull::from(&self.inner)) }
104    }
105
106    /// Returns a mutable iterator over the nodes of the tree.
107    pub fn iter_mut(&mut self) -> NgxRbTreeIter<'_> {
108        unsafe { NgxRbTreeIter::new(NonNull::from(&mut self.inner)) }
109    }
110}
111
112/// Raw iterator over the `ngx_rbtree_t` nodes.
113///
114/// This iterator type can be used to access elements of any correctly initialized `ngx_rbtree_t`
115/// instance, including those already embedded in the nginx structures.  The iterator stores pointer
116/// to the next node and thus remains valid and usable even if the last returned item is removed
117/// from the tree.
118pub struct NgxRbTreeIter<'a> {
119    tree: NonNull<ngx_rbtree_t>,
120    node: *mut ngx_rbtree_node_t,
121    _lifetime: PhantomData<&'a ()>,
122}
123
124impl NgxRbTreeIter<'_> {
125    /// Creates an iterator for the `ngx_rbtree_t`.
126    ///
127    /// # Safety
128    ///
129    /// The tree must outlive the iterator.
130    pub unsafe fn new(tree: NonNull<ngx_rbtree_t>) -> Self {
131        let t = unsafe { tree.as_ref() };
132        let node = if ptr::addr_eq(t.root, t.sentinel) {
133            // empty tree
134            ptr::null_mut()
135        } else {
136            unsafe { ngx_rbtree_min(t.root, t.sentinel) }
137        };
138
139        Self {
140            tree,
141            node,
142            _lifetime: PhantomData,
143        }
144    }
145}
146
147impl Iterator for NgxRbTreeIter<'_> {
148    type Item = NonNull<ngx_rbtree_node_t>;
149
150    fn next(&mut self) -> Option<Self::Item> {
151        let item = NonNull::new(self.node)?;
152        // ngx_rbtree_next does not mutate the tree
153        self.node = unsafe { ngx_rbtree_next(self.tree.as_mut(), self.node) };
154        Some(item)
155    }
156}
157
158#[allow(deprecated)]
159type BuildMapHasher = core::hash::BuildHasherDefault<hash::SipHasher>;
160
161/// A map type based on the `ngx_rbtree_t`.
162///
163/// This map implementation owns the stored keys and values and ensures that the data is dropped.
164/// The order of the elements is an undocumented implementation detail.
165///
166/// This is a `ngx`-specific high-level type with no direct counterpart in the NGINX code.
167#[derive(Debug)]
168pub struct RbTreeMap<K, V, A>
169where
170    A: Allocator,
171{
172    tree: NgxRbTree<MapEntry<K, V>>,
173    sentinel: NonNull<ngx_rbtree_node_t>,
174    alloc: A,
175}
176
177/// Entry type for the [RbTreeMap].
178///
179/// The struct is used from the Rust code only and thus does not need to be compatible with C.
180#[derive(Debug)]
181struct MapEntry<K, V> {
182    node: ngx_rbtree_node_t,
183    key: K,
184    value: V,
185}
186
187impl<K, V> MapEntry<K, V>
188where
189    K: Hash,
190{
191    fn new(key: K, value: V) -> Self {
192        let mut node: ngx_rbtree_node_t = unsafe { mem::zeroed() };
193        node.key = BuildMapHasher::default().hash_one(&key) as ngx_rbtree_key_t;
194
195        Self { node, key, value }
196    }
197
198    fn into_kv(self) -> (K, V) {
199        (self.key, self.value)
200    }
201}
202
203unsafe impl<K, V> NgxRbTreeEntry for MapEntry<K, V> {
204    fn from_rbtree_node(node: NonNull<ngx_rbtree_node_t>) -> NonNull<Self> {
205        unsafe { ngx_rbtree_data!(node, Self, node) }
206    }
207
208    fn to_rbtree_node(&mut self) -> &mut ngx_rbtree_node_t {
209        &mut self.node
210    }
211}
212
213/// An iterator for the [RbTreeMap].
214pub struct MapIter<'a, K: 'a, V: 'a>(NgxRbTreeIter<'a>, PhantomData<(K, V)>);
215
216impl<'a, K: 'a, V: 'a> MapIter<'a, K, V> {
217    /// Creates an iterator for the [RbTreeMap].
218    pub fn new<A: Allocator>(tree: &'a RbTreeMap<K, V, A>) -> Self {
219        // msrv(1.89.0): NonNull::from_ref()
220        let rbtree = NonNull::from(&tree.tree.inner);
221        // SAFETY: Iter borrows from the tree, ensuring that the tree would outlive it.
222        Self(unsafe { NgxRbTreeIter::new(rbtree) }, Default::default())
223    }
224}
225
226impl<'a, K: 'a, V: 'a> Iterator for MapIter<'a, K, V> {
227    type Item = (&'a K, &'a V);
228
229    fn next(&mut self) -> Option<Self::Item> {
230        let item = self.0.next()?;
231        let item = unsafe { ngx_rbtree_data!(item, MapEntry<K, V>, node).as_ref() };
232        Some((&item.key, &item.value))
233    }
234}
235
236/// A mutable iterator for the [RbTreeMap].
237pub struct MapIterMut<'a, K: 'a, V: 'a>(NgxRbTreeIter<'a>, PhantomData<(K, V)>);
238
239impl<'a, K: 'a, V: 'a> MapIterMut<'a, K, V> {
240    /// Creates an iterator for the [RbTreeMap].
241    pub fn new<A: Allocator>(tree: &'a mut RbTreeMap<K, V, A>) -> Self {
242        // msrv(1.89.0): NonNull::from_mut()
243        let rbtree = NonNull::from(&mut tree.tree.inner);
244        // SAFETY: IterMut borrows from the tree, ensuring that the tree would outlive it.
245        Self(unsafe { NgxRbTreeIter::new(rbtree) }, Default::default())
246    }
247}
248
249impl<'a, K: 'a, V: 'a> Iterator for MapIterMut<'a, K, V> {
250    type Item = (&'a K, &'a mut V);
251
252    fn next(&mut self) -> Option<Self::Item> {
253        let mut item = MapEntry::<K, V>::from_rbtree_node(self.0.next()?);
254        let item = unsafe { item.as_mut() };
255        Some((&item.key, &mut item.value))
256    }
257}
258
259impl<K, V, A> RbTreeMap<K, V, A>
260where
261    A: Allocator,
262{
263    /// Returns a reference to the underlying allocator.
264    pub fn allocator(&self) -> &A {
265        &self.alloc
266    }
267
268    /// Clears the tree, removing all elements.
269    pub fn clear(&mut self) {
270        // SAFETY: the iter lives until the end of the scope
271        let iter = unsafe { NgxRbTreeIter::new(NonNull::from(&self.tree.inner)) };
272        let layout = Layout::new::<MapEntry<K, V>>();
273
274        for node in iter {
275            unsafe {
276                let mut data = MapEntry::<K, V>::from_rbtree_node(node);
277
278                ngx_rbtree_delete(&mut self.tree.inner, &mut data.as_mut().node);
279                ptr::drop_in_place(data.as_mut());
280                self.allocator().deallocate(data.cast(), layout)
281            }
282        }
283    }
284
285    /// Returns true if the tree contains no entries.
286    pub fn is_empty(&self) -> bool {
287        self.tree.is_empty()
288    }
289
290    /// Returns an iterator over the entries of the tree.
291    #[inline]
292    pub fn iter(&self) -> MapIter<'_, K, V> {
293        MapIter::new(self)
294    }
295
296    /// Returns a mutable iterator over the entries of the tree.
297    #[inline]
298    pub fn iter_mut(&mut self) -> MapIterMut<'_, K, V> {
299        MapIterMut::new(self)
300    }
301}
302
303impl<K, V, A> RbTreeMap<K, V, A>
304where
305    A: Allocator,
306    K: Hash + Ord,
307{
308    /// Attempts to create and initialize a new RbTreeMap with specified allocator.
309    pub fn try_new_in(alloc: A) -> Result<Self, AllocError> {
310        let layout = Layout::new::<ngx_rbtree_node_t>();
311        let sentinel: NonNull<ngx_rbtree_node_t> = alloc.allocate_zeroed(layout)?.cast();
312
313        let tree = NgxRbTree {
314            inner: unsafe { mem::zeroed() },
315            _type: PhantomData,
316        };
317
318        let mut this = RbTreeMap {
319            tree,
320            sentinel,
321            alloc,
322        };
323
324        unsafe {
325            ngx_rbtree_init(
326                &mut this.tree.inner,
327                this.sentinel.as_ptr(),
328                Some(Self::insert),
329            )
330        };
331
332        Ok(this)
333    }
334
335    /// Returns a reference to the value corresponding to the key.
336    pub fn get<Q>(&self, key: &Q) -> Option<&V>
337    where
338        K: borrow::Borrow<Q>,
339        Q: Hash + Ord + ?Sized,
340    {
341        self.lookup(key).map(|x| unsafe { &x.as_ref().value })
342    }
343
344    /// Returns a mutable reference to the value corresponding to the key.
345    pub fn get_mut<Q>(&mut self, key: &Q) -> Option<&mut V>
346    where
347        K: borrow::Borrow<Q>,
348        Q: Hash + Ord + ?Sized,
349    {
350        self.lookup(key)
351            .map(|mut x| unsafe { &mut x.as_mut().value })
352    }
353
354    /// Removes a key from the tree, returning the value at the key if the key was previously in the
355    /// tree.
356    pub fn remove<Q>(&mut self, key: &Q) -> Option<V>
357    where
358        K: borrow::Borrow<Q>,
359        Q: Hash + Ord + ?Sized,
360    {
361        self.remove_entry(key).map(|(_, v)| v)
362    }
363
364    /// Removes a key from the tree, returning the stored key and value if the key was previously in
365    /// the tree.
366    pub fn remove_entry<Q>(&mut self, key: &Q) -> Option<(K, V)>
367    where
368        K: borrow::Borrow<Q>,
369        Q: Hash + Ord + ?Sized,
370    {
371        let mut node = self.lookup(key)?;
372        unsafe {
373            self.tree.remove(node.as_mut());
374
375            let layout = Layout::for_value(node.as_ref());
376            // SAFETY: we make a bitwise copy of the node and dispose of the original value without
377            // dropping it.
378            let copy = node.as_ptr().read();
379            self.allocator().deallocate(node.cast(), layout);
380            Some(copy.into_kv())
381        }
382    }
383
384    /// Attempts to insert a new element into the tree.
385    pub fn try_insert(&mut self, key: K, value: V) -> Result<&mut V, AllocError> {
386        let mut node = if let Some(mut node) = self.lookup(&key) {
387            unsafe { node.as_mut().value = value };
388            node
389        } else {
390            let node = MapEntry::new(key, value);
391            let mut node = allocator::allocate(node, self.allocator())?;
392            self.tree.insert(unsafe { node.as_mut() });
393            node
394        };
395
396        Ok(unsafe { &mut node.as_mut().value })
397    }
398
399    extern "C" fn insert(
400        mut temp: *mut ngx_rbtree_node_t,
401        node: *mut ngx_rbtree_node_t,
402        sentinel: *mut ngx_rbtree_node_t,
403    ) {
404        let n = unsafe { &mut *ngx_rbtree_data!(node, MapEntry<K, V>, node) };
405
406        loop {
407            let t = unsafe { &mut *ngx_rbtree_data!(temp, MapEntry<K, V>, node) };
408            let p = match Ord::cmp(&n.node.key, &t.node.key) {
409                Ordering::Less => &mut t.node.left,
410                Ordering::Greater => &mut t.node.right,
411                Ordering::Equal => match Ord::cmp(&n.key, &t.key) {
412                    Ordering::Less => &mut t.node.left,
413                    Ordering::Greater => &mut t.node.right,
414                    // should be handled in try_insert
415                    Ordering::Equal => &mut t.node.right,
416                },
417            };
418
419            if ptr::addr_eq(*p, sentinel) {
420                *p = node;
421                break;
422            }
423
424            temp = *p;
425        }
426
427        n.node.parent = temp;
428        n.node.left = sentinel;
429        n.node.right = sentinel;
430        unsafe { ngx_rbt_red(node) };
431    }
432
433    fn lookup<Q>(&self, key: &Q) -> Option<NonNull<MapEntry<K, V>>>
434    where
435        K: borrow::Borrow<Q>,
436        Q: Hash + Ord + ?Sized,
437    {
438        let mut node = self.tree.inner.root;
439        let hash = BuildMapHasher::default().hash_one(key) as ngx_rbtree_key_t;
440
441        while !ptr::addr_eq(node, self.tree.inner.sentinel) {
442            let n = unsafe { NonNull::new_unchecked(ngx_rbtree_data!(node, MapEntry<K, V>, node)) };
443            let nr = unsafe { n.as_ref() };
444
445            node = match Ord::cmp(&hash, &nr.node.key) {
446                Ordering::Less => nr.node.left,
447                Ordering::Greater => nr.node.right,
448                Ordering::Equal => match Ord::cmp(key, nr.key.borrow()) {
449                    Ordering::Less => nr.node.left,
450                    Ordering::Greater => nr.node.right,
451                    Ordering::Equal => return Some(n),
452                },
453            }
454        }
455
456        None
457    }
458}
459
460impl<K, V, A> Drop for RbTreeMap<K, V, A>
461where
462    A: Allocator,
463{
464    fn drop(&mut self) {
465        self.clear();
466
467        unsafe {
468            self.allocator().deallocate(
469                self.sentinel.cast(),
470                Layout::for_value(self.sentinel.as_ref()),
471            )
472        };
473    }
474}
475
476unsafe impl<K, V, A> Send for RbTreeMap<K, V, A>
477where
478    A: Send + Allocator,
479    K: Send,
480    V: Send,
481{
482}
483
484unsafe impl<K, V, A> Sync for RbTreeMap<K, V, A>
485where
486    A: Sync + Allocator,
487    K: Sync,
488    V: Sync,
489{
490}
ngx/collections/rbtree.rs

ngx/collections/
rbtree.rs
