1 // Copyright 2024 The Pigweed Authors
2 //
3 // Licensed under the Apache License, Version 2.0 (the "License"); you may not
4 // use this file except in compliance with the License. You may obtain a copy of
5 // the License at
6 //
7 // https://www.apache.org/licenses/LICENSE-2.0
8 //
9 // Unless required by applicable law or agreed to in writing, software
10 // distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
11 // WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
12 // License for the specific language governing permissions and limitations under
13 // the License.
14
15 #include "pw_containers/internal/aa_tree_item.h"
16
17 namespace pw::containers::internal {
18
GetLevel() const19 uint8_t AATreeItem::GetLevel() const {
20 return parent_.packed_value() | (left_.packed_value() << 2) |
21 (right_.packed_value() << 4);
22 }
23
SetLevel(uint8_t level)24 void AATreeItem::SetLevel(uint8_t level) {
25 parent_.set_packed_value(level & 3);
26 left_.set_packed_value((level >> 2) & 3);
27 right_.set_packed_value((level >> 4) & 3);
28 }
29
IsMapped() const30 bool AATreeItem::IsMapped() const {
31 return GetLevel() != 0 || parent_.get() != nullptr ||
32 left_.get() != nullptr || right_.get() != nullptr;
33 }
34
GetTreeSize()35 size_t AATreeItem::GetTreeSize() {
36 return 1 + (left_.get() == nullptr ? 0 : left_->GetTreeSize()) +
37 (right_.get() == nullptr ? 0 : right_->GetTreeSize());
38 }
39
GetRoot()40 AATreeItem* AATreeItem::GetRoot() {
41 AATreeItem* node = this;
42 while (node->parent_.get() != nullptr) {
43 node = node->parent_.get();
44 }
45 return node;
46 }
47
GetLeftmost()48 AATreeItem* AATreeItem::GetLeftmost() {
49 return left_.get() == nullptr ? this : left_->GetLeftmost();
50 }
51
GetRightmost()52 AATreeItem* AATreeItem::GetRightmost() {
53 return right_.get() == nullptr ? this : right_->GetRightmost();
54 }
55
GetPredecessor()56 AATreeItem* AATreeItem::GetPredecessor() {
57 if (left_.get() != nullptr) {
58 return left_->GetRightmost();
59 }
60 const AATreeItem* current = this;
61 AATreeItem* ancestor = parent_.get();
62 while (ancestor != nullptr && ancestor->left_.get() == current) {
63 current = ancestor;
64 ancestor = ancestor->parent_.get();
65 }
66 return ancestor;
67 }
68
GetSuccessor()69 AATreeItem* AATreeItem::GetSuccessor() {
70 if (right_.get() != nullptr) {
71 return right_->GetLeftmost();
72 }
73 const AATreeItem* current = this;
74 AATreeItem* ancestor = parent_.get();
75 while (ancestor != nullptr && ancestor->right_.get() == current) {
76 current = ancestor;
77 ancestor = ancestor->parent_.get();
78 }
79 return ancestor;
80 }
81
Unmap()82 AATreeItem* AATreeItem::Unmap() {
83 AATreeItem* node;
84 if (left_.get() == nullptr && right_.get() == nullptr) {
85 // Leaf node.
86 node = nullptr;
87
88 } else if (left_.get() == nullptr) {
89 // Replace the node with the next one in value.
90 node = GetSuccessor();
91 right_->parent_.set(nullptr);
92 node->SetRight(node->Unmap());
93
94 } else {
95 // Replace the node with the previous one in value.
96 node = GetPredecessor();
97 left_->parent_.set(nullptr);
98 node->SetLeft(node->Unmap());
99 node->SetRight(right_.get());
100 }
101 if (parent_.get() != nullptr) {
102 parent_->Replace(this, node);
103 } else if (node == nullptr) {
104 // Removing the only node from the tree.
105 Reset();
106 return nullptr;
107 }
108 node = node == nullptr ? GetRoot() : node->Rebalance();
109 Reset();
110 return node;
111 }
112
SetLeft(AATreeItem * left)113 void AATreeItem::SetLeft(AATreeItem* left) {
114 if (left != nullptr) {
115 left->parent_.set(this);
116 }
117 left_.set(left);
118 }
119
SetRight(AATreeItem * right)120 void AATreeItem::SetRight(AATreeItem* right) {
121 if (right != nullptr) {
122 right->parent_.set(this);
123 }
124 right_.set(right);
125 }
126
Replace(AATreeItem * old_child,AATreeItem * new_child)127 void AATreeItem::Replace(AATreeItem* old_child, AATreeItem* new_child) {
128 if (left_.get() == old_child) {
129 SetLeft(new_child);
130 } else if (right_.get() == old_child) {
131 SetRight(new_child);
132 }
133 }
134
Skew()135 AATreeItem* AATreeItem::Skew() {
136 if (left_.get() == nullptr || GetLevel() != left_->GetLevel()) {
137 return this;
138 }
139 AATreeItem* skewed = left_.get();
140 SetLeft(skewed->right_.get());
141 skewed->parent_.set(parent_.get());
142 skewed->SetRight(this);
143 return skewed;
144 }
145
Split()146 AATreeItem* AATreeItem::Split() {
147 if (right_.get() == nullptr || right_->right_.get() == nullptr ||
148 GetLevel() != right_->right_->GetLevel()) {
149 return this;
150 }
151 AATreeItem* split = right_.get();
152 SetRight(split->left_.get());
153 split->parent_.set(parent_.get());
154 split->SetLeft(this);
155 split->SetLevel(split->GetLevel() + 1);
156 return split;
157 }
158
Rebalance()159 AATreeItem* AATreeItem::Rebalance() {
160 AATreeItem* node = this;
161 while (true) {
162 uint8_t left_level =
163 node->left_.get() == nullptr ? 0 : node->left_->GetLevel();
164 uint8_t right_level =
165 node->right_.get() == nullptr ? 0 : node->right_->GetLevel();
166 uint8_t new_level = std::min(left_level, right_level) + 1;
167 if (new_level < node->GetLevel()) {
168 node->SetLevel(new_level);
169 if (new_level < right_level) {
170 node->right_->SetLevel(new_level);
171 }
172 }
173 AATreeItem* parent = node->parent_.get();
174 AATreeItem* orig = node;
175 node = node->Skew();
176 if (node->right_.get() != nullptr) {
177 node->SetRight(node->right_->Skew());
178 if (node->right_->right_.get() != nullptr) {
179 node->right_->SetRight(node->right_->right_->Skew());
180 }
181 }
182 node = node->Split();
183 if (node->right_.get() != nullptr) {
184 node->SetRight(node->right_->Split());
185 }
186 if (parent == nullptr) {
187 return node;
188 }
189 parent->Replace(orig, node);
190 node = parent;
191 }
192 }
193
Clear()194 void AATreeItem::Clear() {
195 if (parent_.get() != nullptr) {
196 parent_->Replace(this, nullptr);
197 }
198 if (left_.get() != nullptr) {
199 left_->Clear();
200 }
201 if (right_.get() != nullptr) {
202 right_->Clear();
203 }
204 Reset();
205 }
206
Reset()207 void AATreeItem::Reset() {
208 parent_ = PackedPtr<AATreeItem>();
209 left_ = PackedPtr<AATreeItem>();
210 right_ = PackedPtr<AATreeItem>();
211 }
212
213 } // namespace pw::containers::internal
214