Skip to main content

multiversx_sc/storage/mappers/
ordered_binary_tree_mapper.rs

1use core::marker::PhantomData;
2
3use codec::Empty;
4
5use crate::{
6    api::{ErrorApiImpl, StorageMapperApi},
7    storage::StorageKey,
8    storage_set,
9    types::ManagedType,
10};
11
12use super::{
13    StorageMapper,
14    source::{CurrentStorage, StorageAddress},
15};
16
17use crate::codec::{
18    self, NestedDecode, NestedEncode,
19    derive::{TopDecode, TopEncode},
20};
21
22pub type NodeId = u64;
23
24pub const NULL_NODE_ID: NodeId = 0;
25
26const ROOT_ID_SUFFIX: &str = "_rootId";
27const ID_SUFFIX: &str = "_id";
28const LAST_ID_KEY_SUFFIX: &str = "_lastId";
29
30const CORRUPT_TREE_ERR_MGS: &str = "Corrupt tree";
31
32// https://en.wikipedia.org/wiki/Binary_search_tree
33
34#[derive(TopEncode, TopDecode, Clone, PartialEq, Debug)]
35pub struct OrderedBinaryTreeNode<T: NestedEncode + NestedDecode + PartialOrd + PartialEq + Clone> {
36    pub current_node_id: NodeId,
37    pub left_id: NodeId,
38    pub right_id: NodeId,
39    pub parent_id: NodeId,
40    pub data: T,
41}
42
43impl<T> OrderedBinaryTreeNode<T>
44where
45    T: NestedEncode + NestedDecode + PartialOrd + PartialEq + Clone,
46{
47    pub fn new(current_node_id: NodeId, data: T) -> Self {
48        Self {
49            data,
50            current_node_id,
51            left_id: NULL_NODE_ID,
52            right_id: NULL_NODE_ID,
53            parent_id: NULL_NODE_ID,
54        }
55    }
56}
57
58/// A storage mapper implementing a binary search tree for ordered element storage.
59/// **Note: This is NOT a self-balancing tree** and can degrade to O(n) in worst-case scenarios.
60///
61/// # Storage Layout
62///
63/// The `OrderedBinaryTreeMapper` stores nodes in a tree structure with parent/child relationships:
64///
65/// 1. **Root tracking**:
66///    - `base_key + "_rootId"` → ID of the root node
67///
68/// 2. **Node storage**:
69///    - `base_key + "_id" + node_id` → `OrderedBinaryTreeNode<T>` containing:
70///      - `current_node_id`: this node's ID
71///      - `left_id`: ID of left child (smaller values)
72///      - `right_id`: ID of right child (larger values)
73///      - `parent_id`: ID of parent node
74///      - `data`: the stored value
75///
76/// 3. **ID counter**:
77///    - `base_key + "_lastId"` → highest assigned node ID (auto-incrementing)
78///
79/// # Binary Search Tree Properties
80///
81/// - **Ordering**: For any node, all left descendants < node < all right descendants
82/// - **Uniqueness**: Duplicate values are rejected on insertion
83/// - **Node IDs**: Auto-incrementing, never reused (similar to `AddressToIdMapper`)
84///
85/// # Main Operations
86///
87/// - **Insert**: `insert_element(data)` - Adds element maintaining BST order. O(log n) average, O(n) worst case.
88/// - **Delete**: `delete_node(data)` - Removes element. O(log n) average, O(n) worst case.
89/// - **Search**: `iterative_search(node, data)` / `recursive_search(node, data)` - Finds element. O(log n) average.
90/// - **Min/Max**: `find_min(node)` / `find_max(node)` - Finds minimum/maximum in subtree. O(log n).
91/// - **Successor/Predecessor**: `find_successor(node)` / `find_predecessor(node)` - Next/previous in order. O(log n).
92/// - **Depth**: `get_depth(node)` - Computes tree depth from node. O(n).
93/// - **Root**: `get_root()` - Returns root node. O(1).
94///
95/// # Trade-offs
96///
97/// - **Pros**: Maintains sorted order; efficient ordered iteration; supports range queries; successor/predecessor lookups.
98/// - **Cons**: NOT self-balancing (can degrade to O(n) in worst case); complex deletion logic; higher storage overhead
99///   than simple collections; removed nodes leave ID gaps.
100///
101/// # Performance Notes
102///
103/// This is a **basic binary search tree**, not a balanced variant (e.g., AVL, Red-Black). Performance depends on
104/// insertion order:
105/// - **Balanced tree** (random insertions): O(log n) operations
106/// - **Degenerate tree** (sorted insertions): O(n) operations (becomes a linked list)
107///
108/// # Use Cases
109///
110/// - Leaderboards requiring sorted access
111/// - Priority systems with ordered iteration
112/// - Range queries (find all elements between X and Y)
113/// - Scenarios needing both membership testing and ordering
114/// - When you need successor/predecessor operations
115///
116/// # Example
117///
118/// ```rust
119/// # use multiversx_sc::storage::mappers::{StorageMapper, OrderedBinaryTreeMapper};
120/// # fn example<SA: multiversx_sc::api::StorageMapperApi>() {
121/// # let mut mapper = OrderedBinaryTreeMapper::<SA, u64>::new(
122/// #     multiversx_sc::storage::StorageKey::new(&b"scores"[..])
123/// # );
124/// // Insert elements (maintains BST ordering)
125/// let node1_id = mapper.insert_element(50);
126/// let node2_id = mapper.insert_element(30);
127/// let node3_id = mapper.insert_element(70);
128/// let node4_id = mapper.insert_element(20);
129/// let node5_id = mapper.insert_element(40);
130///
131/// // Duplicate insertion returns 0 (failure)
132/// let duplicate = mapper.insert_element(50);
133/// assert_eq!(duplicate, 0);
134///
135/// // Search for element
136/// let root = mapper.get_root().unwrap();
137/// let found = mapper.iterative_search(Some(root.clone()), &30);
138/// assert!(found.is_some());
139/// assert_eq!(found.unwrap().data, 30);
140///
141/// // Find minimum and maximum
142/// let min_node = mapper.find_min(root.clone());
143/// assert_eq!(min_node.data, 20);
144///
145/// let max_node = mapper.find_max(root.clone());
146/// assert_eq!(max_node.data, 70);
147///
148/// // Find successor (next larger element)
149/// let node_30 = mapper.iterative_search(Some(root.clone()), &30).unwrap();
150/// let successor = mapper.find_successor(node_30);
151/// assert_eq!(successor.unwrap().data, 40);
152///
153/// // Find predecessor (next smaller element)
154/// let node_50 = root.clone();
155/// let predecessor = mapper.find_predecessor(node_50);
156/// assert_eq!(predecessor.unwrap().data, 40);
157///
158/// // Get tree depth
159/// let depth = mapper.get_depth(&root);
160/// assert!(depth > 0);
161///
162/// // Delete element
163/// mapper.delete_node(30);
164/// let not_found = mapper.iterative_search(mapper.get_root(), &30);
165/// assert!(not_found.is_none());
166///
167/// // Tree structure is maintained after deletion
168/// let new_root = mapper.get_root().unwrap();
169/// let still_there = mapper.iterative_search(Some(new_root), &40);
170/// assert!(still_there.is_some());
171/// # }
172/// ```
173pub struct OrderedBinaryTreeMapper<SA, T, A = CurrentStorage>
174where
175    SA: StorageMapperApi,
176    A: StorageAddress<SA>,
177    T: NestedEncode + NestedDecode + PartialOrd + PartialEq + Clone,
178{
179    address: A,
180    key: StorageKey<SA>,
181    _phantom_api: PhantomData<SA>,
182    _phantom_item: PhantomData<T>,
183}
184
185impl<SA, T> StorageMapper<SA> for OrderedBinaryTreeMapper<SA, T, CurrentStorage>
186where
187    SA: StorageMapperApi,
188    T: NestedEncode + NestedDecode + PartialOrd + PartialEq + Clone + 'static,
189{
190    #[inline]
191    fn new(base_key: StorageKey<SA>) -> Self {
192        OrderedBinaryTreeMapper {
193            address: CurrentStorage,
194            key: base_key,
195            _phantom_api: PhantomData,
196            _phantom_item: PhantomData,
197        }
198    }
199}
200
201impl<SA, T, A> OrderedBinaryTreeMapper<SA, T, A>
202where
203    SA: StorageMapperApi,
204    A: StorageAddress<SA>,
205    T: NestedEncode + NestedDecode + PartialOrd + PartialEq + Clone,
206{
207    pub fn get_root(&self) -> Option<OrderedBinaryTreeNode<T>> {
208        let root_key = self.build_root_key();
209        let storage_len = self.address.address_storage_get_len(root_key.as_ref());
210        if storage_len == 0 {
211            return None;
212        }
213
214        Some(self.address.address_storage_get(root_key.as_ref()))
215    }
216
217    pub fn get_depth(&self, node: &OrderedBinaryTreeNode<T>) -> usize {
218        let opt_left_node = self.get_node_by_id(node.left_id);
219        let opt_right_node = self.get_node_by_id(node.right_id);
220
221        let l_depth = match opt_left_node {
222            Some(left_node) => self.get_depth(&left_node),
223            None => 0,
224        };
225        let r_depth = match opt_right_node {
226            Some(right_node) => self.get_depth(&right_node),
227            None => 0,
228        };
229
230        core::cmp::max(l_depth, r_depth) + 1
231    }
232
233    pub fn recursive_search(
234        &self,
235        opt_node: Option<OrderedBinaryTreeNode<T>>,
236        data: &T,
237    ) -> Option<OrderedBinaryTreeNode<T>> {
238        opt_node.as_ref()?;
239
240        let node = unsafe { opt_node.unwrap_unchecked() };
241        if &node.data == data {
242            return Some(node);
243        }
244
245        if data < &node.data {
246            let opt_left_node = self.get_node_by_id(node.left_id);
247            self.recursive_search(opt_left_node, data)
248        } else {
249            let opt_right_node = self.get_node_by_id(node.right_id);
250            self.recursive_search(opt_right_node, data)
251        }
252    }
253
254    pub fn iterative_search(
255        &self,
256        mut opt_node: Option<OrderedBinaryTreeNode<T>>,
257        data: &T,
258    ) -> Option<OrderedBinaryTreeNode<T>> {
259        while opt_node.is_some() {
260            let node = unsafe { opt_node.unwrap_unchecked() };
261            if &node.data == data {
262                return Some(node);
263            }
264
265            if data < &node.data {
266                opt_node = self.get_node_by_id(node.left_id);
267            } else {
268                opt_node = self.get_node_by_id(node.right_id);
269            }
270        }
271
272        None
273    }
274
275    pub fn find_max(&self, mut node: OrderedBinaryTreeNode<T>) -> OrderedBinaryTreeNode<T> {
276        while node.right_id != NULL_NODE_ID {
277            node = self.try_get_node_by_id(node.right_id);
278        }
279
280        node
281    }
282
283    pub fn find_min(&self, mut node: OrderedBinaryTreeNode<T>) -> OrderedBinaryTreeNode<T> {
284        while node.left_id != NULL_NODE_ID {
285            node = self.try_get_node_by_id(node.left_id);
286        }
287
288        node
289    }
290
291    pub fn find_successor(
292        &self,
293        mut node: OrderedBinaryTreeNode<T>,
294    ) -> Option<OrderedBinaryTreeNode<T>> {
295        if node.right_id != NULL_NODE_ID {
296            let right_node = self.try_get_node_by_id(node.right_id);
297            return Some(self.find_min(right_node));
298        }
299
300        let mut successor_id = node.parent_id;
301        let mut opt_successor = self.get_node_by_id(successor_id);
302        while successor_id != NULL_NODE_ID {
303            if opt_successor.is_none() {
304                SA::error_api_impl().signal_error(CORRUPT_TREE_ERR_MGS.as_bytes());
305            }
306
307            let successor = unsafe { opt_successor.unwrap_unchecked() };
308            if node.current_node_id != successor.right_id {
309                return Some(successor);
310            }
311
312            successor_id = successor.parent_id;
313            opt_successor = self.get_node_by_id(successor_id);
314            node = successor;
315        }
316
317        opt_successor
318    }
319
320    pub fn find_predecessor(
321        &self,
322        mut node: OrderedBinaryTreeNode<T>,
323    ) -> Option<OrderedBinaryTreeNode<T>> {
324        if node.left_id != NULL_NODE_ID {
325            let left_node = self.try_get_node_by_id(node.left_id);
326            return Some(self.find_max(left_node));
327        }
328
329        let mut predecessor_id = node.parent_id;
330        let mut opt_predecessor = self.get_node_by_id(predecessor_id);
331        while predecessor_id != NULL_NODE_ID {
332            if opt_predecessor.is_none() {
333                SA::error_api_impl().signal_error(CORRUPT_TREE_ERR_MGS.as_bytes());
334            }
335
336            let predecessor = unsafe { opt_predecessor.unwrap_unchecked() };
337            if node.current_node_id != predecessor.left_id {
338                return Some(predecessor);
339            }
340
341            predecessor_id = predecessor.parent_id;
342            opt_predecessor = self.get_node_by_id(predecessor_id);
343            node = predecessor;
344        }
345
346        opt_predecessor
347    }
348
349    pub fn insert_element(&mut self, new_data: T) -> u64 {
350        let new_node_id = self.get_and_increment_last_id();
351        let mut new_node = OrderedBinaryTreeNode::new(new_node_id, new_data);
352
353        let mut opt_new_node_parent = None;
354        let mut opt_current_node = self.get_root();
355        while opt_current_node.is_some() {
356            opt_new_node_parent.clone_from(&opt_current_node);
357
358            let current_node = unsafe { opt_current_node.unwrap_unchecked() };
359            if new_node.data == current_node.data {
360                return 0u64;
361            }
362
363            if new_node.data < current_node.data {
364                opt_current_node = self.get_node_by_id(current_node.left_id);
365            } else {
366                opt_current_node = self.get_node_by_id(current_node.right_id);
367            }
368        }
369
370        let new_node_parent_id = match &opt_new_node_parent {
371            Some(node) => node.current_node_id,
372            None => NULL_NODE_ID,
373        };
374        new_node.parent_id = new_node_parent_id;
375
376        if opt_new_node_parent.is_none() {
377            let root_id_key = self.build_root_id_key();
378            storage_set(root_id_key.as_ref(), &new_node.current_node_id);
379
380            let root_key = self.build_root_key();
381            storage_set(root_key.as_ref(), &new_node);
382
383            return 0u64;
384        }
385
386        let mut new_node_parent = unsafe { opt_new_node_parent.unwrap_unchecked() };
387        if new_node.data < new_node_parent.data {
388            new_node_parent.left_id = new_node.current_node_id;
389        } else {
390            new_node_parent.right_id = new_node.current_node_id;
391        }
392
393        self.set_item(new_node_id, &new_node);
394        self.set_item(new_node_parent.current_node_id, &new_node_parent);
395
396        new_node_id
397    }
398
399    pub fn delete_node(&mut self, data: T) {
400        let opt_root = self.get_root();
401        let opt_node = self.iterative_search(opt_root, &data);
402        if opt_node.is_none() {
403            SA::error_api_impl().signal_error(b"Node not found");
404        }
405
406        let node = unsafe { opt_node.unwrap_unchecked() };
407        if node.left_id == NULL_NODE_ID {
408            let opt_to_add = self.get_node_by_id(node.right_id);
409            self.shift_nodes(&node, opt_to_add);
410
411            return;
412        }
413
414        if node.right_id == NULL_NODE_ID {
415            let opt_to_add = self.get_node_by_id(node.left_id);
416            self.shift_nodes(&node, opt_to_add);
417
418            return;
419        }
420
421        let opt_successor = self.find_successor(node.clone());
422        if opt_successor.is_none() {
423            SA::error_api_impl().signal_error(CORRUPT_TREE_ERR_MGS.as_bytes());
424        }
425
426        let mut successor = unsafe { opt_successor.unwrap_unchecked() };
427        if successor.parent_id != node.current_node_id {
428            let opt_right = self.get_node_by_id(successor.right_id);
429            self.shift_nodes(&successor, opt_right);
430
431            successor = self.try_get_node_by_id(successor.current_node_id);
432            successor.right_id = node.right_id;
433
434            let opt_successor_right_node = self.get_node_by_id(successor.right_id);
435            if opt_successor_right_node.is_none() {
436                SA::error_api_impl().signal_error(CORRUPT_TREE_ERR_MGS.as_bytes());
437            }
438
439            let mut successor_right_node = unsafe { opt_successor_right_node.unwrap_unchecked() };
440            successor_right_node.parent_id = successor.current_node_id;
441
442            self.set_item(successor_right_node.current_node_id, &successor_right_node);
443        }
444
445        self.shift_nodes(&node, Some(successor.clone()));
446        successor = self.try_get_node_by_id(successor.current_node_id);
447        successor.left_id = node.left_id;
448
449        let opt_successor_left_node = self.get_node_by_id(successor.left_id);
450        if opt_successor_left_node.is_none() {
451            SA::error_api_impl().signal_error(CORRUPT_TREE_ERR_MGS.as_bytes());
452        }
453
454        let mut successor_left_node = unsafe { opt_successor_left_node.unwrap_unchecked() };
455        successor_left_node.parent_id = successor.current_node_id;
456
457        self.set_item(successor_left_node.current_node_id, &successor_left_node);
458        self.set_item(successor.current_node_id, &successor);
459    }
460
461    fn shift_nodes(
462        &mut self,
463        to_delete: &OrderedBinaryTreeNode<T>,
464        mut opt_to_add: Option<OrderedBinaryTreeNode<T>>,
465    ) {
466        if to_delete.parent_id == NULL_NODE_ID {
467            let root_id_key = self.build_root_id_key();
468            match &mut opt_to_add {
469                Some(to_add) => {
470                    to_add.parent_id = NULL_NODE_ID;
471                    storage_set(root_id_key.as_ref(), &to_add.current_node_id);
472
473                    let root_key = self.build_root_key();
474                    storage_set(root_key.as_ref(), to_add);
475                }
476                None => {
477                    let root_key = self.build_root_key();
478
479                    storage_set(root_id_key.as_ref(), &Empty);
480                    storage_set(root_key.as_ref(), &Empty);
481                }
482            };
483
484            return;
485        }
486
487        let to_add_id = match &opt_to_add {
488            Some(to_add) => to_add.current_node_id,
489            None => NULL_NODE_ID,
490        };
491
492        let mut parent = self.try_get_node_by_id(to_delete.parent_id);
493        if to_delete.current_node_id == parent.left_id {
494            parent.left_id = to_add_id;
495        } else {
496            parent.right_id = to_add_id;
497        }
498
499        if let Some(to_add) = &mut opt_to_add {
500            to_add.parent_id = to_delete.parent_id;
501
502            self.set_item(to_add.current_node_id, to_add);
503        }
504
505        self.set_item(parent.current_node_id, &parent);
506        self.clear_item(to_delete.current_node_id);
507    }
508}
509
510impl<SA, T, A> OrderedBinaryTreeMapper<SA, T, A>
511where
512    SA: StorageMapperApi,
513    A: StorageAddress<SA>,
514    T: NestedEncode + NestedDecode + PartialOrd + PartialEq + Clone,
515{
516    fn get_node_by_id(&self, id: NodeId) -> Option<OrderedBinaryTreeNode<T>> {
517        if id == NULL_NODE_ID {
518            return None;
519        }
520
521        let key = self.build_key_for_item(id);
522        let storage_len = self.address.address_storage_get_len(key.as_ref());
523        if storage_len == 0 {
524            return None;
525        }
526
527        Some(self.address.address_storage_get(key.as_ref()))
528    }
529
530    fn try_get_node_by_id(&self, id: NodeId) -> OrderedBinaryTreeNode<T> {
531        let opt_node = self.get_node_by_id(id);
532        if opt_node.is_none() {
533            SA::error_api_impl().signal_error(CORRUPT_TREE_ERR_MGS.as_bytes());
534        }
535
536        unsafe { opt_node.unwrap_unchecked() }
537    }
538
539    fn build_root_id_key(&self) -> StorageKey<SA> {
540        let mut key = self.key.clone();
541        key.append_bytes(ROOT_ID_SUFFIX.as_bytes());
542
543        key
544    }
545
546    fn build_root_key(&self) -> StorageKey<SA> {
547        let mut key = self.key.clone();
548        key.append_bytes(ROOT_ID_SUFFIX.as_bytes());
549
550        let root_id = self.address.address_storage_get(key.as_ref());
551
552        self.build_key_for_item(root_id)
553    }
554
555    fn build_key_for_item(&self, id: NodeId) -> StorageKey<SA> {
556        let mut item_key = self.key.clone();
557        item_key.append_bytes(ID_SUFFIX.as_bytes());
558        item_key.append_item(&id);
559
560        item_key
561    }
562
563    fn build_last_id_key(&self) -> StorageKey<SA> {
564        let mut key = self.key.clone();
565        key.append_bytes(LAST_ID_KEY_SUFFIX.as_bytes());
566
567        key
568    }
569
570    fn get_and_increment_last_id(&self) -> NodeId {
571        let key = self.build_last_id_key();
572        let last_id: NodeId = self.address.address_storage_get(key.as_ref());
573        let new_id = last_id + 1;
574        storage_set(key.as_ref(), &new_id);
575
576        new_id
577    }
578
579    fn set_item(&mut self, id: NodeId, node: &OrderedBinaryTreeNode<T>) {
580        let key = self.build_key_for_item(id);
581        storage_set(key.as_ref(), node);
582    }
583
584    fn clear_item(&mut self, id: NodeId) {
585        let key = self.build_key_for_item(id);
586        storage_set(key.as_ref(), &Empty);
587    }
588}