Skip to content

Commit c01f18a

Browse files
committed
fixup! feat: wrapper for ngx_rbtree_t
1 parent dd7746a commit c01f18a

File tree

1 file changed

+64
-50
lines changed

1 file changed

+64
-50
lines changed

src/core/rbtree.rs

Lines changed: 64 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -8,9 +8,8 @@ use core::ptr::{self, NonNull};
88
use core::{borrow, mem};
99

1010
use nginx_sys::{
11-
ngx_rbtree_data, ngx_rbtree_delete, ngx_rbtree_init, ngx_rbtree_insert,
12-
ngx_rbtree_insert_value, ngx_rbtree_key_t, ngx_rbtree_min, ngx_rbtree_next, ngx_rbtree_node_t,
13-
ngx_rbtree_t,
11+
ngx_rbt_red, ngx_rbtree_data, ngx_rbtree_delete, ngx_rbtree_init, ngx_rbtree_insert,
12+
ngx_rbtree_key_t, ngx_rbtree_min, ngx_rbtree_next, ngx_rbtree_node_t, ngx_rbtree_t,
1413
};
1514

1615
use crate::allocator::{self, AllocError, Allocator};
@@ -147,29 +146,6 @@ where
147146
&self.alloc
148147
}
149148

150-
/// Attempts to create and initialize a new RbTree with specified allocator.
151-
pub fn try_new_in(alloc: A) -> Result<Self, AllocError> {
152-
let layout = Layout::new::<ngx_rbtree_node_t>();
153-
let sentinel: NonNull<ngx_rbtree_node_t> = alloc.allocate_zeroed(layout)?.cast();
154-
155-
let mut this = RbTree {
156-
tree: unsafe { mem::zeroed() },
157-
sentinel,
158-
alloc,
159-
_ph: Default::default(),
160-
};
161-
162-
unsafe {
163-
ngx_rbtree_init(
164-
&mut this.tree,
165-
this.sentinel.as_ptr(),
166-
Some(ngx_rbtree_insert_value),
167-
)
168-
};
169-
170-
Ok(this)
171-
}
172-
173149
/// Clears the tree, removing all elements.
174150
pub fn clear(&mut self) {
175151
let iter = RawIter::new(NonNull::from(&self.tree));
@@ -211,6 +187,23 @@ where
211187
A: Allocator,
212188
K: Hash + Ord,
213189
{
190+
/// Attempts to create and initialize a new RbTree with specified allocator.
191+
pub fn try_new_in(alloc: A) -> Result<Self, AllocError> {
192+
let layout = Layout::new::<ngx_rbtree_node_t>();
193+
let sentinel: NonNull<ngx_rbtree_node_t> = alloc.allocate_zeroed(layout)?.cast();
194+
195+
let mut this = RbTree {
196+
tree: unsafe { mem::zeroed() },
197+
sentinel,
198+
alloc,
199+
_ph: Default::default(),
200+
};
201+
202+
unsafe { ngx_rbtree_init(&mut this.tree, this.sentinel.as_ptr(), Some(Self::insert)) };
203+
204+
Ok(this)
205+
}
206+
214207
/// Returns a reference to the value corresponding to the key.
215208
pub fn get<Q>(&self, key: &Q) -> Option<&V>
216209
where
@@ -274,6 +267,40 @@ where
274267
Ok(unsafe { &mut node.as_mut().value })
275268
}
276269

270+
unsafe extern "C" fn insert(
271+
mut temp: *mut ngx_rbtree_node_t,
272+
node: *mut ngx_rbtree_node_t,
273+
sentinel: *mut ngx_rbtree_node_t,
274+
) {
275+
let n = &mut *ngx_rbtree_data!(node, Node<K, V>, node);
276+
277+
loop {
278+
let t = &mut *ngx_rbtree_data!(temp, Node<K, V>, node);
279+
let p = match Ord::cmp(&n.node.key, &t.node.key) {
280+
Ordering::Less => &mut t.node.left,
281+
Ordering::Greater => &mut t.node.right,
282+
Ordering::Equal => match Ord::cmp(&n.key, &t.key) {
283+
Ordering::Less => &mut t.node.left,
284+
Ordering::Greater => &mut t.node.right,
285+
// should be handled in try_insert
286+
Ordering::Equal => &mut t.node.right,
287+
},
288+
};
289+
290+
if ptr::addr_eq(*p, sentinel) {
291+
*p = node;
292+
break;
293+
}
294+
295+
temp = *p;
296+
}
297+
298+
n.node.parent = temp;
299+
n.node.left = sentinel;
300+
n.node.right = sentinel;
301+
ngx_rbt_red(node);
302+
}
303+
277304
fn lookup<Q>(&self, key: &Q) -> Option<NonNull<Node<K, V>>>
278305
where
279306
K: borrow::Borrow<Q>,
@@ -283,30 +310,17 @@ where
283310
let hash = RbTreeHasher::default().hash_one(key) as ngx_rbtree_key_t;
284311

285312
while !ptr::addr_eq(node, self.tree.sentinel) {
286-
match Ord::cmp(&hash, unsafe { &(*node).key }) {
287-
Ordering::Less => {
288-
node = unsafe { (*node).left };
289-
continue;
290-
}
291-
Ordering::Greater => {
292-
node = unsafe { (*node).right };
293-
continue;
294-
}
295-
Ordering::Equal => (),
296-
};
297-
298-
let n = unsafe { ngx_rbtree_data!(node, Node<K, V>, node) };
299-
300-
match Ord::cmp(key, unsafe { (*n).key.borrow() }) {
301-
Ordering::Less => {
302-
node = unsafe { (*node).left };
303-
continue;
304-
}
305-
Ordering::Greater => {
306-
node = unsafe { (*node).right };
307-
continue;
308-
}
309-
Ordering::Equal => return Some(unsafe { NonNull::new_unchecked(n) }),
313+
let n = unsafe { NonNull::new_unchecked(ngx_rbtree_data!(node, Node<K, V>, node)) };
314+
let nr = unsafe { n.as_ref() };
315+
316+
node = match Ord::cmp(&hash, &nr.node.key) {
317+
Ordering::Less => nr.node.left,
318+
Ordering::Greater => nr.node.right,
319+
Ordering::Equal => match Ord::cmp(key, nr.key.borrow()) {
320+
Ordering::Less => nr.node.left,
321+
Ordering::Greater => nr.node.right,
322+
Ordering::Equal => return Some(n),
323+
},
310324
}
311325
}
312326

0 commit comments

Comments
 (0)