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