aa_tree.h

// Copyright 2024 The Pigweed Authors
//
// 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.
#pragma once
 
#include <cstddef>
#include <cstdint>
#include <functional>
#include <utility>
 
#include "pw_assert/assert.h"
#include "pw_containers/internal/aa_tree_item.h"
#include "pw_containers/internal/aa_tree_iterator.h"
#include "pw_containers/internal/intrusive_item.h"
#include "pw_function/function.h"
 
namespace pw::containers::internal {
 
class GenericAATree {
 public:
  using iterator = AATreeIterator<>;
 
  constexpr explicit GenericAATree(bool unique_keys)
      : unique_keys_(unique_keys) {}
 
  ~GenericAATree() { CheckIntrusiveContainerIsEmpty(empty()); }
 
  // Intrusive trees cannot be copied, since each Item can only be in one tree.
  GenericAATree(const GenericAATree&) = delete;
  GenericAATree& operator=(const GenericAATree&) = delete;
 
  [[nodiscard]] constexpr bool unique_keys() const { return unique_keys_; }
 
  void SetRoot(AATreeItem* item);
 
  // Iterators
 
  constexpr iterator begin() noexcept {
    return empty() ? iterator(&root_) : iterator(&root_, root_->GetLeftmost());
  }
 
  constexpr iterator end() noexcept {
    return empty() ? iterator(&root_)
                   : ++(iterator(&root_, root_->GetRightmost()));
  }
 
  // Capacity
 
  [[nodiscard]] constexpr bool empty() const { return root_ == nullptr; }
 
  size_t size() const;
 
  constexpr size_t max_size() const noexcept {
    return std::numeric_limits<ptrdiff_t>::max();
  }
 
  // Modifiers
 
  void clear();
 
  iterator erase_one(AATreeItem& item);
 
  iterator erase_range(AATreeItem& first, AATreeItem& last);
 
  void swap(GenericAATree& other);
 
 private:
  // The derived, templated type needs to be able to access the root element.
  template <typename>
  friend class KeyedAATree;
 
  AATreeItem* root_ = nullptr;
 
  const bool unique_keys_;
};
 
template <typename K>
class KeyedAATree : public GenericAATree {
 public:
  using Key = K;
  using Compare = Function<bool(Key, Key)>;
  using GetKey = Function<Key(const AATreeItem&)>;
  using GenericAATree::iterator;
 
  constexpr explicit KeyedAATree(bool unique_keys,
                                 Compare&& compare,
                                 GetKey&& get_key)
      : GenericAATree(unique_keys),
        compare_(std::move(compare)),
        get_key_(std::move(get_key)) {}
 
  // Modifiers
 
  void SetCompare(Compare&& compare) { compare_ = std::move(compare); }
  void SetGetKey(GetKey&& get_key) { get_key_ = std::move(get_key); }
 
  std::pair<iterator, bool> insert(AATreeItem& item);
 
  template <class Iterator>
  void insert(Iterator first, Iterator last);
 
  using GenericAATree::erase_one;
  using GenericAATree::erase_range;
 
  size_t erase_all(Key key);
 
  void merge(KeyedAATree<K>& other);
 
  // Lookup
 
  size_t count(Key key);
 
  iterator find(Key key);
 
  std::pair<iterator, iterator> equal_range(Key key);
 
  iterator lower_bound(Key key);
 
  iterator upper_bound(Key key);
 
 private:
  AATreeItem* InsertImpl(AATreeItem& parent,
                         AATreeItem& child,
                         AATreeItem*& duplicate);
 
  AATreeItem* GetLowerBoundImpl(AATreeItem* item, Key key);
 
  AATreeItem* GetUpperBoundImpl(AATreeItem* item, Key key);
 
  Compare compare_;
  GetKey get_key_;
};
 
template <typename K, typename V>
class AATree final : public KeyedAATree<K> {
 public:
  using Base = KeyedAATree<K>;
  using Key = typename Base::Key;
  using Compare = typename Base::Compare;
  using GetKey = Function<Key(const V&)>;
 
  class Item : public AATreeItem {
   public:
    constexpr explicit Item() = default;
 
   private:
    // IntrusiveItem is used to ensure T inherits from this container's Item
    // type. See also `CheckItemType`.
    template <typename, typename, bool>
    friend struct IntrusiveItem;
    using ItemType = V;
  };
 
  class Pair : public Item {
   public:
    constexpr explicit Pair(Key key) : key_(key) {}
    constexpr Key key() const { return key_; }
 
    constexpr bool operator<(const Pair& rhs) { return key_ < rhs.key(); }
 
   private:
    const Key key_;
  };
 
  constexpr AATree(bool unique_keys, Compare&& compare, GetKey&& get_key)
      : KeyedAATree<Key>(unique_keys,
                         std::move(compare),
                         std::move([this](const AATreeItem& item) -> Key {
                           return get_key_(static_cast<const V&>(item));
                         })),
        get_key_(std::move(get_key)) {}
 
