dynamic_map.h

// Copyright 2025 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 <initializer_list>
#include <iterator>
#include <optional>
#include <tuple>
#include <type_traits>
#include <utility>
 
#include "pw_allocator/allocator.h"
#include "pw_allocator/unique_ptr.h"
#include "pw_assert/assert.h"
#include "pw_containers/internal/traits.h"
#include "pw_containers/intrusive_map.h"
 
namespace pw {
 
 
template <typename Key, typename Value>
class DynamicMap {
 public:
  class Node : public IntrusiveMap<Key, Node>::Item {
   public:
    using key_type = Key;
    using mapped_type = Value;
 
    Node(const Node&) = delete;
    Node& operator=(const Node&) = delete;
 
    template <typename K, typename... Args>
    explicit Node(K&& key, Args&&... args)
        : key_value_(std::piecewise_construct,
                     std::forward_as_tuple(std::forward<K>(key)),
                     std::forward_as_tuple(std::forward<Args>(args)...)) {}
 
    const key_type& key() const { return key_value_.first; }
 
    mapped_type& mapped() { return key_value_.second; }
    const mapped_type& mapped() const { return key_value_.second; }
 
    ~Node() = default;
 
   private:
    std::pair<const key_type, mapped_type> key_value_;
    friend class DynamicMap;
  };
 
 private:
  using Map = IntrusiveMap<Key, Node>;
 
  template <bool kIsConst>
  class Iterator {
   public:
    using iterator_category = std::bidirectional_iterator_tag;
    using value_type = std::pair<const Key, Value>;
    using difference_type = std::ptrdiff_t;
    using reference =
        std::conditional_t<kIsConst, const value_type&, value_type&>;
    using pointer =
        std::conditional_t<kIsConst, const value_type*, value_type*>;
 
    constexpr Iterator() = default;
 
    constexpr reference operator*() const { return it_->key_value_; }
    constexpr pointer operator->() const { return &it_->key_value_; }
 
    constexpr Iterator operator++() {
      ++it_;
      return *this;
    }
    constexpr Iterator operator++(int) {
      Iterator result = *this;
      ++it_;
      return result;
    }
    constexpr Iterator operator--() {
      --it_;
      return *this;
    }
    constexpr Iterator operator--(int) {
      Iterator result = *this;
      --it_;
      return result;
    }
 
    constexpr friend bool operator==(Iterator lhs, Iterator rhs) {
      return lhs.it_ == rhs.it_;
    }
    constexpr friend bool operator!=(Iterator lhs, Iterator rhs) {
      return lhs.it_ != rhs.it_;
    }
 
   private:
    using BaseIterator = std::conditional_t<kIsConst,
                                            typename Map::const_iterator,
                                            typename Map::iterator>;
 
    constexpr explicit Iterator(BaseIterator it) : it_(it) {}
 
    template <bool>
    friend class Iterator;
    friend class DynamicMap;
 
    BaseIterator it_;
  };
 
 public:
  using key_type = Key;
  using mapped_type = Value;
  using value_type = std::pair<const Key, Value>;
  using size_type = typename Map::size_type;
  using difference_type = typename Map::difference_type;
  using key_compare = typename Map::key_compare;
  using allocator_type = Allocator;
  using reference = value_type&;
  using const_reference = const value_type&;
  using pointer = value_type*;
  using const_pointer = const value_type*;
  using iterator = Iterator<false>;
  using const_iterator = Iterator<true>;
  using reverse_iterator = std::reverse_iterator<iterator>;
  using const_reverse_iterator = std::reverse_iterator<const_iterator>;
  using node_type = Node;
 
  struct insert_return_type {
    iterator position;
    bool inserted;
    UniquePtr<node_type> node;
  };
 
  constexpr explicit DynamicMap(Allocator& allocator) noexcept
      : allocator_(&allocator), map_() {}
 
  DynamicMap(const DynamicMap&) = delete;
  DynamicMap& operator=(const DynamicMap&) = delete;
 
  DynamicMap(DynamicMap&& other) noexcept {
    allocator_ = other.allocator_;
    map_.swap(other.map_);
  }
 
  DynamicMap& operator=(DynamicMap&& other) noexcept {
    if (&other == this) {
      return *this;
    }
    clear();
    allocator_ = other.allocator_;
    map_.swap(other.map_);
    return *this;
  }
 
  ~DynamicMap() { clear(); }
 
  constexpr allocator_type& get_allocator() const { return *allocator_; }
 
  // Iterators
 
