slice_rbtree/
tree.rs

1//! A slice-based Red-Black tree
2//!
3//! Let's assume you want to create a tree holding up to 100 pairs of `u8 <-> f64`:
4//! ```
5//! use slice_rbtree::tree::{tree_size, RBTree, TreeParams};
6//! use std::mem::size_of;
7//!
8//! // RBTree requires input slice to have a proper size
9//! let size = tree_size(
10//!     TreeParams {
11//!         k_size: size_of::<u8>(),
12//!         v_size: size_of::<f64>(),
13//!     },
14//!     100,
15//! );
16//!
17//! let mut buffer = vec![0; size];
18//!
19//! let mut tree: RBTree<u8, f64, 1, 8> = RBTree::init_slice(&mut buffer).unwrap();
20//!
21//! // Insert some values
22//! tree.insert(15, 1.245).unwrap();
23//! tree.insert(17, 5.5).unwrap();
24//! tree.insert(19, 6.5).unwrap();
25//!
26//! // This type does not implement `Drop` trait - all the data is immediately serialized in the
27//! // slice, so this is actually a no-op
28//! drop(tree);
29//!
30//! let mut new_tree: RBTree<u8, f64, 1, 8> = unsafe { RBTree::from_slice(&mut buffer).unwrap() };
31//! assert_eq!(new_tree.get(&15), Some(1.245));
32//!
33//! // Create iterator of key-value pairs and collect in a `Vec`
34//! let pairs:Vec<_> = new_tree.pairs().collect();
35//! assert_eq!(pairs[0], (15, 1.245));
36//! assert_eq!(pairs[1], (17, 5.5));
37//! assert_eq!(pairs[2], (19, 6.5));
38//!
39//! // There are 3 ways to remove a value from the tree:
40//! // 1. remove() will deserialize the value to be removed
41//! assert_eq!(new_tree.remove(&17), Some(5.5));
42//! // 2. remove_entry() will deserialize both the key and the value
43//! assert_eq!(new_tree.remove_entry(&15), Some((15,1.245)));
44//! // 2. delete() will not deserialize anything, will return `true` if the value was present
45//! assert_eq!(new_tree.delete(&19), true);
46//! ```
47//!
48//! # Internal structure
49//! Internally, [`RBTree`] is just a wrapper around [`RBForest`](super::forest::RBForest) with `max_roots`
50//! equal to `1`. See [`RBForest`](super::forest::RBForest) docs for description of the internals.
51use borsh::{BorshDeserialize, BorshSerialize};
52use core::borrow::Borrow;
53use core::cmp::Ord;
54use core::fmt;
55
56pub use super::forest::iterators::{KeysIterator, PairsIterator, ValuesIterator};
57use super::forest::{forest_size, init_forest, ForestParams, RBForest};
58use super::Error;
59
60/// Parameters required to calculate [`RBTree`] size
61#[derive(Debug, Eq, PartialEq, Copy, Clone)]
62pub struct TreeParams {
63    ///  key buffer size
64    pub k_size: usize,
65    ///  value buffer size
66    pub v_size: usize,
67}
68/// Returns the required size of the slice
69#[must_use]
70#[inline]
71pub const fn tree_size(params: TreeParams, max_nodes: usize) -> usize {
72    forest_size(
73        ForestParams {
74            k_size: params.k_size,
75            v_size: params.v_size,
76            max_roots: 1,
77        },
78        max_nodes,
79    )
80}
81
82/// Initializes [`RBTree`] in the given slice without returning it
83///
84/// This function can be used than you don't know buffer sizes at compile time.
85#[inline]
86pub fn init_tree(params: TreeParams, slice: &mut [u8]) -> Result<(), Error> {
87    init_forest(
88        ForestParams {
89            k_size: params.k_size,
90            v_size: params.v_size,
91            max_roots: 1,
92        },
93        slice,
94    )
95}
96
97/// A slice-based Red-Black tree
98///
99/// See [module](super::tree) level documentation for more info.
100pub struct RBTree<'a, K, V, const KSIZE: usize, const VSIZE: usize>(
101    RBForest<'a, K, V, KSIZE, VSIZE>,
102)
103where
104    K: Ord + BorshDeserialize + BorshSerialize,
105    V: BorshDeserialize + BorshSerialize;
106
107impl<'a, K, V, const KSIZE: usize, const VSIZE: usize> RBTree<'a, K, V, KSIZE, VSIZE>
108where
109    K: Ord + BorshDeserialize + BorshSerialize,
110    V: BorshDeserialize + BorshSerialize,
111{
112    /// Initializes [`RBTree`] in a given slice
113    pub fn init_slice(slice: &'a mut [u8]) -> Result<Self, Error> {
114        RBForest::<'a, K, V, KSIZE, VSIZE>::init_slice(slice, 1).map(|tree| Self(tree))
115    }
116
117    /// Returns [`RBTree`], contained in the given slice
118    ///
119    /// # Safety
120    /// This function must be called only on slices, previously initialized as [`RBTree`] using
121    /// [`init_tree`] or [`RBTree::init_slice`]
122    pub unsafe fn from_slice(slice: &'a mut [u8]) -> Result<Self, Error> {
123        unsafe { RBForest::<'a, K, V, KSIZE, VSIZE>::from_slice(slice).map(|tree| Self(tree)) }
124    }
125
126    /// Returns the number of occupied nodes
127    ///
128    /// This function runs in `O(n)`, where `n` - is the number of nodes
129    #[must_use]
130    pub fn len(&self) -> usize {
131        self.0.len(0).unwrap()
132    }
133
134    /// Clears the tree
135    ///
136    /// This function runs in `O(n)`, where `n` - is the number of nodes
137    pub fn clear(&mut self) {
138        self.0.clear();
139    }
140
141    /// Returns the number of free nodes
142    ///
143    /// This function runs in `O(n)`, where `n` - is the number of nodes
144    #[must_use]
145    pub fn free_nodes_left(&self) -> usize {
146        self.0.free_nodes_left()
147    }
148
149    /// Returns true if the map contains a value for the specified key
150    ///
151    /// This function runs in `O(log(n))`, where `n` - is the number of nodes
152    #[must_use]
153    pub fn contains_key<Q>(&self, k: &Q) -> bool
154    where
155        K: Borrow<Q> + Ord,
156        Q: Ord + ?Sized,
157    {
158        self.0.contains_key(0, k)
159    }
160
161    /// Returns a key-value pair corresponding to the supplied key
162    ///
163    /// This function runs in `O(log(n))`, where `n` - is the number of nodes
164    #[must_use]
165    pub fn get_entry<Q>(&self, k: &Q) -> Option<(K, V)>
166    where
167        K: Borrow<Q> + Ord,
168        Q: Ord + ?Sized,
169    {
170        self.0.get_entry(0, k)
171    }
172
173    /// Returns the value corresponding to the key
174    ///
175    /// This function runs in `O(log(n))`, where `n` - is the number of nodes
176    #[must_use]
177    pub fn get<Q>(&self, k: &Q) -> Option<V>
178    where
179        K: Borrow<Q> + Ord,
180        Q: Ord + ?Sized,
181    {
182        self.0.get(0, k)
183    }
184
185    /// Inserts a new key-value pair and returns the old value if it was present
186    ///
187    /// This function runs in `O(log(n))`, where `n` - is the number of nodes
188    pub fn insert(&mut self, k: K, v: V) -> Result<Option<V>, Error> {
189        self.0.insert(0, k, v)
190    }
191
192    /// Returns `true` if the tree contains no elements
193    #[must_use]
194    pub fn is_empty(&self) -> bool {
195        self.0.is_empty(0)
196    }
197
198    /// Deletes entry and returns deserialized value
199    ///
200    /// This function runs in `O(log(n))`, where `n` - is the number of nodes
201    pub fn remove<Q>(&mut self, key: &Q) -> Option<V>
202    where
203        K: Borrow<Q> + Ord,
204        Q: Ord + ?Sized,
205    {
206        self.0.remove(0, key)
207    }
208
209    /// Deletes entry and returns deserialized key-value pair
210    ///
211    /// This function runs in `O(log(n))`, where `n` - is the number of nodes
212    pub fn remove_entry<Q>(&mut self, key: &Q) -> Option<(K, V)>
213    where
214        K: Borrow<Q> + Ord,
215        Q: Ord + ?Sized,
216    {
217        self.0.remove_entry(0, key)
218    }
219
220    /// Deletes entry without deserializing the value
221    ///
222    /// Returns `true` if there was a value with the given key.
223    pub fn delete<Q>(&mut self, key: &Q) -> bool
224    where
225        K: Borrow<Q> + Ord,
226        Q: Ord + ?Sized,
227    {
228        self.0.delete(0, key)
229    }
230
231    /// Returns the first key-value pair in the map
232    ///
233    /// This function runs in `O(log(n))`, where `n` - is the number of nodes
234    #[must_use]
235    pub fn first_entry(&self) -> Option<(K, V)> {
236        self.0.first_entry(0)
237    }
238
239    /// Returns the last key-value pair in the map
240    ///
241    /// This function runs in `O(log(n))`, where `n` - is the number of nodes
242    #[must_use]
243    pub fn last_entry(&self) -> Option<(K, V)> {
244        self.0.last_entry(0)
245    }
246
247    /// Creates an iterator over key-value pairs, in order by key
248    #[must_use]
249    pub fn pairs<'b>(&'b self) -> PairsIterator<'b, 'a, K, V, KSIZE, VSIZE> {
250        self.0.pairs(0).unwrap()
251    }
252
253    /// Creates an iterator over keys, from smallest to biggest
254    #[must_use]
255    pub fn keys<'b>(&'b self) -> KeysIterator<'b, 'a, K, V, KSIZE, VSIZE> {
256        self.0.keys(0).unwrap()
257    }
258
259    /// Creates an iterator over values, in order by key
260    #[must_use]
261    pub fn values<'b>(&'b self) -> ValuesIterator<'b, 'a, K, V, KSIZE, VSIZE> {
262        self.0.values(0).unwrap()
263    }
264}
265
266impl<'a, K, V, const KSIZE: usize, const VSIZE: usize> fmt::Debug for RBTree<'a, K, V, KSIZE, VSIZE>
267where
268    K: Ord + BorshDeserialize + BorshSerialize + fmt::Debug,
269    V: BorshDeserialize + BorshSerialize + fmt::Debug,
270{
271    #[cfg(not(test))]
272    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
273        f.debug_map().entries(self.pairs()).finish()
274    }
275    #[cfg(test)]
276    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
277        self.0.fmt(f)
278    }
279}
280
281impl<'a, K, V, const KSIZE: usize, const VSIZE: usize> Extend<(K, V)>
282    for RBTree<'a, K, V, KSIZE, VSIZE>
283where
284    K: Ord + BorshDeserialize + BorshSerialize,
285    V: BorshDeserialize + BorshSerialize,
286{
287    fn extend<T: IntoIterator<Item = (K, V)>>(&mut self, iter: T) {
288        for elem in iter {
289            self.insert(elem.0, elem.1).unwrap();
290        }
291    }
292}
293
294#[cfg(test)]
295mod tests;
296
297#[cfg(any(test, fuzzing))]
298pub mod internal_checks;
slice_rbtree/tree.rs

slice_rbtree/
tree.rs
