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, typename>
  friend class AATree;
 
  AATreeItem* root_ = nullptr;
 
  const bool unique_keys_;
};
 
template <typename K, typename V>
class AATree final : public GenericAATree {
 public:
  using Key = K;
  using Compare = Function<bool(Key, Key)>;
  using GetKey = Function<Key(const V&)>;
  using GenericAATree::iterator;
 
  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)
      : GenericAATree(unique_keys),
        compare_(std::move(compare)),
        get_key_(std::move(get_key)) {}
 
  // Modifiers
 
  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(AATree<K, V>& 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);
 
  constexpr Key DoGetKey(const AATreeItem& item) {
    return get_key_(static_cast<const V&>(item));
  }
 
  Compare compare_;
  GetKey get_key_;
};
 
// Template method implementations.
 
template <typename K, typename V>
std::pair<GenericAATree::iterator, bool> AATree<K, V>::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, typename V>
AATreeItem* AATree<K, V>::InsertImpl(AATreeItem& parent,
                                     AATreeItem& child,
                                     AATreeItem*& duplicate) {
  if (compare_(DoGetKey(child), DoGetKey(parent))) {
    AATreeItem* left = parent.left_.get();
    left = left == nullptr ? &child : InsertImpl(*left, child, duplicate);
    parent.SetLeft(left);
 
  } else if (compare_(DoGetKey(parent), DoGetKey(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, typename V>
template <class Iterator>
void AATree<K, V>::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, typename V>
size_t AATree<K, V>::erase_all(Key key) {
  size_t removed = 0;
  iterator iter = lower_bound(key);
  while (iter != end() && !compare_(key, DoGetKey(*iter))) {
    iter = erase_one(*iter);
    ++removed;
  }
  return removed;
}
 
template <typename K, typename V>
void AATree<K, V>::merge(AATree<K, V>& other) {
  auto iter = other.begin();
  while (!other.empty()) {
    AATreeItem& item = *iter;
    iter = other.erase_one(item);
    insert(item);
  }
}
 
template <typename K, typename V>
size_t AATree<K, V>::count(Key key) {
  return std::distance(lower_bound(key), upper_bound(key));
}
 
template <typename K, typename V>
GenericAATree::iterator AATree<K, V>::find(Key key) {
  iterator iter = lower_bound(key);
  if (iter == end() || compare_(key, DoGetKey(*iter))) {
    return end();
  }
  return iter;
}
 
template <typename K, typename V>
std::pair<GenericAATree::iterator, GenericAATree::iterator>
AATree<K, V>::equal_range(Key key) {
  return std::make_pair(lower_bound(key), upper_bound(key));
}
 
template <typename K, typename V>
GenericAATree::iterator AATree<K, V>::lower_bound(Key key) {
  AATreeItem* item = GetLowerBoundImpl(root_, key);
  return item == nullptr ? end() : iterator(&root_, item);
}
 
template <typename K, typename V>
AATreeItem* AATree<K, V>::GetLowerBoundImpl(AATreeItem* item, Key key) {
  if (item == nullptr) {
    return nullptr;
  }
  // If the item's key is less than the key, go right.
  if (compare_(DoGetKey(*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, typename V>
GenericAATree::iterator AATree<K, V>::upper_bound(Key key) {
  AATreeItem* item = GetUpperBoundImpl(root_, key);
  return item == nullptr ? end() : iterator(&root_, item);
}
 
template <typename K, typename V>
AATreeItem* AATree<K, V>::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, DoGetKey(*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