From c62d52b5c0b5df9e3bb5d2b964b4c33e2ed2a997 Mon Sep 17 00:00:00 2001 From: Justin King Date: Wed, 23 Sep 2026 10:17:30 -0700 Subject: [PATCH] Introduce RBTree implementation tailored for `proto2::Arena` PiperOrigin-RevId: 986842403 --- internal/BUILD | 20 ++ internal/arena_tree.cc | 352 +++++++++++++++++++++++++++++++ internal/arena_tree.h | 398 ++++++++++++++++++++++++++++++++++++ internal/arena_tree_test.cc | 223 ++++++++++++++++++++ 4 files changed, 993 insertions(+) create mode 100644 internal/arena_tree.cc create mode 100644 internal/arena_tree.h create mode 100644 internal/arena_tree_test.cc diff --git a/internal/BUILD b/internal/BUILD index 189853323..5b44efbad 100644 --- a/internal/BUILD +++ b/internal/BUILD @@ -42,6 +42,26 @@ cc_test( ], ) +cc_library( + name = "arena_tree", + srcs = ["arena_tree.cc"], + hdrs = ["arena_tree.h"], + deps = [ + "@com_google_absl//absl/base:nullability", + "@com_google_absl//absl/log:absl_check", + "@com_google_protobuf//:protobuf", + ], +) + +cc_test( + name = "arena_tree_test", + srcs = ["arena_tree_test.cc"], + deps = [ + ":arena_tree", + ":testing", + ], +) + cc_library( name = "new", srcs = ["new.cc"], diff --git a/internal/arena_tree.cc b/internal/arena_tree.cc new file mode 100644 index 000000000..9bd2922c3 --- /dev/null +++ b/internal/arena_tree.cc @@ -0,0 +1,352 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "internal/arena_tree.h" + +#include "absl/base/nullability.h" + +namespace cel::internal { + +namespace { + +void ArenaTreeRotateLeft(ArenaTreeNodeBase** head, ArenaTreeNodeBase* elem) { + ArenaTreeNodeBase* tmp = ArenaTreeNodeGetRight(elem); + if (ArenaTreeNodeSetRight(elem, ArenaTreeNodeGetLeft(tmp)) != nullptr) { + ArenaTreeNodeSetParent(ArenaTreeNodeGetLeft(tmp), elem); + } + ArenaTreeNodeBase* parent = ArenaTreeNodeGetParent(elem); + ArenaTreeNodeSetParent(tmp, parent); + if (parent != nullptr) { + if (elem == ArenaTreeNodeGetLeft(parent)) { + ArenaTreeNodeSetLeft(parent, tmp); + } else { + ArenaTreeNodeSetRight(parent, tmp); + } + } else { + *head = tmp; + } + ArenaTreeNodeSetLeft(tmp, elem); + ArenaTreeNodeSetParent(elem, tmp); +} + +void ArenaTreeRotateRight(ArenaTreeNodeBase** head, ArenaTreeNodeBase* elem) { + ArenaTreeNodeBase* tmp = ArenaTreeNodeGetLeft(elem); + if (ArenaTreeNodeSetLeft(elem, ArenaTreeNodeGetRight(tmp)) != nullptr) { + ArenaTreeNodeSetParent(ArenaTreeNodeGetRight(tmp), elem); + } + ArenaTreeNodeBase* parent = ArenaTreeNodeGetParent(elem); + ArenaTreeNodeSetParent(tmp, parent); + if (parent != nullptr) { + if (elem == ArenaTreeNodeGetLeft(parent)) { + ArenaTreeNodeSetLeft(parent, tmp); + } else { + ArenaTreeNodeSetRight(parent, tmp); + } + } else { + *head = tmp; + } + ArenaTreeNodeSetRight(tmp, elem); + ArenaTreeNodeSetParent(elem, tmp); +} + +void ArenaTreeRemoveColor(ArenaTreeNodeBase** head, ArenaTreeNodeBase* parent, + ArenaTreeNodeBase* elem) { + ArenaTreeNodeBase* tmp; + while ((elem == nullptr || + ArenaTreeNodeGetColor(elem) == ArenaTreeNodeColor::kBlack) && + elem != *head) { + if (ArenaTreeNodeGetLeft(parent) == elem) { + tmp = ArenaTreeNodeGetRight(parent); + if (ArenaTreeNodeGetColor(tmp) == ArenaTreeNodeColor::kRed) { + ArenaTreeNodeSetColor(tmp, ArenaTreeNodeColor::kBlack); + ArenaTreeNodeSetColor(parent, ArenaTreeNodeColor::kRed); + ArenaTreeRotateLeft(head, parent); + tmp = ArenaTreeNodeGetRight(parent); + } + if ((ArenaTreeNodeGetLeft(tmp) == nullptr || + ArenaTreeNodeGetColor(ArenaTreeNodeGetLeft(tmp)) == + ArenaTreeNodeColor::kBlack) && + (ArenaTreeNodeGetRight(tmp) == nullptr || + ArenaTreeNodeGetColor(ArenaTreeNodeGetRight(tmp)) == + ArenaTreeNodeColor::kBlack)) { + ArenaTreeNodeSetColor(tmp, ArenaTreeNodeColor::kRed); + elem = parent; + parent = ArenaTreeNodeGetParent(elem); + } else { + if (ArenaTreeNodeGetRight(tmp) == nullptr || + ArenaTreeNodeGetColor(ArenaTreeNodeGetRight(tmp)) == + ArenaTreeNodeColor::kBlack) { + ArenaTreeNodeBase* left; + if ((left = ArenaTreeNodeGetLeft(tmp)) != nullptr) { + ArenaTreeNodeSetColor(left, ArenaTreeNodeColor::kBlack); + } + ArenaTreeNodeSetColor(tmp, ArenaTreeNodeColor::kRed); + ArenaTreeRotateRight(head, tmp); + tmp = ArenaTreeNodeGetRight(parent); + } + ArenaTreeNodeSetColor(tmp, ArenaTreeNodeGetColor(parent)); + ArenaTreeNodeSetColor(parent, ArenaTreeNodeColor::kBlack); + if (ArenaTreeNodeGetRight(tmp) != nullptr) { + ArenaTreeNodeSetColor(ArenaTreeNodeGetRight(tmp), + ArenaTreeNodeColor::kBlack); + } + ArenaTreeRotateLeft(head, parent); + elem = *head; + break; + } + } else { + tmp = ArenaTreeNodeGetLeft(parent); + if (ArenaTreeNodeGetColor(tmp) == ArenaTreeNodeColor::kRed) { + ArenaTreeNodeSetColor(tmp, ArenaTreeNodeColor::kBlack); + ArenaTreeNodeSetColor(parent, ArenaTreeNodeColor::kRed); + ArenaTreeRotateRight(head, parent); + tmp = ArenaTreeNodeGetLeft(parent); + } + if ((ArenaTreeNodeGetLeft(tmp) == nullptr || + ArenaTreeNodeGetColor(ArenaTreeNodeGetLeft(tmp)) == + ArenaTreeNodeColor::kBlack) && + (ArenaTreeNodeGetRight(tmp) == nullptr || + ArenaTreeNodeGetColor(ArenaTreeNodeGetRight(tmp)) == + ArenaTreeNodeColor::kBlack)) { + ArenaTreeNodeSetColor(tmp, ArenaTreeNodeColor::kRed); + elem = parent; + parent = ArenaTreeNodeGetParent(elem); + } else { + if (ArenaTreeNodeGetLeft(tmp) == nullptr || + ArenaTreeNodeGetColor(ArenaTreeNodeGetLeft(tmp)) == + ArenaTreeNodeColor::kBlack) { + ArenaTreeNodeBase* right; + if ((right = ArenaTreeNodeGetRight(tmp)) != nullptr) { + ArenaTreeNodeSetColor(right, ArenaTreeNodeColor::kBlack); + } + ArenaTreeNodeSetColor(tmp, ArenaTreeNodeColor::kRed); + ArenaTreeRotateLeft(head, tmp); + tmp = ArenaTreeNodeGetLeft(parent); + } + ArenaTreeNodeSetColor(tmp, ArenaTreeNodeGetColor(parent)); + ArenaTreeNodeSetColor(parent, ArenaTreeNodeColor::kBlack); + if (ArenaTreeNodeGetLeft(tmp) != nullptr) { + ArenaTreeNodeSetColor(ArenaTreeNodeGetLeft(tmp), + ArenaTreeNodeColor::kBlack); + } + ArenaTreeRotateRight(head, parent); + elem = *head; + break; + } + } + } + if (elem != nullptr) { + ArenaTreeNodeSetColor(elem, ArenaTreeNodeColor::kBlack); + } +} + +} // namespace + +const ArenaTreeNodeBase* absl_nullable ArenaTreeNext( + const ArenaTreeNodeBase* absl_nullable node) { + if (node != nullptr) { + const ArenaTreeNodeBase* right = ArenaTreeNodeGetRight(node); + if (right != nullptr) { + node = right; + const ArenaTreeNodeBase* left; + while ((left = ArenaTreeNodeGetLeft(node)) != nullptr) { + node = left; + } + } else { + const ArenaTreeNodeBase* parent = ArenaTreeNodeGetParent(node); + if (parent != nullptr && node == ArenaTreeNodeGetLeft(parent)) { + node = parent; + } else { + while (parent != nullptr && ArenaTreeNodeGetRight(parent) == node) { + node = ArenaTreeNodeGetParent(node); + parent = ArenaTreeNodeGetParent(node); + } + node = parent; + } + } + } + return node; +} + +const ArenaTreeNodeBase* absl_nullable ArenaTreePrev( + const ArenaTreeNodeBase* absl_nullable node) { + if (node != nullptr) { + const ArenaTreeNodeBase* left = ArenaTreeNodeGetLeft(node); + if (left != nullptr) { + node = left; + const ArenaTreeNodeBase* right; + while ((right = ArenaTreeNodeGetRight(node)) != nullptr) { + node = right; + } + } else { + const ArenaTreeNodeBase* parent = ArenaTreeNodeGetParent(node); + if (parent != nullptr && node == ArenaTreeNodeGetRight(parent)) { + node = parent; + } else { + while (parent != nullptr && ArenaTreeNodeGetLeft(parent) == node) { + node = ArenaTreeNodeGetParent(node); + parent = ArenaTreeNodeGetParent(node); + } + node = parent; + } + } + } + return node; +} + +const ArenaTreeNodeBase* absl_nullable ArenaTreeMin( + const ArenaTreeNodeBase* absl_nullable node) { + const ArenaTreeNodeBase* tmp = node; + const ArenaTreeNodeBase* parent = nullptr; + while (tmp != nullptr) { + parent = tmp; + tmp = ArenaTreeNodeGetLeft(tmp); + } + return parent; +} + +const ArenaTreeNodeBase* absl_nullable ArenaTreeMax( + const ArenaTreeNodeBase* absl_nullable node) { + const ArenaTreeNodeBase* tmp = node; + const ArenaTreeNodeBase* parent = nullptr; + while (tmp != nullptr) { + parent = tmp; + tmp = ArenaTreeNodeGetRight(tmp); + } + return parent; +} + +void ArenaTreeRemove(ArenaTreeNodeBase* absl_nullable* absl_nonnull head, + ArenaTreeNodeBase* absl_nonnull elem) { + ArenaTreeNodeBase* child; + ArenaTreeNodeBase* parent; + ArenaTreeNodeBase* const old = elem; + ArenaTreeNodeColor color; + if (ArenaTreeNodeGetLeft(elem) == nullptr) { + child = ArenaTreeNodeGetRight(elem); + } else if (ArenaTreeNodeGetRight(elem) == nullptr) { + child = ArenaTreeNodeGetLeft(elem); + } else { + ArenaTreeNodeBase* left; + elem = ArenaTreeNodeGetRight(elem); + while ((left = ArenaTreeNodeGetLeft(elem)) != nullptr) { + elem = left; + } + child = ArenaTreeNodeGetRight(elem); + parent = ArenaTreeNodeGetParent(elem); + color = ArenaTreeNodeGetColor(elem); + if (child != nullptr) { + ArenaTreeNodeSetParent(child, parent); + } + if (parent != nullptr) { + if (ArenaTreeNodeGetLeft(parent) == elem) { + ArenaTreeNodeSetLeft(parent, child); + } else { + ArenaTreeNodeSetRight(parent, child); + } + } else { + *head = child; + } + if (ArenaTreeNodeGetParent(elem) == old) { + parent = elem; + } + ArenaTreeNodeBase* old_parent = ArenaTreeNodeGetParent(old); + if (old_parent != nullptr) { + if (ArenaTreeNodeGetLeft(old_parent) == old) { + ArenaTreeNodeSetLeft(old_parent, elem); + } else { + ArenaTreeNodeSetRight(old_parent, elem); + } + } else { + *head = elem; + } + ArenaTreeNodeSetParent(ArenaTreeNodeGetLeft(old), elem); + if (ArenaTreeNodeGetRight(old) != nullptr) { + ArenaTreeNodeSetParent(ArenaTreeNodeGetRight(old), elem); + } + goto color; + } + parent = ArenaTreeNodeGetParent(elem); + color = ArenaTreeNodeGetColor(elem); + if (child != nullptr) { + ArenaTreeNodeSetParent(child, parent); + } + if (parent != nullptr) { + if (ArenaTreeNodeGetLeft(parent) == elem) { + ArenaTreeNodeSetLeft(parent, child); + } else { + ArenaTreeNodeSetRight(parent, child); + } + } else { + *head = child; + } +color: + if (color == ArenaTreeNodeColor::kBlack) { + ArenaTreeRemoveColor(head, parent, child); + } + ArenaTreeNodeClear(old); +} + +void ArenaTreeInsertColor(ArenaTreeNodeBase* absl_nullable* absl_nonnull head, + ArenaTreeNodeBase* absl_nonnull elem) { + ArenaTreeNodeBase* parent; + ArenaTreeNodeBase* grandparent; + ArenaTreeNodeBase* tmp; + while ((parent = ArenaTreeNodeGetParent(elem)) != nullptr && + ArenaTreeNodeGetColor(parent) == ArenaTreeNodeColor::kRed) { + grandparent = ArenaTreeNodeGetParent(parent); + if (parent == ArenaTreeNodeGetLeft(grandparent)) { + tmp = ArenaTreeNodeGetRight(grandparent); + if (tmp != nullptr && + ArenaTreeNodeGetColor(tmp) == ArenaTreeNodeColor::kRed) { + ArenaTreeNodeSetColor(tmp, ArenaTreeNodeColor::kBlack); + ArenaTreeNodeSetColor(parent, ArenaTreeNodeColor::kBlack); + ArenaTreeNodeSetColor(grandparent, ArenaTreeNodeColor::kRed); + elem = grandparent; + continue; + } + if (ArenaTreeNodeGetRight(parent) == elem) { + ArenaTreeRotateLeft(head, parent); + tmp = parent; + parent = elem; + elem = tmp; + } + ArenaTreeNodeSetColor(parent, ArenaTreeNodeColor::kBlack); + ArenaTreeNodeSetColor(grandparent, ArenaTreeNodeColor::kRed); + ArenaTreeRotateRight(head, grandparent); + } else { + tmp = ArenaTreeNodeGetLeft(grandparent); + if (tmp != nullptr && + ArenaTreeNodeGetColor(tmp) == ArenaTreeNodeColor::kRed) { + ArenaTreeNodeSetColor(tmp, ArenaTreeNodeColor::kBlack); + ArenaTreeNodeSetColor(parent, ArenaTreeNodeColor::kBlack); + ArenaTreeNodeSetColor(grandparent, ArenaTreeNodeColor::kRed); + elem = grandparent; + continue; + } + if (ArenaTreeNodeGetLeft(parent) == elem) { + ArenaTreeRotateRight(head, parent); + tmp = parent; + parent = elem; + elem = tmp; + } + ArenaTreeNodeSetColor(parent, ArenaTreeNodeColor::kBlack); + ArenaTreeNodeSetColor(grandparent, ArenaTreeNodeColor::kRed); + ArenaTreeRotateLeft(head, grandparent); + } + } + ArenaTreeNodeSetColor(*head, ArenaTreeNodeColor::kBlack); +} + +} // namespace cel::internal diff --git a/internal/arena_tree.h b/internal/arena_tree.h new file mode 100644 index 000000000..9e76bd84a --- /dev/null +++ b/internal/arena_tree.h @@ -0,0 +1,398 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// ArenaTree is a low level implementation of an RBTree tailored for use with +// google::protobuf::Arena. + +#ifndef THIRD_PARTY_CEL_CPP_INTERNAL_ARENA_TREE_H_ +#define THIRD_PARTY_CEL_CPP_INTERNAL_ARENA_TREE_H_ + +#include +#include +#include +#include + +#include "absl/base/nullability.h" +#include "absl/log/absl_check.h" +#include "google/protobuf/arena.h" + +namespace cel::internal { + +enum class ArenaTreeNodeColor : uintptr_t { + kBlack = 0, + kRed = 1, +}; + +struct ArenaTreeNodeBase; + +[[nodiscard]] +const ArenaTreeNodeBase* absl_nullable ArenaTreeNext( + const ArenaTreeNodeBase* absl_nullable node); + +[[nodiscard]] +const ArenaTreeNodeBase* absl_nullable ArenaTreePrev( + const ArenaTreeNodeBase* absl_nullable node); + +[[nodiscard]] +const ArenaTreeNodeBase* absl_nullable ArenaTreeMin( + const ArenaTreeNodeBase* absl_nullable node); + +[[nodiscard]] +const ArenaTreeNodeBase* absl_nullable ArenaTreeMax( + const ArenaTreeNodeBase* absl_nullable node); + +template +using IsDerivedFromArenaTreeNodeBase = std::conjunction< + std::is_base_of, + std::negation>>>; + +template +constexpr bool kIsDerivedFromArenaTreeNodeBase = + IsDerivedFromArenaTreeNodeBase::value; + +template +using EnableIfDerivedFromArenaTreeNodeBase = + std::enable_if_t, U>; + +template +[[nodiscard]] +inline EnableIfDerivedFromArenaTreeNodeBase ArenaTreeNext( + T* absl_nullable node) { + return static_cast(const_cast( + (ArenaTreeNext)(static_cast(node)))); +} + +template +[[nodiscard]] +inline EnableIfDerivedFromArenaTreeNodeBase ArenaTreePrev( + T* absl_nullable node) { + return static_cast(const_cast( + (ArenaTreePrev)(static_cast(node)))); +} + +template +[[nodiscard]] +inline EnableIfDerivedFromArenaTreeNodeBase ArenaTreeMin( + T* absl_nullable node) { + return static_cast(const_cast( + (ArenaTreeMin)(static_cast(node)))); +} + +template +[[nodiscard]] +inline EnableIfDerivedFromArenaTreeNodeBase ArenaTreeMax( + T* absl_nullable node) { + return static_cast(const_cast( + (ArenaTreeMax)(static_cast(node)))); +} + +[[nodiscard]] +ArenaTreeNodeBase* absl_nullability_unknown ArenaTreeNodeGetParent( + const ArenaTreeNodeBase* absl_nonnull node); + +void ArenaTreeNodeSetParent(ArenaTreeNodeBase* absl_nonnull node, + ArenaTreeNodeBase* absl_nullability_unknown parent); + +[[nodiscard]] +ArenaTreeNodeColor ArenaTreeNodeGetColor( + const ArenaTreeNodeBase* absl_nonnull node); + +void ArenaTreeNodeSetColor(ArenaTreeNodeBase* absl_nonnull node, + ArenaTreeNodeColor color); + +void ArenaTreeNodeSet(ArenaTreeNodeBase* absl_nonnull node, + ArenaTreeNodeBase* absl_nullable parent); + +void ArenaTreeNodeClear(ArenaTreeNodeBase* absl_nonnull node); + +[[nodiscard]] +ArenaTreeNodeBase* absl_nullability_unknown ArenaTreeNodeGetLeft( + const ArenaTreeNodeBase* absl_nonnull node); + +ArenaTreeNodeBase* absl_nullability_unknown ArenaTreeNodeSetLeft( + ArenaTreeNodeBase* absl_nonnull node, + ArenaTreeNodeBase* absl_nullability_unknown left); + +[[nodiscard]] +ArenaTreeNodeBase* absl_nullability_unknown ArenaTreeNodeGetRight( + const ArenaTreeNodeBase* absl_nonnull node); + +ArenaTreeNodeBase* absl_nullability_unknown ArenaTreeNodeSetRight( + ArenaTreeNodeBase* absl_nonnull node, + ArenaTreeNodeBase* absl_nullability_unknown right); + +struct ArenaTreeNodeBase { + private: + uintptr_t parent_and_color = 0; + ArenaTreeNodeBase* absl_nullable left = nullptr; + ArenaTreeNodeBase* absl_nullable right = nullptr; + + friend ArenaTreeNodeBase* absl_nullability_unknown ArenaTreeNodeGetParent( + const ArenaTreeNodeBase* absl_nonnull node); + friend void ArenaTreeNodeSetParent( + ArenaTreeNodeBase* absl_nonnull node, + ArenaTreeNodeBase* absl_nullability_unknown parent); + friend ArenaTreeNodeColor ArenaTreeNodeGetColor( + const ArenaTreeNodeBase* absl_nonnull node); + friend void ArenaTreeNodeSetColor(ArenaTreeNodeBase* absl_nonnull node, + ArenaTreeNodeColor color); + friend void ArenaTreeNodeSet(ArenaTreeNodeBase* absl_nonnull node, + ArenaTreeNodeBase* absl_nullable parent); + friend void ArenaTreeNodeClear(ArenaTreeNodeBase* absl_nonnull node); + friend ArenaTreeNodeBase* absl_nullability_unknown ArenaTreeNodeGetLeft( + const ArenaTreeNodeBase* absl_nonnull node); + friend ArenaTreeNodeBase* absl_nullability_unknown ArenaTreeNodeSetLeft( + ArenaTreeNodeBase* absl_nonnull node, + ArenaTreeNodeBase* absl_nullability_unknown left); + friend ArenaTreeNodeBase* absl_nullability_unknown ArenaTreeNodeGetRight( + const ArenaTreeNodeBase* absl_nonnull node); + friend ArenaTreeNodeBase* absl_nullability_unknown ArenaTreeNodeSetRight( + ArenaTreeNodeBase* absl_nonnull node, + ArenaTreeNodeBase* absl_nullability_unknown right); +}; + +[[nodiscard]] +inline ArenaTreeNodeBase* absl_nullability_unknown ArenaTreeNodeGetParent( + const ArenaTreeNodeBase* absl_nonnull node) { + return reinterpret_cast(node->parent_and_color & + ~uintptr_t{1}); +} + +inline void ArenaTreeNodeSetParent( + ArenaTreeNodeBase* absl_nonnull node, + ArenaTreeNodeBase* absl_nullability_unknown parent) { + node->parent_and_color = static_cast(ArenaTreeNodeGetColor(node)) | + reinterpret_cast(parent); +} + +[[nodiscard]] +inline ArenaTreeNodeColor ArenaTreeNodeGetColor( + const ArenaTreeNodeBase* absl_nonnull node) { + return static_cast(node->parent_and_color & uintptr_t{1}); +} + +inline void ArenaTreeNodeSetColor(ArenaTreeNodeBase* absl_nonnull node, + ArenaTreeNodeColor color) { + node->parent_and_color = + reinterpret_cast(ArenaTreeNodeGetParent(node)) | + static_cast(color); +} + +inline void ArenaTreeNodeSet(ArenaTreeNodeBase* absl_nonnull node, + ArenaTreeNodeBase* absl_nullable parent) { + node->parent_and_color = reinterpret_cast(parent) | + static_cast(ArenaTreeNodeColor::kRed); + node->left = node->right = nullptr; +} + +inline void ArenaTreeNodeClear(ArenaTreeNodeBase* absl_nonnull node) { + node->parent_and_color = 0; + node->left = node->right = nullptr; +} + +[[nodiscard]] +inline ArenaTreeNodeBase* absl_nullability_unknown ArenaTreeNodeGetLeft( + const ArenaTreeNodeBase* absl_nonnull node) { + return node->left; +} + +inline ArenaTreeNodeBase* absl_nullability_unknown ArenaTreeNodeSetLeft( + ArenaTreeNodeBase* absl_nonnull node, + ArenaTreeNodeBase* absl_nullability_unknown left) { + return node->left = left; +} + +[[nodiscard]] +inline ArenaTreeNodeBase* absl_nullability_unknown ArenaTreeNodeGetRight( + const ArenaTreeNodeBase* absl_nonnull node) { + return node->right; +} + +inline ArenaTreeNodeBase* absl_nullability_unknown ArenaTreeNodeSetRight( + ArenaTreeNodeBase* absl_nonnull node, + ArenaTreeNodeBase* absl_nullability_unknown right) { + return node->right = right; +} + +void ArenaTreeRemove(ArenaTreeNodeBase* absl_nullable* absl_nonnull head, + ArenaTreeNodeBase* absl_nonnull elem); + +template +[[nodiscard]] +inline EnableIfDerivedFromArenaTreeNodeBase ArenaTreeRemove( + T* absl_nullable* absl_nonnull head, T* absl_nonnull elem) { + ArenaTreeNodeBase* head_base = *head; + (ArenaTreeRemove)(&head_base, static_cast(elem)); + *head = static_cast(head_base); +} + +void ArenaTreeInsertColor(ArenaTreeNodeBase* absl_nullable* absl_nonnull head, + ArenaTreeNodeBase* absl_nonnull elem); + +template +struct ArenaTreeNode; + +template +[[nodiscard]] +const T& ArenaTreeNodeGetValue(const ArenaTreeNode* absl_nonnull node); + +template +struct ArenaTreeNode : ArenaTreeNodeBase { + template + explicit ArenaTreeNode(Args&&... args) + : ArenaTreeNodeBase(), value(std::forward(args)...) {} + + private: + template + friend const U& ArenaTreeNodeGetValue( + const ArenaTreeNode* absl_nonnull node); + + T value; +}; + +template +[[nodiscard]] inline const T& ArenaTreeNodeGetValue( + const ArenaTreeNode* absl_nonnull node) { + return node->value; +} + +template +struct ArenaTreeNodeCrtp : ArenaTreeNodeBase { + using ArenaTreeNodeBase::ArenaTreeNodeBase; +}; + +template +[[nodiscard]] inline const T& ArenaTreeNodeGetValue( + const ArenaTreeNodeCrtp* absl_nonnull node) { + return *static_cast(node); +} + +template +[[nodiscard]] +T* absl_nonnull ArenaTreeInsert(T* absl_nullable* absl_nonnull head, + T* absl_nonnull elem, const Compare& compare) { + T* tmp = *head; + T* parent = nullptr; + int diff = 0; + while (tmp != nullptr) { + parent = tmp; + diff = std::invoke(compare, (ArenaTreeNodeGetValue)(elem), + (ArenaTreeNodeGetValue)(parent)); + if (diff < 0) { + tmp = static_cast((ArenaTreeNodeGetLeft)(tmp)); + } else if (diff > 0) { + tmp = static_cast((ArenaTreeNodeGetRight)(tmp)); + } else { + return tmp; + } + } + (ArenaTreeNodeSet)(elem, parent); + if (parent != nullptr) { + if (diff < 0) { + (ArenaTreeNodeSetLeft)(parent, elem); + } else { + (ArenaTreeNodeSetRight)(parent, elem); + } + } else { + *head = elem; + } + ArenaTreeNodeBase* head_base = *head; + (ArenaTreeInsertColor)(&head_base, elem); + *head = static_cast(head_base); + return elem; +} + +template +struct ArenaTreeNodeConstructor { + template + void operator()(Args&&... args) const { + ABSL_DCHECK(*out == nullptr); + *out = google::protobuf::Arena::Create(arena, std::forward(args)...); + } + + google::protobuf::Arena* const absl_nonnull arena; + T** out; +}; + +template +[[nodiscard]] +std::pair ArenaTreeLazyEmplace( + google::protobuf::Arena* absl_nonnull arena, T* absl_nullable* absl_nonnull head, + const K& key, const Compare& compare, Emplacer&& emplacer) { + T* tmp = *head; + T* parent = nullptr; + int diff = 0; + while (tmp != nullptr) { + parent = tmp; + diff = std::invoke(compare, key, (ArenaTreeNodeGetValue)(parent)); + if (diff < 0) { + tmp = static_cast((ArenaTreeNodeGetLeft)(tmp)); + } else if (diff > 0) { + tmp = static_cast((ArenaTreeNodeGetRight)(tmp)); + } else { + return {tmp, false}; + } + } + T* elem = nullptr; + ArenaTreeNodeConstructor constructor{ + .arena = arena, + .out = &elem, + }; + std::invoke(std::forward(emplacer), + static_cast&>(constructor)); + ABSL_DCHECK(elem != nullptr); + (ArenaTreeNodeSet)(elem, parent); + if (parent != nullptr) { + if (diff < 0) { + (ArenaTreeNodeSetLeft)(parent, elem); + } else { + (ArenaTreeNodeSetRight)(parent, elem); + } + } else { + *head = elem; + } + ArenaTreeNodeBase* head_base = *head; + (ArenaTreeInsertColor)(&head_base, elem); + *head = static_cast(head_base); + return {elem, true}; +} + +template +[[nodiscard]] +const T* absl_nullable ArenaTreeFind(const T* absl_nullable head, const K& key, + const Compare& compare) { + const T* tmp = head; + while (tmp != nullptr) { + int diff = std::invoke(compare, key, (ArenaTreeNodeGetValue)(tmp)); + if (diff < 0) { + tmp = static_cast((ArenaTreeNodeGetLeft)(tmp)); + } else if (diff > 0) { + tmp = static_cast((ArenaTreeNodeGetRight)(tmp)); + } else { + return tmp; + } + } + return nullptr; +} + +template +[[nodiscard]] +T* absl_nullable ArenaTreeFind(T* absl_nullable head, const K& key, + const Compare& compare) { + return (ArenaTreeFind)(static_cast(head), key, compare); +} + +} // namespace cel::internal + +#endif // THIRD_PARTY_CEL_CPP_INTERNAL_ARENA_TREE_H_ diff --git a/internal/arena_tree_test.cc b/internal/arena_tree_test.cc new file mode 100644 index 000000000..ddc6b9885 --- /dev/null +++ b/internal/arena_tree_test.cc @@ -0,0 +1,223 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "internal/arena_tree.h" + +#include + +#include "internal/testing.h" + +namespace cel::internal { +namespace { + +using ::testing::IsNull; + +using TestNode = ArenaTreeNode; + +struct TestNodeCompare { + int operator()(int lhs, int rhs) const { + if (lhs < rhs) { + return -1; + } + if (lhs > rhs) { + return 1; + } + return 0; + } +}; + +struct TestNodeCrtp : ArenaTreeNodeCrtp { + explicit TestNodeCrtp(int value) : ArenaTreeNodeCrtp(), value(value) {} + + int value; +}; + +struct TestNodeCrtpCompare { + int operator()(int lhs, int rhs) const { + if (lhs < rhs) { + return -1; + } + if (lhs > rhs) { + return 1; + } + return 0; + } + + int operator()(int lhs, const TestNodeCrtp& rhs) const { + return (*this)(lhs, rhs.value); + } + + int operator()(const TestNodeCrtp& lhs, int rhs) const { + return (*this)(lhs.value, rhs); + } + + int operator()(const TestNodeCrtp& lhs, const TestNodeCrtp& rhs) const { + return (*this)(lhs.value, rhs.value); + } +}; + +TEST(ArenaTree, Empty) { + TestNode* head = nullptr; + EXPECT_THAT(ArenaTreePrev(head), IsNull()); + EXPECT_THAT(ArenaTreeNext(head), IsNull()); + EXPECT_THAT(ArenaTreeMin(head), IsNull()); + EXPECT_THAT(ArenaTreeMax(head), IsNull()); +} + +TEST(ArenaTree, Single) { + TestNode* head = nullptr; + std::unique_ptr node1 = std::make_unique(1); + EXPECT_EQ(ArenaTreeInsert(&head, node1.get(), TestNodeCompare{}), + node1.get()); + EXPECT_EQ(ArenaTreePrev(head), nullptr); + EXPECT_EQ(ArenaTreeNext(head), nullptr); + EXPECT_EQ(ArenaTreeMin(head), node1.get()); + EXPECT_EQ(ArenaTreeMax(head), node1.get()); + ArenaTreeRemove(&head, node1.get()); + EXPECT_THAT(ArenaTreePrev(head), IsNull()); + EXPECT_THAT(ArenaTreeNext(head), IsNull()); + EXPECT_THAT(ArenaTreeMin(head), IsNull()); + EXPECT_THAT(ArenaTreeMax(head), IsNull()); +} + +TEST(ArenaTree, Couple) { + TestNode* head = nullptr; + std::unique_ptr node1 = std::make_unique(1); + std::unique_ptr node2 = std::make_unique(2); + + EXPECT_EQ(ArenaTreeInsert(&head, node1.get(), TestNodeCompare{}), + node1.get()); + EXPECT_EQ(ArenaTreeInsert(&head, node2.get(), TestNodeCompare{}), + node2.get()); + EXPECT_EQ(ArenaTreePrev(node1.get()), nullptr); + EXPECT_EQ(ArenaTreePrev(node2.get()), node1.get()); + EXPECT_EQ(ArenaTreeNext(node1.get()), node2.get()); + EXPECT_EQ(ArenaTreeNext(node2.get()), nullptr); + EXPECT_EQ(ArenaTreeMin(head), node1.get()); + EXPECT_EQ(ArenaTreeMax(head), node2.get()); + + ArenaTreeRemove(&head, node1.get()); + EXPECT_THAT(ArenaTreePrev(head), IsNull()); + EXPECT_THAT(ArenaTreeNext(head), IsNull()); + EXPECT_EQ(ArenaTreeMin(head), node2.get()); + EXPECT_EQ(ArenaTreeMax(head), node2.get()); + + ArenaTreeRemove(&head, node2.get()); + EXPECT_THAT(ArenaTreePrev(head), IsNull()); + EXPECT_THAT(ArenaTreeNext(head), IsNull()); + EXPECT_THAT(ArenaTreeMin(head), IsNull()); + EXPECT_THAT(ArenaTreeMax(head), IsNull()); + + EXPECT_EQ(ArenaTreeInsert(&head, node2.get(), TestNodeCompare{}), + node2.get()); + EXPECT_EQ(ArenaTreeInsert(&head, node1.get(), TestNodeCompare{}), + node1.get()); + EXPECT_EQ(ArenaTreePrev(node1.get()), nullptr); + EXPECT_EQ(ArenaTreePrev(node2.get()), node1.get()); + EXPECT_EQ(ArenaTreeNext(node1.get()), node2.get()); + EXPECT_EQ(ArenaTreeNext(node2.get()), nullptr); + EXPECT_EQ(ArenaTreeMin(head), node1.get()); + EXPECT_EQ(ArenaTreeMax(head), node2.get()); + + ArenaTreeRemove(&head, node2.get()); + EXPECT_THAT(ArenaTreePrev(head), IsNull()); + EXPECT_THAT(ArenaTreeNext(head), IsNull()); + EXPECT_EQ(ArenaTreeMin(head), node1.get()); + EXPECT_EQ(ArenaTreeMax(head), node1.get()); + + ArenaTreeRemove(&head, node1.get()); + EXPECT_THAT(ArenaTreePrev(head), IsNull()); + EXPECT_THAT(ArenaTreeNext(head), IsNull()); + EXPECT_THAT(ArenaTreeMin(head), IsNull()); + EXPECT_THAT(ArenaTreeMax(head), IsNull()); +} + +TEST(ArenaTree, CrtpEmpty) { + TestNodeCrtp* head = nullptr; + EXPECT_THAT(ArenaTreePrev(head), IsNull()); + EXPECT_THAT(ArenaTreeNext(head), IsNull()); + EXPECT_THAT(ArenaTreeMin(head), IsNull()); + EXPECT_THAT(ArenaTreeMax(head), IsNull()); +} + +TEST(ArenaTree, CrtpSingle) { + TestNodeCrtp* head = nullptr; + std::unique_ptr node1 = std::make_unique(1); + EXPECT_EQ(ArenaTreeInsert(&head, node1.get(), TestNodeCrtpCompare{}), + node1.get()); + EXPECT_EQ(ArenaTreePrev(head), nullptr); + EXPECT_EQ(ArenaTreeNext(head), nullptr); + EXPECT_EQ(ArenaTreeMin(head), node1.get()); + EXPECT_EQ(ArenaTreeMax(head), node1.get()); + ArenaTreeRemove(&head, node1.get()); + EXPECT_THAT(ArenaTreePrev(head), IsNull()); + EXPECT_THAT(ArenaTreeNext(head), IsNull()); + EXPECT_THAT(ArenaTreeMin(head), IsNull()); + EXPECT_THAT(ArenaTreeMax(head), IsNull()); +} + +TEST(ArenaTree, CrtpCouple) { + TestNodeCrtp* head = nullptr; + std::unique_ptr node1 = std::make_unique(1); + std::unique_ptr node2 = std::make_unique(2); + + EXPECT_EQ(ArenaTreeInsert(&head, node1.get(), TestNodeCrtpCompare{}), + node1.get()); + EXPECT_EQ(ArenaTreeInsert(&head, node2.get(), TestNodeCrtpCompare{}), + node2.get()); + EXPECT_EQ(ArenaTreePrev(node1.get()), nullptr); + EXPECT_EQ(ArenaTreePrev(node2.get()), node1.get()); + EXPECT_EQ(ArenaTreeNext(node1.get()), node2.get()); + EXPECT_EQ(ArenaTreeNext(node2.get()), nullptr); + EXPECT_EQ(ArenaTreeMin(head), node1.get()); + EXPECT_EQ(ArenaTreeMax(head), node2.get()); + + ArenaTreeRemove(&head, node1.get()); + EXPECT_THAT(ArenaTreePrev(head), IsNull()); + EXPECT_THAT(ArenaTreeNext(head), IsNull()); + EXPECT_EQ(ArenaTreeMin(head), node2.get()); + EXPECT_EQ(ArenaTreeMax(head), node2.get()); + + ArenaTreeRemove(&head, node2.get()); + EXPECT_THAT(ArenaTreePrev(head), IsNull()); + EXPECT_THAT(ArenaTreeNext(head), IsNull()); + EXPECT_THAT(ArenaTreeMin(head), IsNull()); + EXPECT_THAT(ArenaTreeMax(head), IsNull()); + + EXPECT_EQ(ArenaTreeInsert(&head, node2.get(), TestNodeCrtpCompare{}), + node2.get()); + EXPECT_EQ(ArenaTreeInsert(&head, node1.get(), TestNodeCrtpCompare{}), + node1.get()); + EXPECT_EQ(ArenaTreePrev(node1.get()), nullptr); + EXPECT_EQ(ArenaTreePrev(node2.get()), node1.get()); + EXPECT_EQ(ArenaTreeNext(node1.get()), node2.get()); + EXPECT_EQ(ArenaTreeNext(node2.get()), nullptr); + EXPECT_EQ(ArenaTreeMin(head), node1.get()); + EXPECT_EQ(ArenaTreeMax(head), node2.get()); + + ArenaTreeRemove(&head, node2.get()); + EXPECT_THAT(ArenaTreePrev(head), IsNull()); + EXPECT_THAT(ArenaTreeNext(head), IsNull()); + EXPECT_EQ(ArenaTreeMin(head), node1.get()); + EXPECT_EQ(ArenaTreeMax(head), node1.get()); + + ArenaTreeRemove(&head, node1.get()); + EXPECT_THAT(ArenaTreePrev(head), IsNull()); + EXPECT_THAT(ArenaTreeNext(head), IsNull()); + EXPECT_THAT(ArenaTreeMin(head), IsNull()); + EXPECT_THAT(ArenaTreeMax(head), IsNull()); +} + +} // namespace +} // namespace cel::internal