multiversx_sc/storage/mappers/
ordered_binary_tree_mapper.rs1use 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#[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
58pub 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}