uspace/lib/cpp/include/__bits/adt/rbtree.hpp

   1 /*
   2  * Copyright (c) 2018 Jaroslav Jindrak
   3  * All rights reserved.
   4  *
   5  * Redistribution and use in source and binary forms, with or without
   6  * modification, are permitted provided that the following conditions
   7  * are met:
   8  *
   9  * - Redistributions of source code must retain the above copyright
  10  *   notice, this list of conditions and the following disclaimer.
  11  * - Redistributions in binary form must reproduce the above copyright
  12  *   notice, this list of conditions and the following disclaimer in the
  13  *   documentation and/or other materials provided with the distribution.
  14  * - The name of the author may not be used to endorse or promote products
  15  *   derived from this software without specific prior written permission.
  16  *
  17  * THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
  18  * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
  19  * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
  20  * IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT,
  21  * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
  22  * NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
  23  * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
  24  * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
  25  * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
  26  * THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  27  */
  28
  29 #ifndef LIBCPP_BITS_ADT_RBTREE
  30 #define LIBCPP_BITS_ADT_RBTREE
  31
  32 #include <__bits/adt/key_extractors.hpp>
  33 #include <__bits/adt/rbtree_iterators.hpp>
  34 #include <__bits/adt/rbtree_node.hpp>
  35 #include <__bits/adt/rbtree_policies.hpp>
  36
  37 namespace std::aux
  38 {
  39     template<
  40         class Value, class Key, class KeyExtractor,
  41         class KeyComp, class Alloc, class Size,
  42         class Iterator, class ConstIterator,
  43         class Policy, class Node
  44     >
  45     class rbtree
  46     {
  47         public:
  48             using value_type     = Value;
  49             using key_type       = Key;
  50             using size_type      = Size;
  51             using allocator_type = Alloc;
  52             using key_compare    = KeyComp;
  53             using key_extract    = KeyExtractor;
  54
  55             using iterator       = Iterator;
  56             using const_iterator = ConstIterator;
  57
  58             using reverse_iterator       = std::reverse_iterator<iterator>;
  59             using const_reverse_iterator = std::reverse_iterator<const_iterator>;
  60
  61             using node_type = Node;
  62
  63             rbtree(const key_compare& kcmp = key_compare{})
  64                 : root_{nullptr}, size_{}, key_compare_{},
  65                   key_extractor_{}
  66             { /* DUMMY BODY */ }
  67
  68             rbtree(const rbtree& other)
  69                 : rbtree{other.key_compare_}
  70             {
  71                 for (const auto& x: other)
  72                     insert(x);
  73             }
  74
  75             rbtree(rbtree&& other)
  76                 : root_{other.root_}, size_{other.size_},
  77                   key_compare_{move(other.key_compare_)},
  78                   key_extractor_{move(other.key_extractor_)}
  79             {
  80                 other.root_ = nullptr;
  81                 other.size_ = size_type{};
  82             }
  83
  84             rbtree& operator=(const rbtree& other)
  85             {
  86                 auto tmp{other};
  87                 tmp.swap(*this);
  88
  89                 return *this;
  90             }
  91
  92             rbtree& operator=(rbtree&& other)
  93             {
  94                 rbtree tmp{move(other)};
  95                 tmp.swap(*this);
  96
  97                 return *this;
  98             }
  99
 100             bool empty() const noexcept
 101             {
 102                 return size_ == 0U;
 103             }
 104
 105             size_type size() const noexcept
 106             {
 107                 return size_;
 108             }
 109
 110             size_type max_size(allocator_type& alloc)
 111             {
 112                 return allocator_traits<allocator_type>::max_size(alloc);
 113             }
 114
 115             iterator begin()
 116             {
 117                 return iterator{find_smallest_(), false};
 118             }
 119
 120             const_iterator begin() const
 121             {
 122                 return cbegin();
 123             }
 124
 125             iterator end()
 126             {
 127                 /**
 128                  * In case we have lists of nodes
 129                  * we need to get the actual end
 130                  * from the largest node.
 131                  */
 132                 auto res = find_largest_();
 133                 if (res)
 134                     return iterator{res->get_end(), true};
 135                 else
 136                     return iterator{res, true};
 137             }
 138
 139             const_iterator end() const
 140             {
 141                 return cend();
 142             }
 143
 144             reverse_iterator rbegin()
 145             {
 146                 return make_reverse_iterator(end());
 147             }
 148
 149             const_reverse_iterator rbegin() const
 150             {
 151                 return make_reverse_iterator(cend());
 152             }
 153
 154             reverse_iterator rend()
 155             {
 156                 return make_reverse_iterator(begin());
 157             }
 158
 159             const_reverse_iterator rend() const
 160             {
 161                 return make_reverse_iterator(cbegin());
 162             }
 163
 164             const_iterator cbegin() const
 165             {
 166                 return const_iterator{find_smallest_(), false};
 167             }
 168
 169             const_iterator cend() const
 170             {
 171                 auto res = find_largest_();
 172                 if (res)
 173                     return const_iterator{res->get_end(), true};
 174                 else
 175                     return const_iterator{res, true};
 176             }
 177
 178             const_reverse_iterator crbegin() const
 179             {
 180                 return make_reverse_iterator(cend());
 181             }
 182
 183             const_reverse_iterator crend() const
 184             {
 185                 return make_reverse_iterator(cbegin());
 186             }
 187
 188             template<class... Args>
 189             auto emplace(Args&&... args)
 190             {
 191                 return Policy::emplace(*this, forward<Args>(args)...);
 192             }
 193
 194             auto insert(const value_type& val)
 195             {
 196                 return Policy::insert(*this, val);
 197             }
 198
 199             auto insert(value_type&& val)
 200             {
 201                 return Policy::insert(*this, forward<value_type>(val));
 202             }
 203
 204             size_type erase(const key_type& key)
 205             {
 206                 return Policy::erase(*this, key);
 207             }
 208
 209             iterator erase(const_iterator it)
 210             {
 211                 if (it == cend())
 212                     return end();
 213
 214                 auto node = const_cast<node_type*>(it.node());
 215
 216                 node = delete_node(node);
 217                 if (!node)
 218                     return iterator{find_largest_(), true};
 219                 else
 220                     return iterator{const_cast<node_type*>(node), false};
 221             }
 222
 223             void clear() noexcept
 224             {
 225                 if (root_)
 226                 {
 227                     delete root_;
 228                     root_ = nullptr;
 229                     size_ = size_type{};
 230                 }
 231             }
 232
 233             void swap(rbtree& other)
 234                 noexcept(allocator_traits<allocator_type>::is_always_equal::value &&
 235                          noexcept(std::swap(declval<KeyComp&>(), declval<KeyComp&>())))
 236             {
 237                 std::swap(root_, other.root_);
 238                 std::swap(size_, other.size_);
 239                 std::swap(key_compare_, other.key_compare_);
 240                 std::swap(key_extractor_, other.key_extractor_);
 241             }
 242
 243             key_compare key_comp() const
 244             {
 245                 return key_compare_;
 246             }
 247
 248             iterator find(const key_type& key)
 249             {
 250                 auto node = find_(key);
 251                 if (node)
 252                     return iterator{node, false};
 253                 else
 254                     return end();
 255             }
 256
 257             const_iterator find(const key_type& key) const
 258             {
 259                 auto node = find_(key);
 260                 if (node)
 261                     return const_iterator{node, false};
 262                 else
 263                     return end();
 264             }
 265
 266             size_type count(const key_type& key) const
 267             {
 268                 return Policy::count(*this, key);
 269             }
 270
 271             iterator upper_bound(const key_type& key)
 272             {
 273                 return Policy::upper_bound(*this, key);
 274             }
 275
 276             const_iterator upper_bound(const key_type& key) const
 277             {
 278                 return Policy::upper_bound_const(*this, key);
 279             }
 280
 281             iterator lower_bound(const key_type& key)
 282             {
 283                 return Policy::lower_bound(*this, key);
 284             }
 285
 286             const_iterator lower_bound(const key_type& key) const
 287             {
 288                 return Policy::lower_bound_const(*this, key);
 289             }
 290
 291             pair<iterator, iterator> equal_range(const key_type& key)
 292             {
 293                 return Policy::equal_range(*this, key);
 294             }
 295
 296             pair<const_iterator, const_iterator> equal_range(const key_type& key) const
 297             {
 298                 return Policy::equal_range_const(*this, key);
 299             }
 300
 301             bool is_eq_to(const rbtree& other) const
 302             {
 303                 if (size_ != other.size())
 304                     return false;
 305
 306                 auto it1 = begin();
 307                 auto it2 = other.begin();
 308
 309                 // TODO: this doesn't compare values :/
 310                 while (keys_equal(*it1++, *it2++))
 311                 { /* DUMMY BODY */ }
 312
 313                 return (it1 == end()) && (it2 == other.end());
 314             }
 315
 316             const key_type& get_key(const value_type& val) const
 317             {
 318                 return key_extractor_(val);
 319             }
 320
 321             bool keys_comp(const key_type& key, const value_type& val) const
 322             {
 323                 return key_compare_(key, key_extractor_(val));
 324             }
 325
 326             bool keys_equal(const key_type& k1, const key_type& k2) const
 327             {
 328                 return !key_compare_(k1, k2) && !key_compare_(k2, k1);
 329             }
 330
 331             node_type* find_parent_for_insertion(const key_type& key) const
 332             {
 333                 auto current = root_;
 334                 auto parent = current;
 335
 336                 while (current)
 337                 {
 338                     parent = current;
 339                     if (key_compare_(key, key_extractor_(current->value)))
 340                         current = current->left();
 341                     else if (key_compare_(key_extractor_(current->value), key))
 342                         current = current->right();
 343                     else
 344                         return current;
 345                 }
 346
 347                 return parent;
 348             }
 349
 350             node_type* delete_node(const node_type* n)
 351             {
 352                 auto node = const_cast<node_type*>(n);
 353                 if (!node)
 354                     return nullptr;
 355
 356                 --size_;
 357
 358                 auto succ = node->successor();
 359                 if (auto tmp = node->get_node_for_deletion(); tmp != nullptr)
 360                 {
 361                     /**
 362                      * This will kick in multi containers,
 363                      * we popped one node from a list of nodes
 364                      * with equivalent keys and we can delete it
 365                      * and return the successor which was the next
 366                      * in the list.
 367                      */
 368                     delete tmp;
 369
 370                     update_root_(succ); // Incase the first in list was root.
 371                     return succ;
 372                 }
 373                 else if (node == root_)
 374                 { // Only executed if root_ is unique.
 375                     root_ = nullptr;
 376                     delete node;
 377
 378                     return nullptr;
 379                 }
 380
 381                 if (node->left() && node->right())
 382                 {
 383                     node->swap(succ);
 384                     if (succ && !succ->parent())
 385                         root_ = succ;
 386
 387                     // Node now has at most one child.
 388                     // Also: If succ was nullptr, the swap
 389                     //       didn't do anything and we can
 390                     //       safely delete node.
 391                     return delete_node(node);
 392                 }
 393
 394                 auto child = node->right() ? node->right() : node->left();
 395                 if (!child)
 396                 {
 397                     // Simply remove the node.
 398                     // TODO: repair here too?
 399                     node->unlink();
 400                     delete node;
 401                 }
 402                 else
 403                 {
 404                     // Replace with the child.
 405                     child->parent(node->parent());
 406                     if (node->is_left_child())
 407                         child->parent()->left(child);
 408                     else if (node->is_right_child())
 409                         child->parent()->right(child);
 410                     node->parent(nullptr);
 411                     node->left(nullptr);
 412                     node->right(nullptr);
 413
 414                     // Repair if needed.
 415                     repair_after_erase_(node, child);
 416                     update_root_(child);
 417
 418                     delete node;
 419                 }
 420
 421                 return succ;
 422             }
 423
 424             void insert_node(node_type* node, node_type* parent)
 425             {
 426                 Policy::insert(*this, node, parent);
 427             }
 428
 429         private:
 430             node_type* root_;
 431             size_type size_;
 432             key_compare key_compare_;
 433             key_extract key_extractor_;
 434
 435             node_type* find_(const key_type& key) const
 436             {
 437                 auto current = root_;
 438                 while (current != nullptr)
 439                 {
 440                     if (key_compare_(key, key_extractor_(current->value)))
 441                         current = current->left();
 442                     else if (key_compare_(key_extractor_(current->value), key))
 443                         current = current->right();
 444                     else
 445                         return current;
 446                 }
 447
 448                 return nullptr;
 449             }
 450
 451             node_type* find_smallest_() const
 452             {
 453                 if (root_)
 454                     return root_->find_smallest();
 455                 else
 456                     return nullptr;
 457             }
 458
 459             node_type* find_largest_() const
 460             {
 461                 if (root_)
 462                     return root_->find_largest();
 463                 else
 464                     return nullptr;
 465             }
 466
 467             void update_root_(const node_type* node)
 468             {
 469                 if (!node)
 470                     return;
 471
 472                 root_ = const_cast<node_type*>(node);
 473                 while (root_->parent())
 474                     root_ = root_->parent();
 475             }
 476
 477             void repair_after_insert_(const node_type* node)
 478             {
 479                 // TODO: implement
 480             }
 481
 482             void repair_after_erase_(const node_type* node, const node_type* child)
 483             {
 484                 // TODO: implement
 485             }
 486
 487             friend Policy;
 488     };
 489 }
 490
 491 #endif