unc_sdk/collections/
legacy_tree_map.rs

1//! Legacy `TreeMap` implementation that is using `UnorderedMap`.
2//! DEPRECATED. This implementation is deprecated and may be removed in the future.
3#![allow(clippy::all)]
4// This suppresses the depreciation warnings for uses of LegacyTreeMap in this module
5#![allow(deprecated)]
6
7use borsh::{BorshDeserialize, BorshSerialize};
8use std::ops::Bound;
9use unc_sdk_macros::unc;
10
11use crate::collections::UnorderedMap;
12use crate::collections::{append, Vector};
13use crate::IntoStorageKey;
14
15/// TreeMap based on AVL-tree
16///
17/// Runtime complexity (worst case):
18/// - `get`/`contains_key`:     O(1) - UnorderedMap lookup
19/// - `insert`/`remove`:        O(log(N))
20/// - `min`/`max`:              O(log(N))
21/// - `above`/`below`:          O(log(N))
22/// - `range` of K elements:    O(Klog(N))
23///
24#[deprecated(since = "2.1.0", note = "Use unc_sdk::collections::TreeMap")]
25#[unc(inside_uncsdk)]
26pub struct LegacyTreeMap<K, V> {
27    root: u64,
28    // ser/de is independent of `K`,`V` ser/de, `BorshSerialize`/`BorshDeserialize`/`BorshSchema` bounds removed
29    #[cfg_attr(not(feature = "abi"), borsh(bound(serialize = "", deserialize = "")))]
30    #[cfg_attr(
31        feature = "abi",
32        borsh(bound(serialize = "", deserialize = ""), schema(params = ""))
33    )]
34    val: UnorderedMap<K, V>,
35    // ser/de is independent of `K` ser/de, `BorshSerialize`/`BorshDeserialize`/`BorshSchema` bounds removed
36    #[cfg_attr(not(feature = "abi"), borsh(bound(serialize = "", deserialize = "")))]
37    #[cfg_attr(
38        feature = "abi",
39        borsh(bound(serialize = "", deserialize = ""), schema(params = ""))
40    )]
41    tree: Vector<Node<K>>,
42}
43
44#[unc(inside_uncsdk)]
45pub struct Node<K> {
46    id: u64,
47    key: K,           // key stored in a node
48    lft: Option<u64>, // left link of a node
49    rgt: Option<u64>, // right link of a node
50    ht: u64,          // height of a subtree at a node
51}
52
53impl<K> Node<K>
54where
55    K: Ord + Clone + BorshSerialize + BorshDeserialize,
56{
57    fn of(id: u64, key: K) -> Self {
58        Self { id, key, lft: None, rgt: None, ht: 1 }
59    }
60}
61
62#[allow(clippy::len_without_is_empty)]
63impl<K, V> LegacyTreeMap<K, V>
64where
65    K: Ord + Clone + BorshSerialize + BorshDeserialize,
66    V: BorshSerialize + BorshDeserialize,
67{
68    pub fn new<S>(prefix: S) -> Self
69    where
70        S: IntoStorageKey,
71    {
72        let prefix = prefix.into_storage_key();
73        Self {
74            root: 0,
75            val: UnorderedMap::new(append(&prefix, b'v')),
76            tree: Vector::new(append(&prefix, b'n')),
77        }
78    }
79
80    #[allow(clippy::len_without_is_empty)]
81    pub fn len(&self) -> u64 {
82        self.tree.len()
83    }
84
85    pub fn clear(&mut self) {
86        self.root = 0;
87        self.val.clear();
88        self.tree.clear();
89    }
90
91    fn node(&self, id: u64) -> Option<Node<K>> {
92        self.tree.get(id)
93    }
94
95    fn save(&mut self, node: &Node<K>) {
96        if node.id < self.len() {
97            self.tree.replace(node.id, node);
98        } else {
99            self.tree.push(node);
100        }
101    }
102
103    pub fn contains_key(&self, key: &K) -> bool {
104        self.val.get(key).is_some()
105    }
106
107    pub fn get(&self, key: &K) -> Option<V> {
108        self.val.get(key)
109    }
110
111    pub fn insert(&mut self, key: &K, val: &V) -> Option<V> {
112        if !self.contains_key(&key) {
113            self.root = self.insert_at(self.root, self.len(), &key);
114        }
115        self.val.insert(&key, &val)
116    }
117
118    pub fn remove(&mut self, key: &K) -> Option<V> {
119        if self.contains_key(&key) {
120            self.root = self.do_remove(&key);
121            self.val.remove(&key)
122        } else {
123            // no such key, nothing to do
124            None
125        }
126    }
127
128    /// Returns the smallest stored key from the tree
129    pub fn min(&self) -> Option<K> {
130        self.min_at(self.root, self.root).map(|(n, _)| n.key)
131    }
132
133    /// Returns the largest stored key from the tree
134    pub fn max(&self) -> Option<K> {
135        self.max_at(self.root, self.root).map(|(n, _)| n.key)
136    }
137
138    /// Returns the smallest key that is strictly greater than key given as the parameter
139    pub fn higher(&self, key: &K) -> Option<K> {
140        self.above_at(self.root, key)
141    }
142
143    /// Returns the largest key that is strictly less than key given as the parameter
144    pub fn lower(&self, key: &K) -> Option<K> {
145        self.below_at(self.root, key)
146    }
147
148    /// Returns the smallest key that is greater or equal to key given as the parameter
149    pub fn ceil_key(&self, key: &K) -> Option<K> {
150        if self.contains_key(key) {
151            Some(key.clone())
152        } else {
153            self.higher(key)
154        }
155    }
156
157    /// Returns the largest key that is less or equal to key given as the parameter
158    pub fn floor_key(&self, key: &K) -> Option<K> {
159        if self.contains_key(key) {
160            Some(key.clone())
161        } else {
162            self.lower(key)
163        }
164    }
165
166    /// Iterate all entries in ascending order: min to max, both inclusive
167    pub fn iter(&self) -> impl Iterator<Item = (K, V)> + '_ {
168        Cursor::asc(&self)
169    }
170
171    /// Iterate entries in ascending order: given key (exclusive) to max (inclusive)
172    pub fn iter_from(&self, key: K) -> impl Iterator<Item = (K, V)> + '_ {
173        Cursor::asc_from(&self, key)
174    }
175
176    /// Iterate all entries in descending order: max to min, both inclusive
177    pub fn iter_rev(&self) -> impl Iterator<Item = (K, V)> + '_ {
178        Cursor::desc(&self)
179    }
180
181    /// Iterate entries in descending order: given key (exclusive) to min (inclusive)
182    pub fn iter_rev_from(&self, key: K) -> impl Iterator<Item = (K, V)> + '_ {
183        Cursor::desc_from(&self, key)
184    }
185
186    /// Iterate entries in ascending order according to specified bounds.
187    ///
188    /// # Panics
189    ///
190    /// Panics if range start > end.
191    /// Panics if range start == end and both bounds are Excluded.
192    pub fn range(&self, r: (Bound<K>, Bound<K>)) -> impl Iterator<Item = (K, V)> + '_ {
193        let (lo, hi) = match r {
194            (Bound::Included(a), Bound::Included(b)) if a > b => panic!("Invalid range."),
195            (Bound::Excluded(a), Bound::Included(b)) if a > b => panic!("Invalid range."),
196            (Bound::Included(a), Bound::Excluded(b)) if a > b => panic!("Invalid range."),
197            (Bound::Excluded(a), Bound::Excluded(b)) if a >= b => panic!("Invalid range."),
198            (lo, hi) => (lo, hi),
199        };
200
201        Cursor::range(&self, lo, hi)
202    }
203
204    /// Helper function which creates a [`Vec<(K, V)>`] of all items in the [`LegacyTreeMap`].
205    /// This function collects elements from [`LegacyTreeMap::iter`].
206    pub fn to_vec(&self) -> Vec<(K, V)> {
207        self.iter().collect()
208    }
209
210    //
211    // Internal utilities
212    //
213
214    /// Returns (node, parent node) of left-most lower (min) node starting from given node `at`.
215    /// As min_at only traverses the tree down, if a node `at` is the minimum node in a subtree,
216    /// its parent must be explicitly provided in advance.
217    fn min_at(&self, mut at: u64, p: u64) -> Option<(Node<K>, Node<K>)> {
218        let mut parent: Option<Node<K>> = self.node(p);
219        loop {
220            let node = self.node(at);
221            match node.as_ref().and_then(|n| n.lft) {
222                Some(lft) => {
223                    at = lft;
224                    parent = node;
225                }
226                None => {
227                    return node.and_then(|n| parent.map(|p| (n, p)));
228                }
229            }
230        }
231    }
232
233    /// Returns (node, parent node) of right-most lower (max) node starting from given node `at`.
234    /// As min_at only traverses the tree down, if a node `at` is the minimum node in a subtree,
235    /// its parent must be explicitly provided in advance.
236    fn max_at(&self, mut at: u64, p: u64) -> Option<(Node<K>, Node<K>)> {
237        let mut parent: Option<Node<K>> = self.node(p);
238        loop {
239            let node = self.node(at);
240            match node.as_ref().and_then(|n| n.rgt) {
241                Some(rgt) => {
242                    parent = node;
243                    at = rgt;
244                }
245                None => {
246                    return node.and_then(|n| parent.map(|p| (n, p)));
247                }
248            }
249        }
250    }
251
252    fn above_at(&self, mut at: u64, key: &K) -> Option<K> {
253        let mut seen: Option<K> = None;
254        loop {
255            let node = self.node(at);
256            match node.as_ref().map(|n| &n.key) {
257                Some(k) => {
258                    if k.le(key) {
259                        match node.and_then(|n| n.rgt) {
260                            Some(rgt) => at = rgt,
261                            None => break,
262                        }
263                    } else {
264                        seen = Some(k.clone());
265                        match node.and_then(|n| n.lft) {
266                            Some(lft) => at = lft,
267                            None => break,
268                        }
269                    }
270                }
271                None => break,
272            }
273        }
274        seen
275    }
276
277    fn below_at(&self, mut at: u64, key: &K) -> Option<K> {
278        let mut seen: Option<K> = None;
279        loop {
280            let node = self.node(at);
281            match node.as_ref().map(|n| &n.key) {
282                Some(k) => {
283                    if k.lt(key) {
284                        seen = Some(k.clone());
285                        match node.and_then(|n| n.rgt) {
286                            Some(rgt) => at = rgt,
287                            None => break,
288                        }
289                    } else {
290                        match node.and_then(|n| n.lft) {
291                            Some(lft) => at = lft,
292                            None => break,
293                        }
294                    }
295                }
296                None => break,
297            }
298        }
299        seen
300    }
301
302    fn insert_at(&mut self, at: u64, id: u64, key: &K) -> u64 {
303        match self.node(at) {
304            None => {
305                self.save(&Node::of(id, key.clone()));
306                at
307            }
308            Some(mut node) => {
309                if key.eq(&node.key) {
310                    at
311                } else {
312                    if key.lt(&node.key) {
313                        let idx = match node.lft {
314                            Some(lft) => self.insert_at(lft, id, key),
315                            None => self.insert_at(id, id, key),
316                        };
317                        node.lft = Some(idx);
318                    } else {
319                        let idx = match node.rgt {
320                            Some(rgt) => self.insert_at(rgt, id, key),
321                            None => self.insert_at(id, id, key),
322                        };
323                        node.rgt = Some(idx);
324                    };
325
326                    self.update_height(&mut node);
327                    self.enforce_balance(&mut node)
328                }
329            }
330        }
331    }
332
333    // Calculate and save the height of a subtree at node `at`:
334    // height[at] = 1 + max(height[at.L], height[at.R])
335    fn update_height(&mut self, node: &mut Node<K>) {
336        let lft = node.lft.and_then(|id| self.node(id).map(|n| n.ht)).unwrap_or_default();
337        let rgt = node.rgt.and_then(|id| self.node(id).map(|n| n.ht)).unwrap_or_default();
338
339        node.ht = 1 + std::cmp::max(lft, rgt);
340        self.save(&node);
341    }
342
343    // Balance = difference in heights between left and right subtrees at given node.
344    fn get_balance(&self, node: &Node<K>) -> i64 {
345        let lht = node.lft.and_then(|id| self.node(id).map(|n| n.ht)).unwrap_or_default();
346        let rht = node.rgt.and_then(|id| self.node(id).map(|n| n.ht)).unwrap_or_default();
347
348        lht as i64 - rht as i64
349    }
350
351    // Left rotation of an AVL subtree with at node `at`.
352    // New root of subtree is returned, caller is responsible for updating proper link from parent.
353    fn rotate_left(&mut self, node: &mut Node<K>) -> u64 {
354        let mut lft = node.lft.and_then(|id| self.node(id)).unwrap();
355        let lft_rgt = lft.rgt;
356
357        // at.L = at.L.R
358        node.lft = lft_rgt;
359
360        // at.L.R = at
361        lft.rgt = Some(node.id);
362
363        // at = at.L
364        self.update_height(node);
365        self.update_height(&mut lft);
366
367        lft.id
368    }
369
370    // Right rotation of an AVL subtree at node in `at`.
371    // New root of subtree is returned, caller is responsible for updating proper link from parent.
372    fn rotate_right(&mut self, node: &mut Node<K>) -> u64 {
373        let mut rgt = node.rgt.and_then(|id| self.node(id)).unwrap();
374        let rgt_lft = rgt.lft;
375
376        // at.R = at.R.L
377        node.rgt = rgt_lft;
378
379        // at.R.L = at
380        rgt.lft = Some(node.id);
381
382        // at = at.R
383        self.update_height(node);
384        self.update_height(&mut rgt);
385
386        rgt.id
387    }
388
389    // Check balance at a given node and enforce it if necessary with respective rotations.
390    fn enforce_balance(&mut self, node: &mut Node<K>) -> u64 {
391        let balance = self.get_balance(&node);
392        if balance > 1 {
393            let mut lft = node.lft.and_then(|id| self.node(id)).unwrap();
394            if self.get_balance(&lft) < 0 {
395                let rotated = self.rotate_right(&mut lft);
396                node.lft = Some(rotated);
397            }
398            self.rotate_left(node)
399        } else if balance < -1 {
400            let mut rgt = node.rgt.and_then(|id| self.node(id)).unwrap();
401            if self.get_balance(&rgt) > 0 {
402                let rotated = self.rotate_left(&mut rgt);
403                node.rgt = Some(rotated);
404            }
405            self.rotate_right(node)
406        } else {
407            node.id
408        }
409    }
410
411    // Returns (node, parent node) for a node that holds the `key`.
412    // For root node, same node is returned for node and parent node.
413    fn lookup_at(&self, mut at: u64, key: &K) -> Option<(Node<K>, Node<K>)> {
414        let mut p: Node<K> = self.node(at).unwrap();
415        while let Some(node) = self.node(at) {
416            if node.key.eq(key) {
417                return Some((node, p));
418            } else if node.key.lt(key) {
419                match node.rgt {
420                    Some(rgt) => {
421                        p = node;
422                        at = rgt;
423                    }
424                    None => break,
425                }
426            } else {
427                match node.lft {
428                    Some(lft) => {
429                        p = node;
430                        at = lft;
431                    }
432                    None => break,
433                }
434            }
435        }
436        None
437    }
438
439    // Navigate from root to node holding `key` and backtrace back to the root
440    // enforcing balance (if necessary) along the way.
441    fn check_balance(&mut self, at: u64, key: &K) -> u64 {
442        #[allow(clippy::branches_sharing_code)]
443        match self.node(at) {
444            Some(mut node) => {
445                if node.key.eq(key) {
446                    self.update_height(&mut node);
447                    self.enforce_balance(&mut node)
448                } else {
449                    if node.key.gt(key) {
450                        if let Some(l) = node.lft {
451                            let id = self.check_balance(l, key);
452                            node.lft = Some(id);
453                        }
454                    } else if let Some(r) = node.rgt {
455                        let id = self.check_balance(r, key);
456                        node.rgt = Some(id);
457                    }
458                    self.update_height(&mut node);
459                    self.enforce_balance(&mut node)
460                }
461            }
462            None => at,
463        }
464    }
465
466    // Node holding the key is not removed from the tree - instead the substitute node is found,
467    // the key is copied to 'removed' node from substitute node, and then substitute node gets
468    // removed from the tree.
469    //
470    // The substitute node is either:
471    // - right-most (max) node of the left subtree (containing smaller keys) of node holding `key`
472    // - or left-most (min) node of the right subtree (containing larger keys) of node holding `key`
473    //
474    fn do_remove(&mut self, key: &K) -> u64 {
475        // r_node - node containing key of interest
476        // p_node - immediate parent node of r_node
477        let (mut r_node, mut p_node) = match self.lookup_at(self.root, key) {
478            Some(x) => x,
479            None => return self.root, // cannot remove a missing key, no changes to the tree needed
480        };
481
482        let lft_opt = r_node.lft;
483        let rgt_opt = r_node.rgt;
484
485        if lft_opt.is_none() && rgt_opt.is_none() {
486            // remove leaf
487            if p_node.key.lt(key) {
488                p_node.rgt = None;
489            } else {
490                p_node.lft = None;
491            }
492            self.update_height(&mut p_node);
493
494            self.swap_with_last(r_node.id);
495
496            // removing node might have caused a imbalance - balance the tree up to the root,
497            // starting from lowest affected key - the parent of a leaf node in this case
498            self.check_balance(self.root, &p_node.key)
499        } else {
500            // non-leaf node, select subtree to proceed with
501            let b = self.get_balance(&r_node);
502            if b >= 0 {
503                // proceed with left subtree
504                let lft = lft_opt.unwrap();
505
506                // k - max key from left subtree
507                // n - node that holds key k, p - immediate parent of n
508                let (n, mut p) = self.max_at(lft, r_node.id).unwrap();
509                let k = n.key.clone();
510
511                if p.rgt.as_ref().map(|&id| id == n.id).unwrap_or_default() {
512                    // n is on right link of p
513                    p.rgt = n.lft;
514                } else {
515                    // n is on left link of p
516                    p.lft = n.lft;
517                }
518
519                self.update_height(&mut p);
520
521                if r_node.id == p.id {
522                    // r_node.id and p.id can overlap on small trees (2 levels, 2-3 nodes)
523                    // that leads to nasty lost update of the key, refresh below fixes that
524                    r_node = self.node(r_node.id).unwrap();
525                }
526                r_node.key = k;
527                self.save(&r_node);
528
529                self.swap_with_last(n.id);
530
531                // removing node might have caused an imbalance - balance the tree up to the root,
532                // starting from the lowest affected key (max key from left subtree in this case)
533                self.check_balance(self.root, &p.key)
534            } else {
535                // proceed with right subtree
536                let rgt = rgt_opt.unwrap();
537
538                // k - min key from right subtree
539                // n - node that holds key k, p - immediate parent of n
540                let (n, mut p) = self.min_at(rgt, r_node.id).unwrap();
541                let k = n.key.clone();
542
543                if p.lft.map(|id| id == n.id).unwrap_or_default() {
544                    // n is on left link of p
545                    p.lft = n.rgt;
546                } else {
547                    // n is on right link of p
548                    p.rgt = n.rgt;
549                }
550
551                self.update_height(&mut p);
552
553                if r_node.id == p.id {
554                    // r_node.id and p.id can overlap on small trees (2 levels, 2-3 nodes)
555                    // that leads to nasty lost update of the key, refresh below fixes that
556                    r_node = self.node(r_node.id).unwrap();
557                }
558                r_node.key = k;
559                self.save(&r_node);
560
561                self.swap_with_last(n.id);
562
563                // removing node might have caused a imbalance - balance the tree up to the root,
564                // starting from the lowest affected key (min key from right subtree in this case)
565                self.check_balance(self.root, &p.key)
566            }
567        }
568    }
569
570    // Move content of node with id = `len - 1` (parent left or right link, left, right, key, height)
571    // to node with given `id`, and remove node `len - 1` (pop the vector of nodes).
572    // This ensures that among `n` nodes in the tree, max `id` is `n-1`, so when new node is inserted,
573    // it gets an `id` as its position in the vector.
574    fn swap_with_last(&mut self, id: u64) {
575        if id == self.len() - 1 {
576            // noop: id is already last element in the vector
577            self.tree.pop();
578            return;
579        }
580
581        let key = self.node(self.len() - 1).map(|n| n.key).unwrap();
582        let (mut n, mut p) = self.lookup_at(self.root, &key).unwrap();
583
584        if n.id != p.id {
585            if p.lft.map(|id| id == n.id).unwrap_or_default() {
586                p.lft = Some(id);
587            } else {
588                p.rgt = Some(id);
589            }
590            self.save(&p);
591        }
592
593        if self.root == n.id {
594            self.root = id;
595        }
596
597        n.id = id;
598        self.save(&n);
599        self.tree.pop();
600    }
601}
602
603impl<'a, K, V> IntoIterator for &'a LegacyTreeMap<K, V>
604where
605    K: Ord + Clone + BorshSerialize + BorshDeserialize,
606    V: BorshSerialize + BorshDeserialize,
607{
608    type Item = (K, V);
609    type IntoIter = Cursor<'a, K, V>;
610
611    fn into_iter(self) -> Self::IntoIter {
612        Cursor::asc(self)
613    }
614}
615
616impl<K, V> Iterator for Cursor<'_, K, V>
617where
618    K: Ord + Clone + BorshSerialize + BorshDeserialize,
619    V: BorshSerialize + BorshDeserialize,
620{
621    type Item = (K, V);
622
623    fn next(&mut self) -> Option<Self::Item> {
624        let this_key = self.key.clone();
625
626        let next_key = self
627            .key
628            .take()
629            .and_then(|k| if self.asc { self.map.higher(&k) } else { self.map.lower(&k) })
630            .filter(|k| fits(k, &self.lo, &self.hi));
631        self.key = next_key;
632
633        this_key.and_then(|k| self.map.get(&k).map(|v| (k, v)))
634    }
635}
636
637fn fits<K: Ord>(key: &K, lo: &Bound<K>, hi: &Bound<K>) -> bool {
638    (match lo {
639        Bound::Included(ref x) => key >= x,
640        Bound::Excluded(ref x) => key > x,
641        Bound::Unbounded => true,
642    }) && (match hi {
643        Bound::Included(ref x) => key <= x,
644        Bound::Excluded(ref x) => key < x,
645        Bound::Unbounded => true,
646    })
647}
648
649pub struct Cursor<'a, K, V> {
650    asc: bool,
651    lo: Bound<K>,
652    hi: Bound<K>,
653    key: Option<K>,
654    map: &'a LegacyTreeMap<K, V>,
655}
656
657impl<'a, K, V> Cursor<'a, K, V>
658where
659    K: Ord + Clone + BorshSerialize + BorshDeserialize,
660    V: BorshSerialize + BorshDeserialize,
661{
662    fn asc(map: &'a LegacyTreeMap<K, V>) -> Self {
663        let key: Option<K> = map.min();
664        Self { asc: true, key, lo: Bound::Unbounded, hi: Bound::Unbounded, map }
665    }
666
667    fn asc_from(map: &'a LegacyTreeMap<K, V>, key: K) -> Self {
668        let key = map.higher(&key);
669        Self { asc: true, key, lo: Bound::Unbounded, hi: Bound::Unbounded, map }
670    }
671
672    fn desc(map: &'a LegacyTreeMap<K, V>) -> Self {
673        let key: Option<K> = map.max();
674        Self { asc: false, key, lo: Bound::Unbounded, hi: Bound::Unbounded, map }
675    }
676
677    fn desc_from(map: &'a LegacyTreeMap<K, V>, key: K) -> Self {
678        let key = map.lower(&key);
679        Self { asc: false, key, lo: Bound::Unbounded, hi: Bound::Unbounded, map }
680    }
681
682    fn range(map: &'a LegacyTreeMap<K, V>, lo: Bound<K>, hi: Bound<K>) -> Self {
683        let key = match &lo {
684            Bound::Included(k) if map.contains_key(k) => Some(k.clone()),
685            Bound::Included(k) | Bound::Excluded(k) => map.higher(k),
686            _ => None,
687        };
688        let key = key.filter(|k| fits(k, &lo, &hi));
689
690        Self { asc: true, key, lo, hi, map }
691    }
692}
693
694#[cfg(not(target_arch = "wasm32"))]
695#[cfg(test)]
696mod tests {
697    use super::*;
698    use crate::test_utils::{next_trie_id, test_env};
699
700    extern crate rand;
701    use self::rand::RngCore;
702    use quickcheck::QuickCheck;
703    use std::collections::BTreeMap;
704    use std::collections::HashSet;
705    use std::fmt::Formatter;
706    use std::fmt::{Debug, Result};
707
708    /// Return height of the tree - number of nodes on the longest path starting from the root node.
709    fn height<K, V>(tree: &LegacyTreeMap<K, V>) -> u64
710    where
711        K: Ord + Clone + BorshSerialize + BorshDeserialize,
712        V: BorshSerialize + BorshDeserialize,
713    {
714        tree.node(tree.root).map(|n| n.ht).unwrap_or_default()
715    }
716
717    fn random(n: u64) -> Vec<u32> {
718        let mut rng = rand::thread_rng();
719        let mut vec = Vec::with_capacity(n as usize);
720        (0..n).for_each(|_| {
721            vec.push(rng.next_u32() % 1000);
722        });
723        vec
724    }
725
726    fn log2(x: f64) -> f64 {
727        std::primitive::f64::log(x, 2.0f64)
728    }
729
730    fn max_tree_height(n: u64) -> u64 {
731        // h <= C * log2(n + D) + B
732        // where:
733        // C =~ 1.440, D =~ 1.065, B =~ 0.328
734        // (source: https://en.wikipedia.org/wiki/AVL_tree)
735        const B: f64 = -0.328;
736        const C: f64 = 1.440;
737        const D: f64 = 1.065;
738
739        let h = C * log2(n as f64 + D) + B;
740        h.ceil() as u64
741    }
742
743    impl<K> Debug for Node<K>
744    where
745        K: Ord + Clone + Debug + BorshSerialize + BorshDeserialize,
746    {
747        fn fmt(&self, f: &mut Formatter<'_>) -> Result {
748            f.debug_struct("Node")
749                .field("id", &self.id)
750                .field("key", &self.key)
751                .field("lft", &self.lft)
752                .field("rgt", &self.rgt)
753                .field("ht", &self.ht)
754                .finish()
755        }
756    }
757
758    impl<K, V> Debug for LegacyTreeMap<K, V>
759    where
760        K: Ord + Clone + Debug + BorshSerialize + BorshDeserialize,
761        V: Debug + BorshSerialize + BorshDeserialize,
762    {
763        fn fmt(&self, f: &mut Formatter<'_>) -> Result {
764            f.debug_struct("TreeMap")
765                .field("root", &self.root)
766                .field("tree", &self.tree.iter().collect::<Vec<Node<K>>>())
767                .field("val", &self.val.iter().collect::<Vec<(K, V)>>())
768                .finish()
769        }
770    }
771
772    #[test]
773    fn test_empty() {
774        let map: LegacyTreeMap<u8, u8> = LegacyTreeMap::new(b't');
775        assert_eq!(map.len(), 0);
776        assert_eq!(height(&map), 0);
777        assert_eq!(map.get(&42), None);
778        assert!(!map.contains_key(&42));
779        assert_eq!(map.min(), None);
780        assert_eq!(map.max(), None);
781        assert_eq!(map.lower(&42), None);
782        assert_eq!(map.higher(&42), None);
783    }
784
785    #[test]
786    fn test_insert_3_rotate_l_l() {
787        let mut map: LegacyTreeMap<i32, i32> = LegacyTreeMap::new(next_trie_id());
788        assert_eq!(height(&map), 0);
789
790        map.insert(&3, &3);
791        assert_eq!(height(&map), 1);
792
793        map.insert(&2, &2);
794        assert_eq!(height(&map), 2);
795
796        map.insert(&1, &1);
797        assert_eq!(height(&map), 2);
798
799        let root = map.root;
800        assert_eq!(root, 1);
801        assert_eq!(map.node(root).map(|n| n.key), Some(2));
802
803        map.clear();
804    }
805
806    #[test]
807    fn test_insert_3_rotate_r_r() {
808        let mut map: LegacyTreeMap<i32, i32> = LegacyTreeMap::new(next_trie_id());
809        assert_eq!(height(&map), 0);
810
811        map.insert(&1, &1);
812        assert_eq!(height(&map), 1);
813
814        map.insert(&2, &2);
815        assert_eq!(height(&map), 2);
816
817        map.insert(&3, &3);
818
819        let root = map.root;
820        assert_eq!(root, 1);
821        assert_eq!(map.node(root).map(|n| n.key), Some(2));
822        assert_eq!(height(&map), 2);
823
824        map.clear();
825    }
826
827    #[test]
828    fn test_insert_lookup_n_asc() {
829        let mut map: LegacyTreeMap<i32, i32> = LegacyTreeMap::new(next_trie_id());
830
831        let n: u64 = 30;
832        let cases = (0..2 * (n as i32)).collect::<Vec<i32>>();
833
834        let mut counter = 0;
835        for k in &cases {
836            if *k % 2 == 0 {
837                counter += 1;
838                map.insert(k, &counter);
839            }
840        }
841
842        counter = 0;
843        for k in &cases {
844            if *k % 2 == 0 {
845                counter += 1;
846                assert_eq!(map.get(k), Some(counter));
847            } else {
848                assert_eq!(map.get(k), None);
849            }
850        }
851
852        assert!(height(&map) <= max_tree_height(n));
853        map.clear();
854    }
855
856    #[test]
857    fn test_insert_lookup_n_desc() {
858        let mut map: LegacyTreeMap<i32, i32> = LegacyTreeMap::new(next_trie_id());
859
860        let n: u64 = 30;
861        let cases = (0..2 * (n as i32)).rev().collect::<Vec<i32>>();
862
863        let mut counter = 0;
864        for k in &cases {
865            if *k % 2 == 0 {
866                counter += 1;
867                map.insert(k, &counter);
868            }
869        }
870
871        counter = 0;
872        for k in &cases {
873            if *k % 2 == 0 {
874                counter += 1;
875                assert_eq!(map.get(k), Some(counter));
876            } else {
877                assert_eq!(map.get(k), None);
878            }
879        }
880
881        assert!(height(&map) <= max_tree_height(n));
882        map.clear();
883    }
884
885    #[test]
886    fn insert_n_random() {
887        test_env::setup_free();
888
889        for k in 1..10 {
890            // tree size is 2^k
891            let mut map: LegacyTreeMap<u32, u32> = LegacyTreeMap::new(next_trie_id());
892
893            let n = 1 << k;
894            let input: Vec<u32> = random(n);
895
896            for x in &input {
897                map.insert(x, &42);
898            }
899
900            for x in &input {
901                assert_eq!(map.get(x), Some(42));
902            }
903
904            assert!(height(&map) <= max_tree_height(n));
905            map.clear();
906        }
907    }
908
909    #[test]
910    fn test_min() {
911        let n: u64 = 30;
912        let vec = random(n);
913
914        let mut map: LegacyTreeMap<u32, u32> = LegacyTreeMap::new(b't');
915        for x in vec.iter().rev() {
916            map.insert(x, &1);
917        }
918
919        assert_eq!(map.min().unwrap(), *vec.iter().min().unwrap());
920        map.clear();
921    }
922
923    #[test]
924    fn test_max() {
925        let n: u64 = 30;
926        let vec = random(n);
927
928        let mut map: LegacyTreeMap<u32, u32> = LegacyTreeMap::new(b't');
929        for x in vec.iter().rev() {
930            map.insert(x, &1);
931        }
932
933        assert_eq!(map.max().unwrap(), *vec.iter().max().unwrap());
934        map.clear();
935    }
936
937    #[test]
938    fn test_lower() {
939        let mut map: LegacyTreeMap<u32, u32> = LegacyTreeMap::new(next_trie_id());
940        let vec: Vec<u32> = vec![10, 20, 30, 40, 50];
941
942        for x in vec.iter() {
943            map.insert(x, &1);
944        }
945
946        assert_eq!(map.lower(&5), None);
947        assert_eq!(map.lower(&10), None);
948        assert_eq!(map.lower(&11), Some(10));
949        assert_eq!(map.lower(&20), Some(10));
950        assert_eq!(map.lower(&49), Some(40));
951        assert_eq!(map.lower(&50), Some(40));
952        assert_eq!(map.lower(&51), Some(50));
953
954        map.clear();
955    }
956
957    #[test]
958    fn test_higher() {
959        let mut map: LegacyTreeMap<u32, u32> = LegacyTreeMap::new(next_trie_id());
960        let vec: Vec<u32> = vec![10, 20, 30, 40, 50];
961
962        for x in vec.iter() {
963            map.insert(x, &1);
964        }
965
966        assert_eq!(map.higher(&5), Some(10));
967        assert_eq!(map.higher(&10), Some(20));
968        assert_eq!(map.higher(&11), Some(20));
969        assert_eq!(map.higher(&20), Some(30));
970        assert_eq!(map.higher(&49), Some(50));
971        assert_eq!(map.higher(&50), None);
972        assert_eq!(map.higher(&51), None);
973
974        map.clear();
975    }
976
977    #[test]
978    fn test_floor_key() {
979        let mut map: LegacyTreeMap<u32, u32> = LegacyTreeMap::new(next_trie_id());
980        let vec: Vec<u32> = vec![10, 20, 30, 40, 50];
981
982        for x in vec.iter() {
983            map.insert(x, &1);
984        }
985
986        assert_eq!(map.floor_key(&5), None);
987        assert_eq!(map.floor_key(&10), Some(10));
988        assert_eq!(map.floor_key(&11), Some(10));
989        assert_eq!(map.floor_key(&20), Some(20));
990        assert_eq!(map.floor_key(&49), Some(40));
991        assert_eq!(map.floor_key(&50), Some(50));
992        assert_eq!(map.floor_key(&51), Some(50));
993
994        map.clear();
995    }
996
997    #[test]
998    fn test_ceil_key() {
999        let mut map: LegacyTreeMap<u32, u32> = LegacyTreeMap::new(next_trie_id());
1000        let vec: Vec<u32> = vec![10, 20, 30, 40, 50];
1001
1002        for x in vec.iter() {
1003            map.insert(x, &1);
1004        }
1005
1006        assert_eq!(map.ceil_key(&5), Some(10));
1007        assert_eq!(map.ceil_key(&10), Some(10));
1008        assert_eq!(map.ceil_key(&11), Some(20));
1009        assert_eq!(map.ceil_key(&20), Some(20));
1010        assert_eq!(map.ceil_key(&49), Some(50));
1011        assert_eq!(map.ceil_key(&50), Some(50));
1012        assert_eq!(map.ceil_key(&51), None);
1013
1014        map.clear();
1015    }
1016
1017    #[test]
1018    fn test_remove_1() {
1019        let mut map: LegacyTreeMap<u32, u32> = LegacyTreeMap::new(next_trie_id());
1020        map.insert(&1, &1);
1021        assert_eq!(map.get(&1), Some(1));
1022        map.remove(&1);
1023        assert_eq!(map.get(&1), None);
1024        assert_eq!(map.tree.len(), 0);
1025        map.clear();
1026    }
1027
1028    #[test]
1029    fn test_remove_3() {
1030        let map: LegacyTreeMap<u32, u32> = avl(&[(0, 0)], &[0, 0, 1]);
1031
1032        assert_eq!(map.iter().collect::<Vec<(u32, u32)>>(), vec![]);
1033    }
1034
1035    #[test]
1036    fn test_remove_3_desc() {
1037        let vec: Vec<u32> = vec![3, 2, 1];
1038        let mut map: LegacyTreeMap<u32, u32> = LegacyTreeMap::new(next_trie_id());
1039
1040        for x in &vec {
1041            assert_eq!(map.get(x), None);
1042            map.insert(x, &1);
1043            assert_eq!(map.get(x), Some(1));
1044        }
1045
1046        for x in &vec {
1047            assert_eq!(map.get(x), Some(1));
1048            map.remove(x);
1049            assert_eq!(map.get(x), None);
1050        }
1051        map.clear();
1052    }
1053
1054    #[test]
1055    fn test_remove_3_asc() {
1056        let vec: Vec<u32> = vec![1, 2, 3];
1057        let mut map: LegacyTreeMap<u32, u32> = LegacyTreeMap::new(next_trie_id());
1058
1059        for x in &vec {
1060            assert_eq!(map.get(x), None);
1061            map.insert(x, &1);
1062            assert_eq!(map.get(x), Some(1));
1063        }
1064
1065        for x in &vec {
1066            assert_eq!(map.get(x), Some(1));
1067            map.remove(x);
1068            assert_eq!(map.get(x), None);
1069        }
1070        map.clear();
1071    }
1072
1073    #[test]
1074    fn test_remove_7_regression_1() {
1075        let vec: Vec<u32> =
1076            vec![2104297040, 552624607, 4269683389, 3382615941, 155419892, 4102023417, 1795725075];
1077        let mut map: LegacyTreeMap<u32, u32> = LegacyTreeMap::new(next_trie_id());
1078
1079        for x in &vec {
1080            assert_eq!(map.get(x), None);
1081            map.insert(x, &1);
1082            assert_eq!(map.get(x), Some(1));
1083        }
1084
1085        for x in &vec {
1086            assert_eq!(map.get(x), Some(1));
1087            map.remove(x);
1088            assert_eq!(map.get(x), None);
1089        }
1090        map.clear();
1091    }
1092
1093    #[test]
1094    fn test_remove_7_regression_2() {
1095        let vec: Vec<u32> =
1096            vec![700623085, 87488544, 1500140781, 1111706290, 3187278102, 4042663151, 3731533080];
1097        let mut map: LegacyTreeMap<u32, u32> = LegacyTreeMap::new(next_trie_id());
1098
1099        for x in &vec {
1100            assert_eq!(map.get(x), None);
1101            map.insert(x, &1);
1102            assert_eq!(map.get(x), Some(1));
1103        }
1104
1105        for x in &vec {
1106            assert_eq!(map.get(x), Some(1));
1107            map.remove(x);
1108            assert_eq!(map.get(x), None);
1109        }
1110        map.clear();
1111    }
1112
1113    #[test]
1114    fn test_remove_9_regression() {
1115        let vec: Vec<u32> = vec![
1116            1186903464, 506371929, 1738679820, 1883936615, 1815331350, 1512669683, 3581743264,
1117            1396738166, 1902061760,
1118        ];
1119        let mut map: LegacyTreeMap<u32, u32> = LegacyTreeMap::new(next_trie_id());
1120
1121        for x in &vec {
1122            assert_eq!(map.get(x), None);
1123            map.insert(x, &1);
1124            assert_eq!(map.get(x), Some(1));
1125        }
1126
1127        for x in &vec {
1128            assert_eq!(map.get(x), Some(1));
1129            map.remove(x);
1130            assert_eq!(map.get(x), None);
1131        }
1132        map.clear();
1133    }
1134
1135    #[test]
1136    fn test_remove_20_regression_1() {
1137        let vec: Vec<u32> = vec![
1138            552517392, 3638992158, 1015727752, 2500937532, 638716734, 586360620, 2476692174,
1139            1425948996, 3608478547, 757735878, 2709959928, 2092169539, 3620770200, 783020918,
1140            1986928932, 200210441, 1972255302, 533239929, 497054557, 2137924638,
1141        ];
1142        let mut map: LegacyTreeMap<u32, u32> = LegacyTreeMap::new(next_trie_id());
1143
1144        for x in &vec {
1145            assert_eq!(map.get(x), None);
1146            map.insert(x, &1);
1147            assert_eq!(map.get(x), Some(1));
1148        }
1149
1150        for x in &vec {
1151            assert_eq!(map.get(x), Some(1));
1152            map.remove(x);
1153            assert_eq!(map.get(x), None);
1154        }
1155        map.clear();
1156    }
1157
1158    #[test]
1159    fn test_remove_7_regression() {
1160        let vec: Vec<u32> = vec![280, 606, 163, 857, 436, 508, 44, 801];
1161
1162        let mut map: LegacyTreeMap<u32, u32> = LegacyTreeMap::new(next_trie_id());
1163
1164        for x in &vec {
1165            assert_eq!(map.get(x), None);
1166            map.insert(x, &1);
1167            assert_eq!(map.get(x), Some(1));
1168        }
1169
1170        for x in &vec {
1171            assert_eq!(map.get(x), Some(1));
1172            map.remove(x);
1173            assert_eq!(map.get(x), None);
1174        }
1175
1176        assert_eq!(map.len(), 0, "map.len() > 0");
1177        assert_eq!(map.val.len(), 0, "map.val is not empty");
1178        assert_eq!(map.tree.len(), 0, "map.tree is not empty");
1179        map.clear();
1180    }
1181
1182    #[test]
1183    fn test_insert_8_remove_4_regression() {
1184        let insert = vec![882, 398, 161, 76];
1185        let remove = vec![242, 687, 860, 811];
1186
1187        let mut map: LegacyTreeMap<u32, u32> = LegacyTreeMap::new(next_trie_id());
1188
1189        for (i, (k1, k2)) in insert.iter().zip(remove.iter()).enumerate() {
1190            let v = i as u32;
1191            map.insert(k1, &v);
1192            map.insert(k2, &v);
1193        }
1194
1195        for k in remove.iter() {
1196            map.remove(k);
1197        }
1198
1199        assert_eq!(map.len(), insert.len() as u64);
1200
1201        for (i, k) in insert.iter().enumerate() {
1202            assert_eq!(map.get(k), Some(i as u32));
1203        }
1204    }
1205
1206    #[test]
1207    fn test_remove_n() {
1208        let n: u64 = 20;
1209        let vec = random(n);
1210
1211        let mut set: HashSet<u32> = HashSet::new();
1212        let mut map: LegacyTreeMap<u32, u32> = LegacyTreeMap::new(next_trie_id());
1213        for x in &vec {
1214            map.insert(x, &1);
1215            set.insert(*x);
1216        }
1217
1218        assert_eq!(map.len(), set.len() as u64);
1219
1220        for x in &set {
1221            assert_eq!(map.get(x), Some(1));
1222            map.remove(x);
1223            assert_eq!(map.get(x), None);
1224        }
1225
1226        assert_eq!(map.len(), 0, "map.len() > 0");
1227        assert_eq!(map.tree.len(), 0, "map.tree is not empty");
1228        assert_eq!(map.val.len(), 0, "map.val is not empty");
1229        map.clear();
1230    }
1231
1232    #[test]
1233    fn test_remove_root_3() {
1234        let mut map: LegacyTreeMap<u32, u32> = LegacyTreeMap::new(next_trie_id());
1235        map.insert(&2, &1);
1236        map.insert(&3, &1);
1237        map.insert(&1, &1);
1238        map.insert(&4, &1);
1239
1240        map.remove(&2);
1241
1242        assert_eq!(map.get(&1), Some(1));
1243        assert_eq!(map.get(&2), None);
1244        assert_eq!(map.get(&3), Some(1));
1245        assert_eq!(map.get(&4), Some(1));
1246        map.clear();
1247    }
1248
1249    #[test]
1250    fn test_insert_2_remove_2_regression() {
1251        let ins: Vec<u32> = vec![11760225, 611327897];
1252        let rem: Vec<u32> = vec![2982517385, 1833990072];
1253
1254        let mut map: LegacyTreeMap<u32, u32> = LegacyTreeMap::new(next_trie_id());
1255        map.insert(&ins[0], &1);
1256        map.insert(&ins[1], &1);
1257
1258        map.remove(&rem[0]);
1259        map.remove(&rem[1]);
1260
1261        let h = height(&map);
1262        let h_max = max_tree_height(map.len());
1263        assert!(h <= h_max, "h={} h_max={}", h, h_max);
1264        map.clear();
1265    }
1266
1267    #[test]
1268    fn test_insert_n_duplicates() {
1269        let mut map: LegacyTreeMap<u32, u32> = LegacyTreeMap::new(next_trie_id());
1270
1271        for x in 0..30 {
1272            map.insert(&x, &x);
1273            map.insert(&42, &x);
1274        }
1275
1276        assert_eq!(map.get(&42), Some(29));
1277        assert_eq!(map.len(), 31);
1278        assert_eq!(map.val.len(), 31);
1279        assert_eq!(map.tree.len(), 31);
1280
1281        map.clear();
1282    }
1283
1284    #[test]
1285    fn test_insert_2n_remove_n_random() {
1286        for k in 1..4 {
1287            let mut map: LegacyTreeMap<u32, u32> = LegacyTreeMap::new(next_trie_id());
1288            let mut set: HashSet<u32> = HashSet::new();
1289
1290            let n = 1 << k;
1291            let ins: Vec<u32> = random(n);
1292            let rem: Vec<u32> = random(n);
1293
1294            for x in &ins {
1295                set.insert(*x);
1296                map.insert(x, &42);
1297            }
1298
1299            for x in &rem {
1300                set.insert(*x);
1301                map.insert(x, &42);
1302            }
1303
1304            for x in &rem {
1305                set.remove(x);
1306                map.remove(x);
1307            }
1308
1309            assert_eq!(map.len(), set.len() as u64);
1310
1311            let h = height(&map);
1312            let h_max = max_tree_height(n);
1313            assert!(h <= h_max, "[n={}] tree is too high: {} (max is {}).", n, h, h_max);
1314
1315            map.clear();
1316        }
1317    }
1318
1319    #[test]
1320    fn test_remove_empty() {
1321        let mut map: LegacyTreeMap<u32, u32> = LegacyTreeMap::new(next_trie_id());
1322        assert_eq!(map.remove(&1), None);
1323    }
1324
1325    #[test]
1326    fn test_to_vec() {
1327        let mut map: LegacyTreeMap<u32, u32> = LegacyTreeMap::new(next_trie_id());
1328        map.insert(&1, &41);
1329        map.insert(&2, &42);
1330        map.insert(&3, &43);
1331
1332        assert_eq!(map.to_vec(), vec![(1, 41), (2, 42), (3, 43)]);
1333        map.clear();
1334    }
1335
1336    #[test]
1337    fn test_to_vec_empty() {
1338        let map: LegacyTreeMap<u32, u32> = LegacyTreeMap::new(next_trie_id());
1339        assert!(map.to_vec().is_empty());
1340    }
1341
1342    #[test]
1343    fn test_iter() {
1344        let mut map: LegacyTreeMap<u32, u32> = LegacyTreeMap::new(next_trie_id());
1345        map.insert(&1, &41);
1346        map.insert(&2, &42);
1347        map.insert(&3, &43);
1348
1349        assert_eq!(map.iter().collect::<Vec<(u32, u32)>>(), vec![(1, 41), (2, 42), (3, 43)]);
1350        map.clear();
1351    }
1352
1353    #[test]
1354    fn test_iter_empty() {
1355        let map: LegacyTreeMap<u32, u32> = LegacyTreeMap::new(next_trie_id());
1356        assert_eq!(map.iter().count(), 0);
1357    }
1358
1359    #[test]
1360    fn test_iter_rev() {
1361        let mut map: LegacyTreeMap<u32, u32> = LegacyTreeMap::new(next_trie_id());
1362        map.insert(&1, &41);
1363        map.insert(&2, &42);
1364        map.insert(&3, &43);
1365
1366        assert_eq!(map.iter_rev().collect::<Vec<(u32, u32)>>(), vec![(3, 43), (2, 42), (1, 41)]);
1367        map.clear();
1368    }
1369
1370    #[test]
1371    fn test_iter_rev_empty() {
1372        let map: LegacyTreeMap<u32, u32> = LegacyTreeMap::new(next_trie_id());
1373        assert_eq!(map.iter_rev().count(), 0);
1374    }
1375
1376    #[test]
1377    fn test_iter_from() {
1378        let mut map: LegacyTreeMap<u32, u32> = LegacyTreeMap::new(next_trie_id());
1379
1380        let one: Vec<u32> = vec![10, 20, 30, 40, 50];
1381        let two: Vec<u32> = vec![45, 35, 25, 15, 5];
1382
1383        for x in &one {
1384            map.insert(x, &42);
1385        }
1386
1387        for x in &two {
1388            map.insert(x, &42);
1389        }
1390
1391        assert_eq!(
1392            map.iter_from(29).collect::<Vec<(u32, u32)>>(),
1393            vec![(30, 42), (35, 42), (40, 42), (45, 42), (50, 42)]
1394        );
1395
1396        assert_eq!(
1397            map.iter_from(30).collect::<Vec<(u32, u32)>>(),
1398            vec![(35, 42), (40, 42), (45, 42), (50, 42)]
1399        );
1400
1401        assert_eq!(
1402            map.iter_from(31).collect::<Vec<(u32, u32)>>(),
1403            vec![(35, 42), (40, 42), (45, 42), (50, 42)]
1404        );
1405        map.clear();
1406    }
1407
1408    #[test]
1409    fn test_iter_from_empty() {
1410        let map: LegacyTreeMap<u32, u32> = LegacyTreeMap::new(next_trie_id());
1411        assert_eq!(map.iter_from(42).count(), 0);
1412    }
1413
1414    #[test]
1415    fn test_iter_rev_from() {
1416        let mut map: LegacyTreeMap<u32, u32> = LegacyTreeMap::new(next_trie_id());
1417
1418        let one: Vec<u32> = vec![10, 20, 30, 40, 50];
1419        let two: Vec<u32> = vec![45, 35, 25, 15, 5];
1420
1421        for x in &one {
1422            map.insert(x, &42);
1423        }
1424
1425        for x in &two {
1426            map.insert(x, &42);
1427        }
1428
1429        assert_eq!(
1430            map.iter_rev_from(29).collect::<Vec<(u32, u32)>>(),
1431            vec![(25, 42), (20, 42), (15, 42), (10, 42), (5, 42)]
1432        );
1433
1434        assert_eq!(
1435            map.iter_rev_from(30).collect::<Vec<(u32, u32)>>(),
1436            vec![(25, 42), (20, 42), (15, 42), (10, 42), (5, 42)]
1437        );
1438
1439        assert_eq!(
1440            map.iter_rev_from(31).collect::<Vec<(u32, u32)>>(),
1441            vec![(30, 42), (25, 42), (20, 42), (15, 42), (10, 42), (5, 42)]
1442        );
1443        map.clear();
1444    }
1445
1446    #[test]
1447    fn test_range() {
1448        let mut map: LegacyTreeMap<u32, u32> = LegacyTreeMap::new(next_trie_id());
1449
1450        let one: Vec<u32> = vec![10, 20, 30, 40, 50];
1451        let two: Vec<u32> = vec![45, 35, 25, 15, 5];
1452
1453        for x in &one {
1454            map.insert(x, &42);
1455        }
1456
1457        for x in &two {
1458            map.insert(x, &42);
1459        }
1460
1461        assert_eq!(
1462            map.range((Bound::Included(20), Bound::Excluded(30))).collect::<Vec<(u32, u32)>>(),
1463            vec![(20, 42), (25, 42)]
1464        );
1465
1466        assert_eq!(
1467            map.range((Bound::Excluded(10), Bound::Included(40))).collect::<Vec<(u32, u32)>>(),
1468            vec![(15, 42), (20, 42), (25, 42), (30, 42), (35, 42), (40, 42)]
1469        );
1470
1471        assert_eq!(
1472            map.range((Bound::Included(20), Bound::Included(40))).collect::<Vec<(u32, u32)>>(),
1473            vec![(20, 42), (25, 42), (30, 42), (35, 42), (40, 42)]
1474        );
1475
1476        assert_eq!(
1477            map.range((Bound::Excluded(20), Bound::Excluded(45))).collect::<Vec<(u32, u32)>>(),
1478            vec![(25, 42), (30, 42), (35, 42), (40, 42)]
1479        );
1480
1481        assert_eq!(
1482            map.range((Bound::Excluded(20), Bound::Excluded(45))).collect::<Vec<(u32, u32)>>(),
1483            vec![(25, 42), (30, 42), (35, 42), (40, 42)]
1484        );
1485
1486        assert_eq!(
1487            map.range((Bound::Excluded(25), Bound::Excluded(30))).collect::<Vec<(u32, u32)>>(),
1488            vec![]
1489        );
1490
1491        assert_eq!(
1492            map.range((Bound::Included(25), Bound::Included(25))).collect::<Vec<(u32, u32)>>(),
1493            vec![(25, 42)]
1494        );
1495
1496        assert_eq!(
1497            map.range((Bound::Excluded(25), Bound::Included(25))).collect::<Vec<(u32, u32)>>(),
1498            vec![]
1499        ); // the range makes no sense, but `BTreeMap` does not panic in this case
1500
1501        map.clear();
1502    }
1503
1504    #[test]
1505    #[should_panic(expected = "Invalid range.")]
1506    fn test_range_panics_same_excluded() {
1507        let map: LegacyTreeMap<u32, u32> = LegacyTreeMap::new(next_trie_id());
1508        let _ = map.range((Bound::Excluded(1), Bound::Excluded(1)));
1509    }
1510
1511    #[test]
1512    #[should_panic(expected = "Invalid range.")]
1513    fn test_range_panics_non_overlap_incl_exlc() {
1514        let map: LegacyTreeMap<u32, u32> = LegacyTreeMap::new(next_trie_id());
1515        let _ = map.range((Bound::Included(2), Bound::Excluded(1)));
1516    }
1517
1518    #[test]
1519    #[should_panic(expected = "Invalid range.")]
1520    fn test_range_panics_non_overlap_excl_incl() {
1521        let map: LegacyTreeMap<u32, u32> = LegacyTreeMap::new(next_trie_id());
1522        let _ = map.range((Bound::Excluded(2), Bound::Included(1)));
1523    }
1524
1525    #[test]
1526    #[should_panic(expected = "Invalid range.")]
1527    fn test_range_panics_non_overlap_excl_excl() {
1528        let map: LegacyTreeMap<u32, u32> = LegacyTreeMap::new(next_trie_id());
1529        let _ = map.range((Bound::Excluded(2), Bound::Excluded(1)));
1530    }
1531
1532    #[test]
1533    #[should_panic(expected = "Invalid range.")]
1534    fn test_range_panics_non_overlap_incl_incl() {
1535        let map: LegacyTreeMap<u32, u32> = LegacyTreeMap::new(next_trie_id());
1536        let _ = map.range((Bound::Included(2), Bound::Included(1)));
1537    }
1538
1539    #[test]
1540    fn test_iter_rev_from_empty() {
1541        let map: LegacyTreeMap<u32, u32> = LegacyTreeMap::new(next_trie_id());
1542        assert_eq!(map.iter_rev_from(42).count(), 0);
1543    }
1544
1545    #[test]
1546    fn test_balance_regression_1() {
1547        let insert = vec![(2, 0), (3, 0), (4, 0)];
1548        let remove = vec![0, 0, 0, 1];
1549
1550        let map = avl(&insert, &remove);
1551        assert!(is_balanced(&map, map.root));
1552    }
1553
1554    #[test]
1555    fn test_balance_regression_2() {
1556        let insert = vec![(1, 0), (2, 0), (0, 0), (3, 0), (5, 0), (6, 0)];
1557        let remove = vec![0, 0, 0, 3, 5, 6, 7, 4];
1558
1559        let map = avl(&insert, &remove);
1560        assert!(is_balanced(&map, map.root));
1561    }
1562
1563    //
1564    // Property-based tests of AVL-based TreeMap against std::collections::BTreeMap
1565    //
1566
1567    fn avl<K, V>(insert: &[(K, V)], remove: &[K]) -> LegacyTreeMap<K, V>
1568    where
1569        K: Ord + Clone + BorshSerialize + BorshDeserialize,
1570        V: Default + BorshSerialize + BorshDeserialize,
1571    {
1572        test_env::setup_free();
1573        let mut map: LegacyTreeMap<K, V> = LegacyTreeMap::new(next_trie_id());
1574        for k in remove {
1575            map.insert(k, &Default::default());
1576        }
1577        let n = insert.len().max(remove.len());
1578        for i in 0..n {
1579            if i < remove.len() {
1580                map.remove(&remove[i]);
1581            }
1582            if i < insert.len() {
1583                let (k, v) = &insert[i];
1584                map.insert(k, v);
1585            }
1586        }
1587        map
1588    }
1589
1590    fn rb<K, V>(insert: &[(K, V)], remove: &[K]) -> BTreeMap<K, V>
1591    where
1592        K: Ord + Clone + BorshSerialize + BorshDeserialize,
1593        V: Clone + Default + BorshSerialize + BorshDeserialize,
1594    {
1595        let mut map: BTreeMap<K, V> = BTreeMap::default();
1596        for k in remove {
1597            map.insert(k.clone(), Default::default());
1598        }
1599        let n = insert.len().max(remove.len());
1600        for i in 0..n {
1601            if i < remove.len() {
1602                map.remove(&remove[i]);
1603            }
1604            if i < insert.len() {
1605                let (k, v) = &insert[i];
1606                map.insert(k.clone(), v.clone());
1607            }
1608        }
1609        map
1610    }
1611
1612    #[test]
1613    fn prop_avl_vs_rb() {
1614        fn prop(insert: Vec<(u32, u32)>, remove: Vec<u32>) -> bool {
1615            let a = avl(&insert, &remove);
1616            let b = rb(&insert, &remove);
1617            let v1: Vec<(u32, u32)> = a.iter().collect();
1618            let v2: Vec<(u32, u32)> = b.into_iter().collect();
1619            v1 == v2
1620        }
1621
1622        QuickCheck::new()
1623            .tests(300)
1624            .quickcheck(prop as fn(std::vec::Vec<(u32, u32)>, std::vec::Vec<u32>) -> bool);
1625    }
1626
1627    fn is_balanced<K, V>(map: &LegacyTreeMap<K, V>, root: u64) -> bool
1628    where
1629        K: Debug + Ord + Clone + BorshSerialize + BorshDeserialize,
1630        V: Debug + BorshSerialize + BorshDeserialize,
1631    {
1632        let node = map.node(root).unwrap();
1633        let balance = map.get_balance(&node);
1634
1635        (balance >= -1 && balance <= 1)
1636            && node.lft.map(|id| is_balanced(map, id)).unwrap_or(true)
1637            && node.rgt.map(|id| is_balanced(map, id)).unwrap_or(true)
1638    }
1639
1640    #[test]
1641    fn prop_avl_balance() {
1642        test_env::setup_free();
1643
1644        fn prop(insert: Vec<(u32, u32)>, remove: Vec<u32>) -> bool {
1645            let map = avl(&insert, &remove);
1646            map.len() == 0 || is_balanced(&map, map.root)
1647        }
1648
1649        QuickCheck::new()
1650            .tests(300)
1651            .quickcheck(prop as fn(std::vec::Vec<(u32, u32)>, std::vec::Vec<u32>) -> bool);
1652    }
1653
1654    #[test]
1655    fn prop_avl_height() {
1656        test_env::setup_free();
1657
1658        fn prop(insert: Vec<(u32, u32)>, remove: Vec<u32>) -> bool {
1659            let map = avl(&insert, &remove);
1660            height(&map) <= max_tree_height(map.len())
1661        }
1662
1663        QuickCheck::new()
1664            .tests(300)
1665            .quickcheck(prop as fn(std::vec::Vec<(u32, u32)>, std::vec::Vec<u32>) -> bool);
1666    }
1667
1668    fn range_prop(
1669        insert: Vec<(u32, u32)>,
1670        remove: Vec<u32>,
1671        range: (Bound<u32>, Bound<u32>),
1672    ) -> bool {
1673        let a = avl(&insert, &remove);
1674        let b = rb(&insert, &remove);
1675        let v1: Vec<(u32, u32)> = a.range(range).collect();
1676        let v2: Vec<(u32, u32)> = b.range(range).map(|(k, v)| (*k, *v)).collect();
1677        v1 == v2
1678    }
1679
1680    type Prop = fn(std::vec::Vec<(u32, u32)>, std::vec::Vec<u32>, u32, u32) -> bool;
1681
1682    #[test]
1683    fn prop_avl_vs_rb_range_incl_incl() {
1684        fn prop(insert: Vec<(u32, u32)>, remove: Vec<u32>, r1: u32, r2: u32) -> bool {
1685            let range = (Bound::Included(r1.min(r2)), Bound::Included(r1.max(r2)));
1686            range_prop(insert, remove, range)
1687        }
1688
1689        QuickCheck::new().tests(300).quickcheck(prop as Prop);
1690    }
1691
1692    #[test]
1693    fn prop_avl_vs_rb_range_incl_excl() {
1694        fn prop(insert: Vec<(u32, u32)>, remove: Vec<u32>, r1: u32, r2: u32) -> bool {
1695            let range = (Bound::Included(r1.min(r2)), Bound::Excluded(r1.max(r2)));
1696            range_prop(insert, remove, range)
1697        }
1698
1699        QuickCheck::new().tests(300).quickcheck(prop as Prop);
1700    }
1701
1702    #[test]
1703    fn prop_avl_vs_rb_range_excl_incl() {
1704        fn prop(insert: Vec<(u32, u32)>, remove: Vec<u32>, r1: u32, r2: u32) -> bool {
1705            let range = (Bound::Excluded(r1.min(r2)), Bound::Included(r1.max(r2)));
1706            range_prop(insert, remove, range)
1707        }
1708
1709        QuickCheck::new().tests(300).quickcheck(prop as Prop);
1710    }
1711
1712    #[test]
1713    fn prop_avl_vs_rb_range_excl_excl() {
1714        fn prop(insert: Vec<(u32, u32)>, remove: Vec<u32>, r1: u32, r2: u32) -> bool {
1715            // (Excluded(x), Excluded(x)) is invalid range, checking against it makes no sense
1716            r1 == r2 || {
1717                let range = (Bound::Excluded(r1.min(r2)), Bound::Excluded(r1.max(r2)));
1718                range_prop(insert, remove, range)
1719            }
1720        }
1721
1722        QuickCheck::new().tests(300).quickcheck(prop as Prop);
1723    }
1724}