  iterator begin() noexcept { return iterator(map_.begin()); }
  const_iterator begin() const noexcept { return const_iterator(map_.begin()); }
  const_iterator cbegin() const noexcept { return begin(); }
  iterator end() noexcept { return iterator(map_.end()); }
  const_iterator end() const noexcept { return const_iterator(map_.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(); }
 
  // Element access
 
  mapped_type& at(const key_type& key) { return map_.at(key).mapped(); }
 
  const mapped_type& at(const key_type& key) const {
    return map_.at(key).mapped();
  }
 
  template <typename U = mapped_type,
            typename = std::enable_if_t<std::is_default_constructible_v<U>>>
  mapped_type& operator[](const key_type& key) {
    return emplace(key).first->second;
  }
 
  // Capacity
 
  [[nodiscard]] bool empty() const { return map_.empty(); }
  size_type size() const { return map_.size(); }
  constexpr size_type max_size() const noexcept { return map_.max_size(); }
 
  void clear() {
    while (!empty()) {
      erase(begin());
    }
  }
 
  // Modifiers
 
  [[nodiscard]] std::optional<std::pair<iterator, bool>> try_insert(
      const value_type& value) {
    return try_emplace(value.first, value.second);
  }
 
  // Moving into a fallible insertion is deleted to prevent "ghost moves."
  // If allocation fails, the object would be moved-from but not stored.
  // Use try_emplace instead to ensure moves only occur on success.
  std::optional<std::pair<iterator, bool>> try_insert(value_type&&) = delete;
 
  std::pair<iterator, bool> insert(const value_type& value) {
    return emplace(value.first, value.second);
  }
 
  std::pair<iterator, bool> insert(value_type&& value) {
    return emplace(std::move(value.first), std::move(value.second));
  }
 
  template <typename InputIt,
            typename = containers::internal::EnableIfInputIterator<InputIt>>
  void insert(InputIt first, InputIt last) {
    for (auto it = first; it != last; ++it) {
      emplace(it->first, it->second);
    }
  }
 
  void insert(std::initializer_list<value_type> ilist) {
    insert(ilist.begin(), ilist.end());
  }
 
  insert_return_type insert(UniquePtr<node_type>&& node) {
    if (node == nullptr) {
      return {end(), false, nullptr};
    }
    PW_ASSERT(node.deallocator() == allocator_);
    auto [node_it, inserted] = map_.insert(*node);
    if (inserted) {
      node.Release();
      return {iterator(node_it), true, nullptr};
    }
    return {iterator(node_it), false, std::move(node)};
  }
 
  // TODO: b/483691699 - Support hint-optimized insertion
 
  template <typename K, typename... Args>
  [[nodiscard]] std::optional<std::pair<iterator, bool>> try_emplace(
      K&& key, Args&&... args) {
    return TryEmplaceImpl(std::forward<K>(key), std::forward<Args>(args)...);
  }
 
  template <typename K, typename... Args>
  std::pair<iterator, bool> emplace(K&& key, Args&&... args) {
    auto result =
        try_emplace(std::forward<K>(key), std::forward<Args>(args)...);
    PW_ASSERT(result.has_value());
    return result.value();
  }
 
  iterator erase(iterator pos) {
    node_type* node = &(*pos.it_);
    auto next_it = map_.erase(pos.it_);
    allocator_->Delete(node);
    return iterator(next_it);
  }
 
  iterator erase(const_iterator pos) {
    return erase(iterator(typename Map::iterator(pos.it_)));
  }
 
  iterator erase(iterator first, iterator last) {
    while (first != last) {
      first = erase(first);
    }
    return last;
  }
 
  template <typename K>
  size_type erase(K&& key) {
    auto it = find(key);
    if (it == end()) {
      return 0;
    }
    erase(it);
    return 1;
  }
 
  UniquePtr<node_type> take(iterator pos) {
    node_type* node = &(*pos.it_);
    map_.erase(pos.it_);
    return UniquePtr<node_type>(node, *allocator_);
  }
 
  UniquePtr<node_type> take(const_iterator pos) {
    return take(iterator(typename Map::iterator(pos.it_)));
  }
 
  template <typename K>
  UniquePtr<node_type> take(K&& key) {
    auto it = find(key);
    if (it == end()) {
      return nullptr;
    }
    return take(it);
  }
 
  void swap(DynamicMap& other) {
    std::swap(allocator_, other.allocator_);
    map_.swap(other.map_);
  }
 
  void merge(DynamicMap& other) {
    if (this == &other) {
      return;
    }
    PW_ASSERT(allocator_ == other.allocator_);
    if (empty()) {
      swap(other);
      return;
    }
    for (auto it = other.begin(); it != other.end();) {
      if (contains(it->first)) {
        ++it;
      } else {
        insert(other.take(it++));
      }
    }
  }
 
  void merge(DynamicMap&& other) { merge(other); }
 
  // Lookup
 
  size_type count(const key_type& key) { return contains(key) ? 1 : 0; }
 
  iterator find(const key_type& key) { return iterator(map_.find(key)); }
 
  const_iterator find(const key_type& key) const {
    return const_iterator(map_.find(key));
  }
 
  [[nodiscard]] bool contains(const key_type& key) const {
    return find(key) != end();
  }
 
  std::pair<iterator, iterator> equal_range(const key_type& key) {
    auto result = map_.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 = map_.equal_range(key);
    return std::make_pair(const_iterator(result.first),
                          const_iterator(result.second));
  }
 
  iterator lower_bound(const key_type& key) {
    return iterator(map_.lower_bound(key));
  }
 
  const_iterator lower_bound(const key_type& key) const {
    return const_iterator(map_.lower_bound(key));
  }
 
  iterator upper_bound(const key_type& key) {
    return iterator(map_.upper_bound(key));
  }
 
  const_iterator upper_bound(const key_type& key) const {
    return const_iterator(map_.upper_bound(key));
  }
 
 private:
  template <typename K, typename... Args>
  [[nodiscard]] std::optional<std::pair<iterator, bool>> TryEmplaceImpl(
      K&& key, Args&&... args) {
    // Perform lookup before allocation to avoid unnecessary heap usage.
    if (auto it = map_.find(key); it != map_.end()) {
      return std::make_pair(iterator(it), false);
    }
    node_type* node = allocator_->template New<node_type>(
        std::forward<K>(key), std::forward<Args>(args)...);
    if (node == nullptr) {
      return std::nullopt;
    }
    auto [node_it, inserted] = map_.insert(*node);
    PW_ASSERT(inserted);
    return std::make_pair(iterator(node_it), true);
  }
 
  Allocator* allocator_;
  Map map_;
};
 
}  // namespace pw