  GetKey get_key_;
};
 
// Template method implementations.
 
template <typename K>
std::pair<GenericAATree::iterator, bool> KeyedAATree<K>::insert(
    AATreeItem& item) {
  CheckIntrusiveItemIsUncontained(!item.IsMapped());
  item.SetLevel(1);
  if (empty()) {
    SetRoot(&item);
    return std::make_pair(iterator(&root_, &item), true);
  }
  AATreeItem* duplicate = nullptr;
  SetRoot(InsertImpl(*root_, item, duplicate));
  if (duplicate == nullptr) {
    return std::make_pair(iterator(&root_, &item), true);
  }
  item.Reset();
  return std::make_pair(iterator(&root_, duplicate), false);
}
 
template <typename K>
AATreeItem* KeyedAATree<K>::InsertImpl(AATreeItem& parent,
                                       AATreeItem& child,
                                       AATreeItem*& duplicate) {
  if (compare_(get_key_(child), get_key_(parent))) {
    AATreeItem* left = parent.left_.get();
    left = left == nullptr ? &child : InsertImpl(*left, child, duplicate);
    parent.SetLeft(left);
 
  } else if (compare_(get_key_(parent), get_key_(child)) || !unique_keys()) {
    AATreeItem* right = parent.right_.get();
    right = right == nullptr ? &child : InsertImpl(*right, child, duplicate);
    parent.SetRight(right);
 
  } else {
    duplicate = &parent;
    return &parent;
  }
 
  return parent.Skew()->Split();
}
 
template <typename K>
template <class Iterator>
void KeyedAATree<K>::insert(Iterator first, Iterator last) {
  for (Iterator iter = first; iter != last; ++iter) {
    if constexpr (std::is_pointer_v<
                      typename std::iterator_traits<Iterator>::value_type>) {
      insert(**iter);
    } else {
      insert(*iter);
    }
  }
}
 
template <typename K>
size_t KeyedAATree<K>::erase_all(Key key) {
  size_t removed = 0;
  iterator iter = lower_bound(key);
  while (iter != end() && !compare_(key, get_key_(*iter))) {
    iter = erase_one(*iter);
    ++removed;
  }
  return removed;
}
 
template <typename K>
void KeyedAATree<K>::merge(KeyedAATree<K>& other) {
  auto iter = other.begin();
  while (!other.empty()) {
    AATreeItem& item = *iter;
    iter = other.erase_one(item);
    insert(item);
  }
}
 
template <typename K>
size_t KeyedAATree<K>::count(Key key) {
  return std::distance(lower_bound(key), upper_bound(key));
}
 
template <typename K>
GenericAATree::iterator KeyedAATree<K>::find(Key key) {
  iterator iter = lower_bound(key);
  if (iter == end() || compare_(key, get_key_(*iter))) {
    return end();
  }
  return iter;
}
 
template <typename K>
std::pair<GenericAATree::iterator, GenericAATree::iterator>
KeyedAATree<K>::equal_range(Key key) {
  return std::make_pair(lower_bound(key), upper_bound(key));
}
 
template <typename K>
GenericAATree::iterator KeyedAATree<K>::lower_bound(Key key) {
  AATreeItem* item = GetLowerBoundImpl(root_, key);
  return item == nullptr ? end() : iterator(&root_, item);
}
 
template <typename K>
AATreeItem* KeyedAATree<K>::GetLowerBoundImpl(AATreeItem* item, Key key) {
  if (item == nullptr) {
    return nullptr;
  }
  // If the item's key is less than the key, go right.
  if (compare_(get_key_(*item), key)) {
    return GetLowerBoundImpl(item->right_.get(), key);
  }
  // Otherwise the item's key is greater than the key. Try to go left.
  AATreeItem* left = item->left_.get();
  if (left == nullptr) {
    return item;
  }
  AATreeItem* next = GetLowerBoundImpl(left, key);
  return next == nullptr ? item : next;
}
 
template <typename K>
GenericAATree::iterator KeyedAATree<K>::upper_bound(Key key) {
  AATreeItem* item = GetUpperBoundImpl(root_, key);
  return item == nullptr ? end() : iterator(&root_, item);
}
 
template <typename K>
AATreeItem* KeyedAATree<K>::GetUpperBoundImpl(AATreeItem* item, Key key) {
  if (item == nullptr) {
    return nullptr;
  }
  // If the item's key is less than or equal to the key, go right.
  if (!compare_(key, get_key_(*item))) {
    return GetUpperBoundImpl(item->right_.get(), key);
  }
  // Otherwise the item's key is greater than the key. Try to go left.
  AATreeItem* left = item->left_.get();
  if (left == nullptr) {
    return item;
  }
  AATreeItem* next = GetUpperBoundImpl(left, key);
  return next == nullptr ? item : next;
}
 
}  // namespace pw::containers::internal