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
26static ROOT_ID_SUFFIX: &[u8] = b"_rootId";
27static ID_SUFFIX: &[u8] = b"_id";
28static LAST_ID_KEY_SUFFIX: &[u8] = b"_lastId";
29
30static CORRUPT_TREE_ERR_MGS: &[u8] = b"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>
59where
60 SA: StorageMapperApi,
61 A: StorageAddress<SA>,
62 T: NestedEncode + NestedDecode + PartialOrd + PartialEq + Clone,
63{
64 address: A,
65 key: StorageKey<SA>,
66 _phantom_api: PhantomData<SA>,
67 _phantom_item: PhantomData<T>,
68}
69
70impl<SA, T> StorageMapper<SA> for OrderedBinaryTreeMapper<SA, T, CurrentStorage>
71where
72 SA: StorageMapperApi,
73 T: NestedEncode + NestedDecode + PartialOrd + PartialEq + Clone + 'static,
74{
75 #[inline]
76 fn new(base_key: StorageKey<SA>) -> Self {
77 OrderedBinaryTreeMapper {
78 address: CurrentStorage,
79 key: base_key,
80 _phantom_api: PhantomData,
81 _phantom_item: PhantomData,
82 }
83 }
84}
85
86impl<SA, T, A> OrderedBinaryTreeMapper<SA, T, A>
87where
88 SA: StorageMapperApi,
89 A: StorageAddress<SA>,
90 T: NestedEncode + NestedDecode + PartialOrd + PartialEq + Clone,
91{
92 pub fn get_root(&self) -> Option<OrderedBinaryTreeNode<T>> {
93 let root_key = self.build_root_key();
94 let storage_len = self.address.address_storage_get_len(root_key.as_ref());
95 if storage_len == 0 {
96 return None;
97 }
98
99 Some(self.address.address_storage_get(root_key.as_ref()))
100 }
101
102 pub fn get_depth(&self, node: &OrderedBinaryTreeNode<T>) -> usize {
103 let opt_left_node = self.get_node_by_id(node.left_id);
104 let opt_right_node = self.get_node_by_id(node.right_id);
105
106 let l_depth = match opt_left_node {
107 Some(left_node) => self.get_depth(&left_node),
108 None => 0,
109 };
110 let r_depth = match opt_right_node {
111 Some(right_node) => self.get_depth(&right_node),
112 None => 0,
113 };
114
115 core::cmp::max(l_depth, r_depth) + 1
116 }
117
118 pub fn recursive_search(
119 &self,
120 opt_node: Option<OrderedBinaryTreeNode<T>>,
121 data: &T,
122 ) -> Option<OrderedBinaryTreeNode<T>> {
123 opt_node.as_ref()?;
124
125 let node = unsafe { opt_node.unwrap_unchecked() };
126 if &node.data == data {
127 return Some(node);
128 }
129
130 if data < &node.data {
131 let opt_left_node = self.get_node_by_id(node.left_id);
132 self.recursive_search(opt_left_node, data)
133 } else {
134 let opt_right_node = self.get_node_by_id(node.right_id);
135 self.recursive_search(opt_right_node, data)
136 }
137 }
138
139 pub fn iterative_search(
140 &self,
141 mut opt_node: Option<OrderedBinaryTreeNode<T>>,
142 data: &T,
143 ) -> Option<OrderedBinaryTreeNode<T>> {
144 while opt_node.is_some() {
145 let node = unsafe { opt_node.unwrap_unchecked() };
146 if &node.data == data {
147 return Some(node);
148 }
149
150 if data < &node.data {
151 opt_node = self.get_node_by_id(node.left_id);
152 } else {
153 opt_node = self.get_node_by_id(node.right_id);
154 }
155 }
156
157 None
158 }
159
160 pub fn find_max(&self, mut node: OrderedBinaryTreeNode<T>) -> OrderedBinaryTreeNode<T> {
161 while node.right_id != NULL_NODE_ID {
162 node = self.try_get_node_by_id(node.right_id);
163 }
164
165 node
166 }
167
168 pub fn find_min(&self, mut node: OrderedBinaryTreeNode<T>) -> OrderedBinaryTreeNode<T> {
169 while node.left_id != NULL_NODE_ID {
170 node = self.try_get_node_by_id(node.left_id);
171 }
172
173 node
174 }
175
176 pub fn find_successor(
177 &self,
178 mut node: OrderedBinaryTreeNode<T>,
179 ) -> Option<OrderedBinaryTreeNode<T>> {
180 if node.right_id != NULL_NODE_ID {
181 let right_node = self.try_get_node_by_id(node.right_id);
182 return Some(self.find_min(right_node));
183 }
184
185 let mut successor_id = node.parent_id;
186 let mut opt_successor = self.get_node_by_id(successor_id);
187 while successor_id != NULL_NODE_ID {
188 if opt_successor.is_none() {
189 SA::error_api_impl().signal_error(CORRUPT_TREE_ERR_MGS);
190 }
191
192 let successor = unsafe { opt_successor.unwrap_unchecked() };
193 if node.current_node_id != successor.right_id {
194 return Some(successor);
195 }
196
197 successor_id = successor.parent_id;
198 opt_successor = self.get_node_by_id(successor_id);
199 node = successor;
200 }
201
202 opt_successor
203 }
204
205 pub fn find_predecessor(
206 &self,
207 mut node: OrderedBinaryTreeNode<T>,
208 ) -> Option<OrderedBinaryTreeNode<T>> {
209 if node.left_id != NULL_NODE_ID {
210 let left_node = self.try_get_node_by_id(node.left_id);
211 return Some(self.find_max(left_node));
212 }
213
214 let mut predecessor_id = node.parent_id;
215 let mut opt_predecessor = self.get_node_by_id(predecessor_id);
216 while predecessor_id != NULL_NODE_ID {
217 if opt_predecessor.is_none() {
218 SA::error_api_impl().signal_error(CORRUPT_TREE_ERR_MGS);
219 }
220
221 let predecessor = unsafe { opt_predecessor.unwrap_unchecked() };
222 if node.current_node_id != predecessor.left_id {
223 return Some(predecessor);
224 }
225
226 predecessor_id = predecessor.parent_id;
227 opt_predecessor = self.get_node_by_id(predecessor_id);
228 node = predecessor;
229 }
230
231 opt_predecessor
232 }
233
234 pub fn insert_element(&mut self, new_data: T) -> u64 {
235 let new_node_id = self.get_and_increment_last_id();
236 let mut new_node = OrderedBinaryTreeNode::new(new_node_id, new_data);
237
238 let mut opt_new_node_parent = None;
239 let mut opt_current_node = self.get_root();
240 while opt_current_node.is_some() {
241 opt_new_node_parent.clone_from(&opt_current_node);
242
243 let current_node = unsafe { opt_current_node.unwrap_unchecked() };
244 if new_node.data == current_node.data {
245 return 0u64;
246 }
247
248 if new_node.data < current_node.data {
249 opt_current_node = self.get_node_by_id(current_node.left_id);
250 } else {
251 opt_current_node = self.get_node_by_id(current_node.right_id);
252 }
253 }
254
255 let new_node_parent_id = match &opt_new_node_parent {
256 Some(node) => node.current_node_id,
257 None => NULL_NODE_ID,
258 };
259 new_node.parent_id = new_node_parent_id;
260
261 if opt_new_node_parent.is_none() {
262 let root_id_key = self.build_root_id_key();
263 storage_set(root_id_key.as_ref(), &new_node.current_node_id);
264
265 let root_key = self.build_root_key();
266 storage_set(root_key.as_ref(), &new_node);
267
268 return 0u64;
269 }
270
271 let mut new_node_parent = unsafe { opt_new_node_parent.unwrap_unchecked() };
272 if new_node.data < new_node_parent.data {
273 new_node_parent.left_id = new_node.current_node_id;
274 } else {
275 new_node_parent.right_id = new_node.current_node_id;
276 }
277
278 self.set_item(new_node_id, &new_node);
279 self.set_item(new_node_parent.current_node_id, &new_node_parent);
280
281 new_node_id
282 }
283
284 pub fn delete_node(&mut self, data: T) {
285 let opt_root = self.get_root();
286 let opt_node = self.iterative_search(opt_root, &data);
287 if opt_node.is_none() {
288 SA::error_api_impl().signal_error(b"Node not found");
289 }
290
291 let node = unsafe { opt_node.unwrap_unchecked() };
292 if node.left_id == NULL_NODE_ID {
293 let opt_to_add = self.get_node_by_id(node.right_id);
294 self.shift_nodes(&node, opt_to_add);
295
296 return;
297 }
298
299 if node.right_id == NULL_NODE_ID {
300 let opt_to_add = self.get_node_by_id(node.left_id);
301 self.shift_nodes(&node, opt_to_add);
302
303 return;
304 }
305
306 let opt_successor = self.find_successor(node.clone());
307 if opt_successor.is_none() {
308 SA::error_api_impl().signal_error(CORRUPT_TREE_ERR_MGS);
309 }
310
311 let mut successor = unsafe { opt_successor.unwrap_unchecked() };
312 if successor.parent_id != node.current_node_id {
313 let opt_right = self.get_node_by_id(successor.right_id);
314 self.shift_nodes(&successor, opt_right);
315
316 successor = self.try_get_node_by_id(successor.current_node_id);
317 successor.right_id = node.right_id;
318
319 let opt_successor_right_node = self.get_node_by_id(successor.right_id);
320 if opt_successor_right_node.is_none() {
321 SA::error_api_impl().signal_error(CORRUPT_TREE_ERR_MGS);
322 }
323
324 let mut successor_right_node = unsafe { opt_successor_right_node.unwrap_unchecked() };
325 successor_right_node.parent_id = successor.current_node_id;
326
327 self.set_item(successor_right_node.current_node_id, &successor_right_node);
328 }
329
330 self.shift_nodes(&node, Some(successor.clone()));
331 successor = self.try_get_node_by_id(successor.current_node_id);
332 successor.left_id = node.left_id;
333
334 let opt_successor_left_node = self.get_node_by_id(successor.left_id);
335 if opt_successor_left_node.is_none() {
336 SA::error_api_impl().signal_error(CORRUPT_TREE_ERR_MGS);
337 }
338
339 let mut successor_left_node = unsafe { opt_successor_left_node.unwrap_unchecked() };
340 successor_left_node.parent_id = successor.current_node_id;
341
342 self.set_item(successor_left_node.current_node_id, &successor_left_node);
343 self.set_item(successor.current_node_id, &successor);
344 }
345
346 fn shift_nodes(
347 &mut self,
348 to_delete: &OrderedBinaryTreeNode<T>,
349 mut opt_to_add: Option<OrderedBinaryTreeNode<T>>,
350 ) {
351 if to_delete.parent_id == NULL_NODE_ID {
352 let root_id_key = self.build_root_id_key();
353 match &mut opt_to_add {
354 Some(to_add) => {
355 to_add.parent_id = NULL_NODE_ID;
356 storage_set(root_id_key.as_ref(), &to_add.current_node_id);
357
358 let root_key = self.build_root_key();
359 storage_set(root_key.as_ref(), to_add);
360 }
361 None => {
362 let root_key = self.build_root_key();
363
364 storage_set(root_id_key.as_ref(), &Empty);
365 storage_set(root_key.as_ref(), &Empty);
366 }
367 };
368
369 return;
370 }
371
372 let to_add_id = match &opt_to_add {
373 Some(to_add) => to_add.current_node_id,
374 None => NULL_NODE_ID,
375 };
376
377 let mut parent = self.try_get_node_by_id(to_delete.parent_id);
378 if to_delete.current_node_id == parent.left_id {
379 parent.left_id = to_add_id;
380 } else {
381 parent.right_id = to_add_id;
382 }
383
384 if let Some(to_add) = &mut opt_to_add {
385 to_add.parent_id = to_delete.parent_id;
386
387 self.set_item(to_add.current_node_id, to_add);
388 }
389
390 self.set_item(parent.current_node_id, &parent);
391 self.clear_item(to_delete.current_node_id);
392 }
393}
394
395impl<SA, T, A> OrderedBinaryTreeMapper<SA, T, A>
396where
397 SA: StorageMapperApi,
398 A: StorageAddress<SA>,
399 T: NestedEncode + NestedDecode + PartialOrd + PartialEq + Clone,
400{
401 fn get_node_by_id(&self, id: NodeId) -> Option<OrderedBinaryTreeNode<T>> {
402 if id == NULL_NODE_ID {
403 return None;
404 }
405
406 let key = self.build_key_for_item(id);
407 let storage_len = self.address.address_storage_get_len(key.as_ref());
408 if storage_len == 0 {
409 return None;
410 }
411
412 Some(self.address.address_storage_get(key.as_ref()))
413 }
414
415 fn try_get_node_by_id(&self, id: NodeId) -> OrderedBinaryTreeNode<T> {
416 let opt_node = self.get_node_by_id(id);
417 if opt_node.is_none() {
418 SA::error_api_impl().signal_error(CORRUPT_TREE_ERR_MGS);
419 }
420
421 unsafe { opt_node.unwrap_unchecked() }
422 }
423
424 fn build_root_id_key(&self) -> StorageKey<SA> {
425 let mut key = self.key.clone();
426 key.append_bytes(ROOT_ID_SUFFIX);
427
428 key
429 }
430
431 fn build_root_key(&self) -> StorageKey<SA> {
432 let mut key = self.key.clone();
433 key.append_bytes(ROOT_ID_SUFFIX);
434
435 let root_id = self.address.address_storage_get(key.as_ref());
436
437 self.build_key_for_item(root_id)
438 }
439
440 fn build_key_for_item(&self, id: NodeId) -> StorageKey<SA> {
441 let mut item_key = self.key.clone();
442 item_key.append_bytes(ID_SUFFIX);
443 item_key.append_item(&id);
444
445 item_key
446 }
447
448 fn build_last_id_key(&self) -> StorageKey<SA> {
449 let mut key = self.key.clone();
450 key.append_bytes(LAST_ID_KEY_SUFFIX);
451
452 key
453 }
454
455 fn get_and_increment_last_id(&self) -> NodeId {
456 let key = self.build_last_id_key();
457 let last_id: NodeId = self.address.address_storage_get(key.as_ref());
458 let new_id = last_id + 1;
459 storage_set(key.as_ref(), &new_id);
460
461 new_id
462 }
463
464 fn set_item(&mut self, id: NodeId, node: &OrderedBinaryTreeNode<T>) {
465 let key = self.build_key_for_item(id);
466 storage_set(key.as_ref(), node);
467 }
468
469 fn clear_item(&mut self, id: NodeId) {
470 let key = self.build_key_for_item(id);
471 storage_set(key.as_ref(), &Empty);
472 }
473}