Skip to main content

platform_trees/trees/
iterative_size_balanced_tree.rs

1use crate::RecursiveSizeBalancedTree;
2use platform_num::LinkReference;
3
4/// Extends [`RecursiveSizeBalancedTree`] with iterative `attach` and
5/// `detach` operations.
6///
7/// The iterative approach avoids stack overflow on deep trees by using
8/// a loop with pointer-chasing instead of recursion.
9///
10/// # Usage
11///
12/// Implement [`RecursiveSizeBalancedTree`] for your storage type, then
13/// add an empty `impl IterativeSizeBalancedTree<T> for MyStorage {}`.
14///
15/// # Safety
16///
17/// All methods are `unsafe` — the caller must ensure that `root` points
18/// to a valid, writable location and that `node` indices are valid.
19pub trait IterativeSizeBalancedTree<T: LinkReference>: RecursiveSizeBalancedTree<T> {
20    /// Inserts `node` into the tree rooted at `*root`.
21    ///
22    /// If `*root` is zero (empty tree), `node` becomes the root with
23    /// size 1. Otherwise the tree is rebalanced after insertion.
24    unsafe fn attach(&mut self, root: *mut T, node: T) {
25        unsafe {
26            if *root == T::from_byte(0) {
27                self.set_size(node, T::from_byte(1));
28                *root = node;
29                return;
30            }
31            self.attach_core(root, node);
32        }
33    }
34    /// Removes `node` from the tree rooted at `*root`.
35    ///
36    /// The tree is rebalanced after removal. After detach, `node`'s
37    /// left, right, and size fields are cleared to zero.
38    unsafe fn detach(&mut self, root: *mut T, node: T) {
39        unsafe {
40            self.detach_core(root, node);
41        }
42    }
43
44    unsafe fn attach_core(&mut self, mut root: *mut T, node: T) {
45        unsafe {
46            loop {
47                let left = self.get_mut_left_reference(*root);
48                let left_size = self.get_size_or_zero(*left);
49                let right = self.get_mut_right_reference(*root);
50                let right_size = self.get_size_or_zero(*right);
51                if self.first_is_to_the_left_of_second(node, *root) {
52                    if *left == T::from_byte(0) {
53                        self.inc_size(*root);
54                        self.set_size(node, T::from_byte(1));
55                        *left = node;
56                        return;
57                    }
58                    if self.first_is_to_the_left_of_second(node, *left) {
59                        if (left_size + T::from_byte(1)) > right_size {
60                            self.right_rotate(root);
61                        } else {
62                            self.inc_size(*root);
63                            root = left;
64                        }
65                    } else {
66                        let left_right_size = self.get_size_or_zero(self.get_right(*left));
67                        if (left_right_size + T::from_byte(1)) > right_size {
68                            if left_right_size == T::from_byte(0) && right_size == T::from_byte(0) {
69                                self.set_left(node, *left);
70                                self.set_right(node, *root);
71                                self.set_size(node, left_size + T::from_byte(1) + T::from_byte(1));
72                                self.set_left(*root, T::from_byte(0));
73                                self.set_size(*root, T::from_byte(1));
74                                *root = node;
75                                return;
76                            }
77                            self.left_rotate(left);
78                            self.right_rotate(root);
79                        } else {
80                            self.inc_size(*root);
81                            root = left;
82                        }
83                    }
84                } else {
85                    if *right == T::from_byte(0) {
86                        self.inc_size(*root);
87                        self.set_size(node, T::from_byte(1));
88                        *right = node;
89                        return;
90                    }
91                    if self.first_is_to_the_right_of_second(node, *right) {
92                        if (right_size + T::from_byte(1)) > left_size {
93                            self.left_rotate(root);
94                        } else {
95                            self.inc_size(*root);
96                            root = right;
97                        }
98                    } else {
99                        let right_left_size = self.get_size_or_zero(self.get_left(*right));
100                        if (right_left_size + T::from_byte(1)) > left_size {
101                            if right_left_size == T::from_byte(0) && left_size == T::from_byte(0) {
102                                self.set_left(node, *root);
103                                self.set_right(node, *right);
104                                self.set_size(node, right_size + T::from_byte(1) + T::from_byte(1));
105                                self.set_right(*root, T::from_byte(0));
106                                self.set_size(*root, T::from_byte(1));
107                                *root = node;
108                                return;
109                            }
110                            self.right_rotate(right);
111                            self.left_rotate(root);
112                        } else {
113                            self.inc_size(*root);
114                            root = right;
115                        }
116                    }
117                }
118            }
119        }
120    }
121
122    unsafe fn detach_core(&mut self, mut root: *mut T, node: T) {
123        unsafe {
124            loop {
125                let left = self.get_mut_left_reference(*root);
126                let left_size = self.get_size_or_zero(*left);
127                let right = self.get_mut_right_reference(*root);
128                let right_size = self.get_size_or_zero(*right);
129                if self.first_is_to_the_left_of_second(node, *root) {
130                    let decremented_left_size = left_size - T::from_byte(1);
131                    if self.get_size_or_zero(self.get_right_or_default(*right))
132                        > decremented_left_size
133                    {
134                        self.left_rotate(root);
135                    } else if self.get_size_or_zero(self.get_left_or_default(*right))
136                        > decremented_left_size
137                    {
138                        self.right_rotate(right);
139                        self.left_rotate(root);
140                    } else {
141                        self.dec_size(*root);
142                        root = left;
143                    }
144                } else if self.first_is_to_the_right_of_second(node, *root) {
145                    let decremented_right_size = right_size - T::from_byte(1);
146                    if self.get_size_or_zero(self.get_left_or_default(*left))
147                        > decremented_right_size
148                    {
149                        self.right_rotate(root);
150                    } else if self.get_size_or_zero(self.get_right_or_default(*left))
151                        > decremented_right_size
152                    {
153                        self.left_rotate(left);
154                        self.right_rotate(root);
155                    } else {
156                        self.dec_size(*root);
157                        root = right;
158                    }
159                } else {
160                    if left_size > T::from_byte(0) && right_size > T::from_byte(0) {
161                        let replacement;
162                        if left_size > right_size {
163                            replacement = self.get_rightest(*left);
164                            self.detach_core(left, replacement);
165                        } else {
166                            replacement = self.get_leftest(*right);
167                            self.detach_core(right, replacement);
168                        }
169                        self.set_left(replacement, *left);
170                        self.set_right(replacement, *right);
171                        self.set_size(replacement, left_size + right_size);
172                        *root = replacement;
173                    } else if left_size > T::from_byte(0) {
174                        *root = *left;
175                    } else if right_size > T::from_byte(0) {
176                        *root = *right;
177                    } else {
178                        *root = T::from_byte(0);
179                    }
180                    self.clear_node(node);
181                    return;
182                }
183            }
184        }
185    }
186}