carbon-lang/toolchain/base/value_store.h

// Part of the Carbon Language project, under the Apache License v2.0 with LLVM
// Exceptions. See /LICENSE for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

#ifndef CARBON_TOOLCHAIN_BASE_VALUE_STORE_H_
#define CARBON_TOOLCHAIN_BASE_VALUE_STORE_H_

#include <bit>
#include <compare>
#include <cstddef>
#include <limits>
#include <memory>
#include <type_traits>
#include <utility>

#include "common/check.h"
#include "common/ostream.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/Sequence.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/iterator.h"
#include "llvm/Support/MemAlloc.h"
#include "toolchain/base/id_tag.h"
#include "toolchain/base/mem_usage.h"
#include "toolchain/base/value_store_types.h"
#include "toolchain/base/yaml.h"

namespace Carbon {

template <typename IdT, typename ValueT, typename TagIdT>
class ValueStore;

namespace Internal {

// Used as a parent class for non-printable types. This is just for
// std::conditional, not as an API.
class ValueStoreNotPrintable {};

// A random-access iterator over the values in a ValueStore.
//
// Supports both const and non-const iteration, and allows implicit
// conversion from non-const to const iterators.
template <typename ValueStoreT, typename ElementT, typename PointerT,
          typename ReferenceT>
class ValueStoreIterator
    : public llvm::iterator_facade_base<
          ValueStoreIterator<ValueStoreT, ElementT, PointerT, ReferenceT>,
          std::random_access_iterator_tag, ElementT, int32_t, PointerT,
          ReferenceT> {
 public:
  // Implicit conversion constructor from non-const to const iterator.
  // Only enabled if the source iterator's ValueStore pointer is convertible
  // to the target iterator's ValueStore pointer (i.e., ValueStore* to const
  // ValueStore*).
  template <typename OtherValueStoreT, typename OtherElementT,
            typename OtherPointerT, typename OtherReferenceT>
    requires(std::is_convertible_v<OtherValueStoreT*, ValueStoreT*>)
  // NOLINTNEXTLINE(google-explicit-constructor)
  ValueStoreIterator(
      const ValueStoreIterator<OtherValueStoreT, OtherElementT, OtherPointerT,
                               OtherReferenceT>& other)
      : store_(other.store_), index_(other.index_) {}

  auto operator*() const -> ReferenceT {
    return store_->Get(
        typename ValueStoreT::IdType(store_->tag_.Apply(index_)));
  }

  friend auto operator==(const ValueStoreIterator& lhs,
                         const ValueStoreIterator& rhs) -> bool {
    CARBON_DCHECK(lhs.store_ == rhs.store_);
    return lhs.index_ == rhs.index_;
  }

  friend auto operator<=>(const ValueStoreIterator& lhs,
                          const ValueStoreIterator& rhs)
      -> std::strong_ordering {
    CARBON_DCHECK(lhs.store_ == rhs.store_);
    return lhs.index_ <=> rhs.index_;
  }

  friend auto operator-(const ValueStoreIterator& lhs,
                        const ValueStoreIterator& rhs) -> int32_t {
    return lhs.index_ - rhs.index_;
  }

  auto operator+=(int32_t n) -> ValueStoreIterator& {
    index_ += n;
    return *this;
  }

  auto operator-=(int32_t n) -> ValueStoreIterator& { return *this += -n; }

 private:
  template <typename, typename, typename, typename>
  friend class ValueStoreIterator;

  template <typename, typename, typename>
  friend class Carbon::ValueStore;

  template <typename, typename>
  friend class ValueStoreRange;

  // Private constructor to prevent arbitrary construction.
  // Iterators should only be obtained from Range begin()/end().
  ValueStoreIterator(ValueStoreT* store, int32_t index)
      : store_(store), index_(index) {}

  ValueStoreT* store_;
  int32_t index_;
};

// A random-access iterator for enumerating a ValueStore, yielding pairs of
// (ID, Value).
template <typename ValueStoreT>
class ValueStoreEnumerateIterator
    : public llvm::iterator_facade_base<
          ValueStoreEnumerateIterator<ValueStoreT>,
          std::random_access_iterator_tag,
          std::pair<typename ValueStoreT::IdType,
                    typename ValueStoreT::ConstRefType>,
          int32_t,
          const std::pair<typename ValueStoreT::IdType,
                          typename ValueStoreT::ConstRefType>*,
          std::pair<typename ValueStoreT::IdType,
                    typename ValueStoreT::ConstRefType>> {
 public:
  using IdType = ValueStoreT::IdType;
  using ConstRefType = ValueStoreT::ConstRefType;
  using ValueType = std::pair<IdType, ConstRefType>;

  auto operator*() const -> ValueType {
    IdType id = store_->tag_.Apply(index_);
    return {id, store_->Get(id)};
  }

  friend auto operator==(const ValueStoreEnumerateIterator& lhs,
                         const ValueStoreEnumerateIterator& rhs) -> bool {
    CARBON_DCHECK(lhs.store_ == rhs.store_);
    return lhs.index_ == rhs.index_;
  }

  friend auto operator<=>(const ValueStoreEnumerateIterator& lhs,
                          const ValueStoreEnumerateIterator& rhs)
      -> std::strong_ordering {
    CARBON_DCHECK(lhs.store_ == rhs.store_);
    return lhs.index_ <=> rhs.index_;
  }

  friend auto operator-(const ValueStoreEnumerateIterator& lhs,
                        const ValueStoreEnumerateIterator& rhs) -> int32_t {
    return lhs.index_ - rhs.index_;
  }

  auto operator+=(int32_t n) -> ValueStoreEnumerateIterator& {
    index_ += n;
    return *this;
  }

  auto operator-=(int32_t n) -> ValueStoreEnumerateIterator& {
    return *this += -n;
  }

 private:
  template <typename, typename, typename>
  friend class Carbon::ValueStore;

  template <typename, typename>
  friend class ValueStoreRange;

  // Private constructor. Enumerate iterators should only be obtained from
  // EnumerateRange.
  ValueStoreEnumerateIterator(ValueStoreT* store, int32_t index)
      : store_(store), index_(index) {}

  ValueStoreT* store_;
  int32_t index_;
};

// A range over a ValueStore, vending begin() and end() iterators.
template <typename ValueStoreT, typename IteratorT>
class ValueStoreRange {
 public:
  auto begin() const -> IteratorT { return IteratorT(store_, 0); }
  auto end() const -> IteratorT { return IteratorT(store_, store_->size_); }

