intrusive_map.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 "pw_containers/internal/aa_tree.h"
#include "pw_containers/internal/aa_tree_item.h"
 
namespace pw {
 
template <typename Key, typename T>
class IntrusiveMap {
 private:
  using GenericIterator = containers::internal::GenericAATree::iterator;
  using Tree = containers::internal::AATree<Key, T>;
  using Compare = typename Tree::Compare;
  using GetKey = typename Tree::GetKey;
 
 public:
  using Item = typename Tree::Item;
  using Pair = typename Tree::Pair;
 
  using key_type = Key;
  using mapped_type = std::remove_cv_t<T>;
  using value_type = Item;
  using size_type = std::size_t;
  using difference_type = std::ptrdiff_t;
  using key_compare = Compare;
  using reference = value_type&;
  using const_reference = const value_type&;
  using pointer = value_type*;
  using const_pointer = const value_type*;
 
 public:
  class iterator : public containers::internal::AATreeIterator<T> {
   public:
    constexpr iterator() = default;
 
   private:
    friend IntrusiveMap;
    constexpr explicit iterator(GenericIterator iter)
        : containers::internal::AATreeIterator<T>(iter) {}
  };
 
  class const_iterator
      : public containers::internal::AATreeIterator<std::add_const_t<T>> {
   public:
    constexpr const_iterator() = default;
 
   private:
    friend IntrusiveMap;
    constexpr explicit const_iterator(GenericIterator iter)
        : containers::internal::AATreeIterator<std::add_const_t<T>>(iter) {}
  };
 
  using reverse_iterator = std::reverse_iterator<iterator>;
  using const_reverse_iterator = std::reverse_iterator<const_iterator>;
 
  constexpr explicit IntrusiveMap() : IntrusiveMap(std::less<Key>()) {}
 
  template <typename Comparator>
  constexpr explicit IntrusiveMap(Comparator compare)
      : IntrusiveMap(std::move(compare), [](const T& t) { return t.key(); }) {}
 
  template <typename Comparator, typename KeyRetriever>
  constexpr IntrusiveMap(Comparator&& compare, KeyRetriever&& get_key)
      : tree_(true,
              std::forward<Comparator>(compare),
              std::forward<KeyRetriever>(get_key)) {
    CheckItemType();
  }
 
  template <typename Iterator, typename... Functors>
  IntrusiveMap(Iterator first, Iterator last, Functors&&... functors)
      : IntrusiveMap(std::forward<Functors>(functors)...) {
    tree_.insert(first, last);
  }
 
  template <typename... Functors>
  explicit IntrusiveMap(std::initializer_list<T*> items, Functors&&... functors)
      : IntrusiveMap(
            items.begin(), items.end(), std::forward<Functors>(functors)...) {}
 
  // Element access
 
  T& at(const key_type& key) {
    auto iter = tree_.find(key);
    PW_DASSERT(iter != tree_.end());
    return static_cast<T&>(*iter);
  }
 
  const T& at(const key_type& key) const {
    auto iter = tree_.find(key);
    PW_DASSERT(iter != tree_.end());
    return static_cast<const T&>(*iter);
  }
 
  // Iterators
 
  iterator begin() noexcept { return iterator(tree_.begin()); }
  const_iterator begin() const noexcept {
    return const_iterator(tree_.begin());
  }
  const_iterator cbegin() const noexcept { return begin(); }
 
  iterator end() noexcept { return iterator(tree_.end()); }
  const_iterator end() const noexcept { return const_iterator(tree_.end()); }
  const_iterator cend() const noexcept { return end(); }
 
  reverse_iterator rbegin() noexcept { return reverse_iterator(end()); }
  const_reverse_iterator rbegin() const noexcept {
    return const_reverse_iterator(end());
  }
  const_reverse_iterator crbegin() const noexcept { return rbegin(); }
 
  reverse_iterator rend() noexcept { return reverse_iterator(begin()); }
  const_reverse_iterator rend() const noexcept {
    return const_reverse_iterator(begin());
  }
  const_reverse_iterator crend() const noexcept { return rend(); }
 
  // Capacity
 
  [[nodiscard]] bool empty() const noexcept { return tree_.empty(); }
 
  size_t size() const { return tree_.size(); }
 
  constexpr size_t max_size() const noexcept { return tree_.max_size(); }
 
  // Modifiers
 
  void clear() { tree_.clear(); }
 
  std::pair<iterator, bool> insert(T& item) {
    auto result = tree_.insert(item);
    return std::make_pair(iterator(result.first), result.second);
  }
 
  iterator insert(iterator, T& item) {
    // Disregard the hint.
    return insert(item).first;
  }
 
  template <class Iterator>
  void insert(Iterator first, Iterator last) {
    tree_.insert(first, last);
  }
 
  void insert(std::initializer_list<T*> ilist) {
    tree_.insert(ilist.begin(), ilist.end());
  }
 
  iterator erase(T& item) { return iterator(tree_.erase_one(item)); }
 
  iterator erase(iterator pos) { return iterator(tree_.erase_one(*pos)); }
 
  iterator erase(iterator first, iterator last) {
    return iterator(tree_.erase_range(*first, *last));
  }
 
  size_t erase(const key_type& key) { return tree_.erase_all(key); }
 
  void swap(IntrusiveMap<Key, T>& other) { tree_.swap(other.tree_); }
 
  template <typename MapType>
  void merge(MapType& other) {
    tree_.merge(other.tree_);
  }
 
  size_t count(const key_type& key) const { return tree_.count(key); }
 
  iterator find(const key_type& key) { return iterator(tree_.find(key)); }
 
  const_iterator find(const key_type& key) const {
    return const_iterator(tree_.find(key));
  }
 
  std::pair<iterator, iterator> equal_range(const key_type& key) {
    auto result = tree_.equal_range(key);
    return std::make_pair(iterator(result.first), iterator(result.second));
  }
  std::pair<const_iterator, const_iterator> equal_range(
      const key_type& key) const {
    auto result = tree_.equal_range(key);
    return std::make_pair(const_iterator(result.first),
                          const_iterator(result.second));
  }
 
  iterator lower_bound(const key_type& key) {
    return iterator(tree_.lower_bound(key));
  }
  const_iterator lower_bound(const key_type& key) const {
    return const_iterator(tree_.lower_bound(key));
  }
 
  iterator upper_bound(const key_type& key) {
    return iterator(tree_.upper_bound(key));
  }
  const_iterator upper_bound(const key_type& key) const {
    return const_iterator(tree_.upper_bound(key));
  }
 
 private:
  // Check that T is an Item in a function, since the class T will not be fully
  // defined when the IntrusiveList<T> class is instantiated.
  static constexpr void CheckItemType() {
    using ItemBase = containers::internal::AATreeItem;
    using IntrusiveItemType =
        typename containers::internal::IntrusiveItem<ItemBase, T>::Type;
    static_assert(
        std::is_base_of<IntrusiveItemType, T>(),
        "IntrusiveMap items must be derived from IntrusiveMap<Key, T>::Item, "
        "where T is the item or one of its bases.");
  }
 
  // Allow multimaps to access the tree for `merge`.
  template <typename, typename>
  friend class IntrusiveMultiMap;
 
  // The AA tree that stores the map.
  //
  // This field is mutable so that it doesn't need const overloads.
  mutable Tree tree_;
};
 
}  // namespace pw