congee/lib.rs
1#![doc = include_str!("../README.md")]
2#![allow(clippy::comparison_chain)]
3#![allow(clippy::len_without_is_empty)]
4#![cfg_attr(docsrs, feature(doc_cfg))]
5
6mod base_node;
7mod congee_arc;
8mod error;
9mod lock;
10mod node_16;
11mod node_256;
12mod node_4;
13mod node_48;
14mod node_ptr;
15mod tree;
16mod utils;
17
18mod range_scan;
19
20mod stats;
21
22#[cfg(test)]
23mod tests;
24
25use std::{marker::PhantomData, sync::Arc};
26
27pub use congee_arc::CongeeArc;
28use error::OOMError;
29use tree::RawCongee;
30
31/// Types needed to safely access shared data concurrently.
32pub mod epoch {
33 pub use crossbeam_epoch::{Guard, pin};
34}
35
36#[derive(Clone)]
37pub struct DefaultAllocator {}
38
39unsafe impl Send for DefaultAllocator {}
40unsafe impl Sync for DefaultAllocator {}
41
42/// We should use the `Allocator` trait in the std, but it is not stable yet.
43/// https://github.com/rust-lang/rust/issues/32838
44pub trait Allocator {
45 fn allocate(&self, layout: std::alloc::Layout) -> Result<std::ptr::NonNull<[u8]>, OOMError>;
46 fn allocate_zeroed(
47 &self,
48 layout: std::alloc::Layout,
49 ) -> Result<std::ptr::NonNull<[u8]>, OOMError> {
50 let ptr = self.allocate(layout)?;
51 unsafe {
52 std::ptr::write_bytes(ptr.as_ptr() as *mut u8, 0, layout.size());
53 }
54 Ok(ptr)
55 }
56 /// # Safety
57 /// The caller must ensure that the pointer is valid and that the layout is correct.
58 /// The pointer must allocated by this allocator.
59 unsafe fn deallocate(&self, ptr: std::ptr::NonNull<u8>, layout: std::alloc::Layout);
60}
61
62impl Allocator for DefaultAllocator {
63 fn allocate(&self, layout: std::alloc::Layout) -> Result<std::ptr::NonNull<[u8]>, OOMError> {
64 let ptr = unsafe { std::alloc::alloc(layout) };
65 let ptr_slice = std::ptr::slice_from_raw_parts_mut(ptr, layout.size());
66 Ok(std::ptr::NonNull::new(ptr_slice).unwrap())
67 }
68
69 unsafe fn deallocate(&self, ptr: std::ptr::NonNull<u8>, layout: std::alloc::Layout) {
70 unsafe {
71 std::alloc::dealloc(ptr.as_ptr(), layout);
72 }
73 }
74}
75
76pub struct U64Congee<
77 V: Clone + From<usize> + Into<usize>,
78 A: Allocator + Clone + Send + 'static = DefaultAllocator,
79> {
80 inner: RawCongee<8, A>,
81 pt_val: PhantomData<V>,
82}
83
84impl<V: Clone + From<usize> + Into<usize>> Default for U64Congee<V> {
85 fn default() -> Self {
86 Self::new()
87 }
88}
89
90impl<V: Clone + From<usize> + Into<usize>> U64Congee<V> {
91 pub fn new() -> Self {
92 Self {
93 inner: RawCongee::new(DefaultAllocator {}, Arc::new(|_k, _v| {})),
94 pt_val: PhantomData,
95 }
96 }
97
98 pub fn get(&self, key: u64, guard: &epoch::Guard) -> Option<V> {
99 let key: [u8; 8] = key.to_be_bytes();
100 let v = self.inner.get(&key, guard)?;
101 Some(V::from(v))
102 }
103
104 pub fn insert(&self, key: u64, val: V, guard: &epoch::Guard) -> Result<Option<V>, OOMError> {
105 let key: [u8; 8] = key.to_be_bytes();
106 let val = val.into();
107 self.inner
108 .insert(&key, val, guard)
109 .map(|v| v.map(|v| V::from(v)))
110 }
111
112 pub fn remove(&self, key: u64, guard: &epoch::Guard) -> Option<V> {
113 let key: [u8; 8] = key.to_be_bytes();
114 let (old, new) = self.inner.compute_if_present(&key, &mut |_v| None, guard)?;
115 debug_assert!(new.is_none());
116 Some(V::from(old))
117 }
118
119 pub fn range(
120 &self,
121 start: u64,
122 end: u64,
123 result: &mut [([u8; 8], usize)],
124 guard: &epoch::Guard,
125 ) -> usize {
126 let start: [u8; 8] = start.to_be_bytes();
127 let end: [u8; 8] = end.to_be_bytes();
128 self.inner.range(&start, &end, result, guard)
129 }
130}
131
132/// The adaptive radix tree.
133pub struct Congee<
134 K: Copy + From<usize>,
135 V: Copy + From<usize>,
136 A: Allocator + Clone + Send + 'static = DefaultAllocator,
137> where
138 usize: From<K>,
139 usize: From<V>,
140{
141 inner: RawCongee<8, A>,
142 pt_key: PhantomData<K>,
143 pt_val: PhantomData<V>,
144}
145
146impl<K: Copy + From<usize>, V: Copy + From<usize>> Default for Congee<K, V>
147where
148 usize: From<K>,
149 usize: From<V>,
150{
151 fn default() -> Self {
152 Self::new(DefaultAllocator {})
153 }
154}
155
156impl<K: Copy + From<usize>, V: Copy + From<usize>, A: Allocator + Clone + Send> Congee<K, V, A>
157where
158 usize: From<K>,
159 usize: From<V>,
160{
161 /// Returns a copy of the value corresponding to the key.
162 ///
163 /// # Examples
164 ///
165 /// ```
166 /// use congee::Congee;
167 /// let tree = Congee::default();
168 /// let guard = tree.pin();
169 ///
170 /// tree.insert(1, 42, &guard);
171 /// assert_eq!(tree.get(&1, &guard).unwrap(), 42);
172 /// ```
173 #[inline]
174 pub fn get(&self, key: &K, guard: &epoch::Guard) -> Option<V> {
175 let key = usize::from(*key);
176 let key: [u8; 8] = key.to_be_bytes();
177 let v = self.inner.get(&key, guard)?;
178 Some(V::from(v))
179 }
180
181 /// Enters an epoch.
182 /// Note: this can be expensive, try to reuse it.
183 ///
184 /// # Examples
185 ///
186 /// ```
187 /// use congee::Congee;
188 /// let tree = Congee::<usize, usize>::default();
189 /// let guard = tree.pin();
190 /// ```
191 #[inline]
192 pub fn pin(&self) -> epoch::Guard {
193 crossbeam_epoch::pin()
194 }
195
196 /// Create an empty [Art] tree.
197 ///
198 /// # Examples
199 ///
200 /// ```
201 /// use congee::Congee;
202 /// let tree = Congee::<usize, usize>::default();
203 /// ```
204 #[inline]
205 pub fn new(allocator: A) -> Self {
206 Self::new_with_drainer(allocator, |_k, _v| {})
207 }
208
209 /// Create an empty [Art] tree with a drainer.
210 ///
211 /// The drainer is called on each of the value when the tree is dropped.
212 ///
213 /// # Examples
214 ///
215 /// ```
216 /// use congee::{Congee, DefaultAllocator};
217 /// use std::sync::Arc;
218 ///
219 /// let deleted_key = Arc::new(std::sync::atomic::AtomicUsize::new(0));
220 /// let deleted_value = Arc::new(std::sync::atomic::AtomicUsize::new(0));
221 /// let deleted_key_inner = deleted_key.clone();
222 /// let deleted_value_inner = deleted_value.clone();
223 ///
224 /// let drainer = move |k: usize, v: usize| {
225 /// deleted_key_inner.store(k, std::sync::atomic::Ordering::Relaxed);
226 /// deleted_value_inner.store(v, std::sync::atomic::Ordering::Relaxed);
227 /// };
228 ///
229 /// let tree = Congee::<usize, usize>::new_with_drainer(DefaultAllocator {}, drainer);
230 /// let pin = tree.pin();
231 /// tree.insert(1, 42, &pin).unwrap();
232 /// drop(tree);
233 /// assert_eq!(deleted_key.load(std::sync::atomic::Ordering::Relaxed), 1);
234 /// assert_eq!(deleted_value.load(std::sync::atomic::Ordering::Relaxed), 42);
235 /// ```
236 pub fn new_with_drainer(allocator: A, drainer: impl Fn(K, V) + 'static) -> Self {
237 let drainer = Arc::new(move |k: [u8; 8], v: usize| {
238 drainer(K::from(usize::from_be_bytes(k)), V::from(v))
239 });
240 Congee {
241 inner: RawCongee::new(allocator, drainer),
242 pt_key: PhantomData,
243 pt_val: PhantomData,
244 }
245 }
246
247 /// Returns if the tree is empty.
248 ///
249 /// # Examples
250 ///
251 /// ```
252 /// use congee::Congee;
253 /// let tree = Congee::default();
254 /// let guard = tree.pin();
255 /// assert!(tree.is_empty(&guard));
256 /// tree.insert(1, 42, &guard);
257 /// assert!(!tree.is_empty(&guard));
258 /// ```
259 pub fn is_empty(&self, guard: &epoch::Guard) -> bool {
260 self.inner.is_empty(guard)
261 }
262
263 /// Removes key-value pair from the tree, returns the value if the key was found.
264 ///
265 /// # Examples
266 ///
267 /// ```
268 /// use congee::Congee;
269 /// let tree = Congee::default();
270 /// let guard = tree.pin();
271 ///
272 /// tree.insert(1, 42, &guard);
273 /// let removed = tree.remove(&1, &guard);
274 /// assert_eq!(removed, Some(42));
275 /// assert!(tree.get(&1, &guard).is_none());
276 /// ```
277 #[inline]
278 pub fn remove(&self, k: &K, guard: &epoch::Guard) -> Option<V> {
279 let key = usize::from(*k);
280 let key: [u8; 8] = key.to_be_bytes();
281 let (old, new) = self.inner.compute_if_present(&key, &mut |_v| None, guard)?;
282 debug_assert!(new.is_none());
283 Some(V::from(old))
284 }
285
286 /// Insert a key-value pair to the tree, returns the previous value if the key was already present.
287 ///
288 /// # Examples
289 ///
290 /// ```
291 /// use congee::Congee;
292 /// let tree = Congee::default();
293 /// let guard = tree.pin();
294 ///
295 /// tree.insert(1, 42, &guard);
296 /// assert_eq!(tree.get(&1, &guard).unwrap(), 42);
297 /// let old = tree.insert(1, 43, &guard).unwrap();
298 /// assert_eq!(old, Some(42));
299 /// ```
300 #[inline]
301 pub fn insert(&self, k: K, v: V, guard: &epoch::Guard) -> Result<Option<V>, OOMError> {
302 let key = usize::from(k);
303 let key: [u8; 8] = key.to_be_bytes();
304 let val = self.inner.insert(&key, usize::from(v), guard);
305 val.map(|inner| inner.map(|v| V::from(v)))
306 }
307
308 /// Scan the tree with the range of [start, end], write the result to the
309 /// `result` buffer.
310 /// It scans the length of `result` or the number of the keys within the range, whichever is smaller;
311 /// returns the number of the keys scanned.
312 ///
313 /// # Examples
314 ///
315 /// ```
316 /// use congee::Congee;
317 /// let tree = Congee::default();
318 /// let guard = tree.pin();
319 ///
320 /// tree.insert(1, 42, &guard);
321 ///
322 /// let low_key = 1;
323 /// let high_key = 2;
324 /// let mut result = [(0, 0); 2];
325 /// let scanned = tree.range(&low_key, &high_key, &mut result, &guard);
326 /// assert_eq!(scanned, 1);
327 /// assert_eq!(result, [(1, 42), (0, 0)]);
328 /// ```
329 #[inline]
330 pub fn range(
331 &self,
332 start: &K,
333 end: &K,
334 result: &mut [(usize, usize)],
335 guard: &epoch::Guard,
336 ) -> usize {
337 let start = usize::from(*start);
338 let end = usize::from(*end);
339 let start: [u8; 8] = start.to_be_bytes();
340 let end: [u8; 8] = end.to_be_bytes();
341 let result_ref = unsafe {
342 std::slice::from_raw_parts_mut(
343 result.as_mut_ptr() as *mut ([u8; 8], usize),
344 result.len(),
345 )
346 };
347 let v = self.inner.range(&start, &end, result_ref, guard);
348 for i in 0..v {
349 result[i].0 = usize::from_be_bytes(result_ref[i].0);
350 }
351 v
352 }
353
354 /// Compute and update the value if the key presents in the tree.
355 /// Returns the (old, new) value
356 ///
357 /// Note that the function `f` is a FnMut and it must be safe to execute multiple times.
358 /// The `f` is expected to be short and fast as it will hold a exclusive lock on the leaf node.
359 ///
360 /// # Examples
361 ///
362 /// ```
363 /// use congee::Congee;
364 /// let tree = Congee::default();
365 /// let guard = tree.pin();
366 ///
367 /// tree.insert(1, 42, &guard);
368 /// let old = tree.compute_if_present(&1, |v| Some(v+1), &guard).unwrap();
369 /// assert_eq!(old, (42, Some(43)));
370 /// let val = tree.get(&1, &guard).unwrap();
371 /// assert_eq!(val, 43);
372 /// ```
373 #[inline]
374 pub fn compute_if_present<F>(
375 &self,
376 key: &K,
377 mut f: F,
378 guard: &epoch::Guard,
379 ) -> Option<(usize, Option<usize>)>
380 where
381 F: FnMut(usize) -> Option<usize>,
382 {
383 let key = usize::from(*key);
384 let key: [u8; 8] = key.to_be_bytes();
385 self.inner.compute_if_present(&key, &mut f, guard)
386 }
387
388 /// Compute or insert the value if the key is not in the tree.
389 /// Returns the Option(old) value
390 ///
391 /// Note that the function `f` is a FnMut and it must be safe to execute multiple times.
392 /// The `f` is expected to be short and fast as it will hold a exclusive lock on the leaf node.
393 ///
394 /// # Examples
395 ///
396 /// ```
397 /// use congee::Congee;
398 /// let tree = Congee::default();
399 /// let guard = tree.pin();
400 ///
401 /// tree.insert(1, 42, &guard);
402 /// let old = tree.compute_or_insert(1, |v| v.unwrap() + 1, &guard).unwrap().unwrap();
403 /// assert_eq!(old, 42);
404 /// let val = tree.get(&1, &guard).unwrap();
405 /// assert_eq!(val, 43);
406 ///
407 /// let old = tree.compute_or_insert(2, |v| {
408 /// assert!(v.is_none());
409 /// 2
410 /// }, &guard).unwrap();
411 /// assert!(old.is_none());
412 /// let val = tree.get(&2, &guard).unwrap();
413 /// assert_eq!(val, 2);
414 /// ```
415 pub fn compute_or_insert<F>(
416 &self,
417 key: K,
418 mut f: F,
419 guard: &epoch::Guard,
420 ) -> Result<Option<V>, OOMError>
421 where
422 F: FnMut(Option<usize>) -> usize,
423 {
424 let key = usize::from(key);
425 let key: [u8; 8] = key.to_be_bytes();
426 let u_val = self.inner.compute_or_insert(&key, &mut f, guard)?;
427 Ok(u_val.map(|v| V::from(v)))
428 }
429
430 /// Display the internal node statistics
431 pub fn stats(&self) -> stats::NodeStats {
432 self.inner.stats()
433 }
434
435 /// Update the value if the old value matches with the new one.
436 /// Returns the current value.
437 ///
438 /// # Examples:
439 /// ```
440 /// use congee::Congee;
441 /// let tree = Congee::default();
442 /// let guard = tree.pin();
443 /// tree.insert(1, 42, &guard);
444 ///
445 ///
446 /// let v = tree.compare_exchange(&1, &42, Some(43), &guard).unwrap();
447 /// assert_eq!(v, Some(43));
448 /// ```
449 pub fn compare_exchange(
450 &self,
451 key: &K,
452 old: &V,
453 new: Option<V>,
454 guard: &epoch::Guard,
455 ) -> Result<Option<V>, Option<V>> {
456 let key = usize::from(*key);
457 let key: [u8; 8] = key.to_be_bytes();
458 let new_v = new.map(|v| usize::from(v));
459 let mut fc = |v: usize| -> Option<usize> {
460 if v == usize::from(*old) {
461 new_v
462 } else {
463 Some(v)
464 }
465 };
466 let v = self.inner.compute_if_present(&key, &mut fc, guard);
467 match v {
468 Some((actual_old, actual_new)) => {
469 if actual_old == usize::from(*old) && actual_new == new_v {
470 Ok(new)
471 } else {
472 Err(actual_new.map(|v| V::from(v)))
473 }
474 }
475 None => Err(None),
476 }
477 }
478
479 /// Retrieve all keys from ART.
480 /// Isolation level: read committed.
481 ///
482 /// # Examples:
483 /// ```
484 /// use congee::Congee;
485 /// let tree = Congee::default();
486 /// let guard = tree.pin();
487 /// tree.insert(1, 42, &guard);
488 /// tree.insert(2, 43, &guard);
489 ///
490 /// let keys = tree.keys();
491 /// assert_eq!(keys, vec![1, 2]);
492 /// ```
493 pub fn keys(&self) -> Vec<K> {
494 self.inner
495 .keys()
496 .into_iter()
497 .map(|k| {
498 let key = usize::from_be_bytes(k);
499 K::from(key)
500 })
501 .collect()
502 }
503}