 private:
  template <typename, typename, typename>
  friend class Carbon::ValueStore;

  // Private constructor to ensure ranges are only created by ValueStore.
  explicit ValueStoreRange(ValueStoreT& store [[clang::lifetimebound]])
      : store_(&store) {}

  ValueStoreT* store_;
};

}  // namespace Internal

// A simple wrapper for accumulating values, providing IDs to later retrieve the
// value. This does not do deduplication.
//
// This template class is designed to be explicitly instantiated to improve
// compile times. Heavy method definitions are in `value_store_impl.h`.
//
// To use a new instantiation:
// 1. Add an `extern template class ValueStore<...>;` declaration in the header
//    where the instantiation is anchored.
// 2. Add `template class ValueStore<...>;` definition in the associated `.cpp`
//    file.
// 3. Include `toolchain/base/value_store_impl.h` in that `.cpp` file.
template <typename IdT, typename ValueT, typename TagIdT = Untagged>
class ValueStore
    : public std::conditional<std::is_base_of_v<Printable<ValueT>, ValueT>,
                              Yaml::Printable<ValueStore<IdT, ValueT, TagIdT>>,
                              Internal::ValueStoreNotPrintable> {
 public:
  using IdType = IdT;
  using IdTagType = IdTag<IdT, TagIdT>;
  using ValueType = ValueStoreTypes<ValueT>::ValueType;
  using RefType = ValueStoreTypes<ValueT>::RefType;
  using ConstRefType = ValueStoreTypes<ValueT>::ConstRefType;

  using ConstIterator =
      Internal::ValueStoreIterator<const ValueStore, const ValueT,
                                   const ValueT*, ConstRefType>;
  using Iterator = Internal::ValueStoreIterator<
      ValueStore, ValueT,
      std::conditional_t<std::is_reference_v<RefType>, ValueT*, const ValueT*>,
      RefType>;

  using Range = Internal::ValueStoreRange<const ValueStore, ConstIterator>;
  using MutableRange = Internal::ValueStoreRange<ValueStore, Iterator>;

  using EnumerateIterator =
      Internal::ValueStoreEnumerateIterator<const ValueStore>;
  using EnumerateRange =
      Internal::ValueStoreRange<const ValueStore, EnumerateIterator>;

  ValueStore()
    requires(IdTagIsUntagged<IdTagType>);

  explicit ValueStore(IdTagType tag)
    requires(!IdTagIsUntagged<IdTagType>);

  explicit ValueStore(IdTagType::TagIdType id, int32_t initial_reserved_ids = 0)
    requires(!IdTagIsUntagged<IdTagType>);

  ~ValueStore();
  ValueStore(ValueStore&&) noexcept;
  auto operator=(ValueStore&&) noexcept -> ValueStore&;

  // Stores the value and returns an ID to reference it.
  auto Add(ValueType value) -> IdType {
    // This routine is especially hot and the check here relatively expensive
    // for the value provided, so only do this in non-optimized builds to make
    // tracking down issues easier.
    CARBON_DCHECK(size_ < std::numeric_limits<int32_t>::max(), "Id overflow");

    IdType id = tag_.Apply(size_);
    auto [chunk_index, pos] = RawIndexToChunkIndices(size_);
    ++size_;

    CARBON_DCHECK(static_cast<size_t>(chunk_index) <= chunks_.size(),
                  "{0} <= {1}", chunk_index, chunks_.size());
    if (static_cast<size_t>(chunk_index) == chunks_.size()) {
      chunks_.emplace_back();
    }

    CARBON_DCHECK(pos == chunks_[chunk_index].size());
    chunks_[chunk_index].Add(std::move(value));
    return id;
  }

  // Returns a mutable value for an ID.
  auto Get(IdType id) -> RefType {
    CARBON_DCHECK(id.index >= 0, "{0}", id);
    auto [chunk_index, pos] = IdToChunkIndices(id);
    return chunks_[chunk_index].Get(pos);
  }

  // Returns the value for an ID.
  auto Get(IdType id) const -> ConstRefType {
    CARBON_DCHECK(id.index >= 0, "{0}", id);
    auto [chunk_index, pos] = IdToChunkIndices(id);
    return chunks_[chunk_index].Get(pos);
  }

  // Returns the value for an ID, or a specified default value if a value has
  // not yet been added for this ID.
  auto GetWithDefault(IdType id,  //
                      ConstRefType default_value [[clang::lifetimebound]]) const
      -> ConstRefType {
    CARBON_DCHECK(id.index >= 0, "{0}", id);
    auto index = tag_.Remove(id);
    if (index >= size_) {
      return default_value;
    }
    auto [chunk_index, pos] = RawIndexToChunkIndices(index);
    return chunks_[chunk_index].Get(pos);
  }

  // Reserves space. Note that this method has surprising performance impact on
  // the lexer unless available for inlining.
  auto Reserve(int32_t size) -> void {
    if (size <= size_) {
      return;
    }
    auto [final_chunk_index, _] = RawIndexToChunkIndices(size - 1);
    chunks_.resize(final_chunk_index + 1);
  }

  // Grows the ValueStore to `size`. Fills entries with `default_value`.
  auto Resize(int32_t size, ConstRefType default_value) -> void
    requires(std::is_copy_constructible_v<ValueT>);

  // These are to support printable structures, and are not guaranteed.
  auto OutputYaml() const -> Yaml::OutputMapping;

  // Collects memory usage of the values.
  auto CollectMemUsage(MemUsage& mem_usage, llvm::StringRef label) const
      -> void;

  auto size() const -> size_t { return size_; }

  // Makes an iterable range over the IDs in the ValueStore.
  auto ids() const -> auto {
    return llvm::map_range(llvm::seq(size_), [tag = tag_](int32_t i) -> IdType {
      return tag.Apply(i);
    });
  }

  // Makes an iterable range over references to all values in the ValueStore.
  auto values() [[clang::lifetimebound]] -> MutableRange {
    return MutableRange(*this);
  }
  auto values() const [[clang::lifetimebound]] -> Range { return Range(*this); }

  // Makes an iterable range over pairs of the index and a reference to the
  // value for each value in the store.
  //
  // The range is over references to the values in the store, even if used with
  // `auto` to destructure the pair. In this example, the `value` is a
  // `ConstRefType`:
  // ```
  // for (auto [id, value] : store.enumerate()) { ... }
  // ```
  auto enumerate() const [[clang::lifetimebound]] -> EnumerateRange {
    return EnumerateRange(*this);
  }

  auto GetIdTag() const -> IdTagType { return tag_; }
  auto GetRawIndex(IdT id) const -> int32_t {
    CARBON_DCHECK(id.index >= 0, "{0}", id.index);
    auto index = tag_.Remove(id);
#ifndef NDEBUG
    if (index >= size_) {
      // Attempt to decompose id.index to include extra detail in the check
      // here.
      //
      // TODO: Teach ValueStore the type of the tag id with a template, then we
      // can print it with proper formatting instead of just as an integer.
      auto [id_tag, id_untagged_index] = IdTagType::DecomposeWithBestEffort(id);
      CARBON_DCHECK(
          index < size_,
          "Untagged index was outside of container range. Tagged index {0}. "
          "Best-effort decomposition: Tag: {1}, Index: {2}. "
          "Container size: {3}. "
          "Expected Tag for this container: {4}.",
          id.index, id_tag, id_untagged_index, size_, tag_.GetContainerTag());
    }
#endif
    return index;
  }

 private:
  template <typename, typename, typename, typename>
  friend class Internal::ValueStoreIterator;

  template <typename, typename>
  friend class Internal::ValueStoreRange;

  template <typename>
  friend class Internal::ValueStoreEnumerateIterator;
  // A chunk of `ValueType`s which has a fixed capacity, but variable size.
  // Tracks the size internally for verifying bounds.
  struct Chunk {
   public:
    // The max size of each chunk allocation for `ValueStore`. This is based on
    // TLB page sizes for the target platform.
    //
    // See https://docs.kernel.org/admin-guide/mm/hugetlbpage.html
    //
    // A 4K chunk size outperforms a 1M chunk size on Linux-x64 and MacOS-arm64
    // in benchmarks and when running file_test.
    //
    // Linux-x64: x64 CPUs support 4K and 2M page sizes, but we see 1M is slower
    // than 4K with tcmalloc in opt builds for our tests.
    //
    // Mac-arm64: arm64 CPUs support 4K, 8K, 64K, 256K, 1M, 4M and up. Like for
    // Linux-x64, 4K outperformed 1M. We didn't try other sizes yet.
    //
    // TODO: Is there a more optimize size for Mac-arm64? What should
    // Linux-arm64 and Mac-x64 use? What should Windows use?
    //
    // TODO: The previous SmallVector<ValueType> seems to outperform 4K chunks
    // (they may be slower by up to 5%) in benchmarks. Find ways to make
    // chunking faster. Should successive chunks get larger in size? That will
    // greatly complicate math for choosing a chunk though.
    static constexpr auto MaxAllocationBytes() -> int32_t {
#if !defined(NDEBUG) || LLVM_ADDRESS_SANITIZER_BUILD
      // Use a small size in unoptimized builds to ensure multiple chunks get
      // used. And do the same in ASAN builds to reduce bookkeeping overheads.
      // Using large allocations (e.g. 1M+) incurs a 10x runtime cost for our
      // tests under ASAN.
      return sizeof(ValueType) * 5;
#else
      return 4 * 1024;
#endif
    }

    // The number of elements stored in each chunk allocation.
    //
    // The number must be a power of two so that that there are no unused values
    // in bits indexing into the allocation.
    static constexpr auto Capacity() -> int32_t {
      constexpr auto MaxElements = MaxAllocationBytes() / sizeof(ValueType);
      return std::bit_floor(MaxElements);
    }

    // The number of bits needed to index each element in a chunk allocation.
    static constexpr auto IndexBits() -> int32_t {
      static_assert(Capacity() > 0);
      return std::bit_width(uint32_t{Capacity() - 1});
    }

    static constexpr auto CapacityBytes = Capacity() * sizeof(ValueType);

    explicit Chunk()
        : buf_(reinterpret_cast<ValueType*>(
              llvm::allocate_buffer(CapacityBytes, alignof(ValueType)))) {}

    // Moving leaves nullptr behind in the moved-from object so that the
    // destructor is a no-op (preventing double free).
    Chunk(Chunk&& rhs) noexcept
        : buf_(std::exchange(rhs.buf_, nullptr)), num_(rhs.num_) {}

    auto operator=(Chunk&& rhs) noexcept -> Chunk& {
      buf_ = std::exchange(rhs.buf_, nullptr);
      num_ = rhs.num_;
      return *this;
    }

    ~Chunk() {
      if (buf_) {
        if constexpr (!std::is_trivially_destructible_v<ValueType>) {
          std::destroy_n(buf_, num_);
        }
        llvm::deallocate_buffer(buf_, CapacityBytes, alignof(ValueType));
      }
    }

    auto Get(int32_t i) -> ValueType& {
      CARBON_DCHECK(i < num_, "{0}", i);
      return buf_[i];
    }
    auto Get(int32_t i) const -> const ValueType& {
      CARBON_DCHECK(i < num_, "{0}", i);
      return buf_[i];
    }

    auto Add(ValueType&& value) -> void {
      CARBON_DCHECK(num_ < Capacity());
      std::construct_at(buf_ + num_, std::move(value));
      ++num_;
    }

    // Fills `fill_count` entries with `default_value`, increasing the size
    // respectively.
    auto UninitializedFill(int32_t fill_count, ConstRefType default_value)
        -> void
      requires(std::is_copy_constructible_v<ValueT>);

    auto size() const -> int32_t { return num_; }

   private:
    // Verify using an `int32_t` for `num_` is sound.
    static_assert(Capacity() <= std::numeric_limits<int32_t>::max());

    ValueType* buf_;
    int32_t num_ = 0;
  };

  // Converts a raw index into an index into the set of chunks, and an offset
  // into that specific chunk. Looks for index overflow in non-optimized builds.
  static auto RawIndexToChunkIndices(int32_t index)
      -> std::pair<int32_t, int32_t> {
    constexpr auto LowBits = Chunk::IndexBits();

    // Verify there are no unused bits when indexing up to the `Capacity`. This
    // ensures that ids are contiguous values from 0, as if the values were all
    // stored in a single array, and allows using the ids to index into other
    // arrays.
    static_assert((1 << LowBits) == Chunk::Capacity());
    // Simple check to make sure nothing went wildly wrong with the `Capacity`,
    // and we have some room for a chunk index, and that shifting by the number
    // of bits won't be UB in an int32_t.
    static_assert(LowBits < 30);

    // The index of the chunk is the high bits.
    auto chunk = index >> LowBits;
    // The index into the chunk is the low bits.
    auto pos = index & ((1 << LowBits) - 1);
    return {chunk, pos};
  }

  // Converts an id into an index into the set of chunks, and an offset into
  // that specific chunk.
  auto IdToChunkIndices(IdType id) const -> std::pair<int32_t, int32_t> {
    return RawIndexToChunkIndices(GetRawIndex(id));
  }

  // Number of elements added to the store. The number should never exceed what
  // fits in an `int32_t`, which is checked in non-optimized builds in Add().
  int32_t size_ = 0;

  IdTagType tag_;

  // Storage for the `ValueType` objects, indexed by the id. We use a vector of
  // chunks of `ValueType` instead of just a vector of `ValueType` so that
  // addresses of `ValueType` objects are stable. This allows the rest of the
  // toolchain to hold references into `ValueStore` without having to worry
  // about invalidation and use-after-free. We ensure at least one Chunk is held
  // inline so that in the case where there is only a single Chunk (i.e. small
  // files) we can avoid one indirection.
  llvm::SmallVector<Chunk, 1> chunks_;
};

}  // namespace Carbon

#endif  // CARBON_TOOLCHAIN_BASE_VALUE_STORE_H_