// 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_CANONICAL_VALUE_STORE_H_
#define CARBON_TOOLCHAIN_BASE_CANONICAL_VALUE_STORE_H_

#include "common/hashtable_key_context.h"
#include "common/set.h"
#include "toolchain/base/mem_usage.h"
#include "toolchain/base/value_store.h"
#include "toolchain/base/value_store_types.h"
#include "toolchain/base/yaml.h"

namespace Carbon {

// A wrapper for accumulating immutable values with deduplication, providing IDs
// to later retrieve the value.
//
// `ValueT` represents the type being stored.
//
// `KeyT` can optionally be different from `ValueT`, and if so is used for the
// argument to `Lookup`. In this case, `ValueT` must provide a `GetAsKey` member
// function that returns the corresponding key.
template <typename IdT, typename KeyT, typename TagIdT = Untagged,
          typename ValueT = KeyT>
class CanonicalValueStore {
 public:
  using IdType = IdT;
  using IdTagType = IdTag<IdT, TagIdT>;
  using KeyType = std::remove_cvref_t<KeyT>;
  using ValueType = ValueStoreTypes<ValueT>::ValueType;
  using RefType = ValueStoreTypes<ValueT>::RefType;
  using ConstRefType = ValueStoreTypes<ValueT>::ConstRefType;

  CanonicalValueStore() = default;
  template <typename Id>
  explicit CanonicalValueStore(Id id, int32_t initial_reserved_ids = 0)
      : values_(id, initial_reserved_ids) {}

  // Stores a canonical copy of the value and returns an ID to reference it. If
  // the value is already in the store, returns the ID of the existing value.
  auto Add(ValueType value) -> IdT;

  // Returns the value for an ID.
  auto Get(IdT id) const -> ConstRefType { return values_.Get(id); }

  // Looks up the canonical ID for a value, or returns `None` if not in the
  // store.
  auto Lookup(KeyType key) const -> IdT;

  // Reserves space.
  auto Reserve(size_t size) -> void;

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

  auto values() const [[clang::lifetimebound]]
  -> ValueStore<IdT, ValueType, TagIdT>::Range {
    return values_.values();
  }
  auto size() const -> size_t { return values_.size(); }
  auto enumerate() const -> auto { return values_.enumerate(); }

  // Collects memory usage of the values and deduplication set.
  auto CollectMemUsage(MemUsage& mem_usage, llvm::StringRef label) const
      -> void {
    mem_usage.Collect(MemUsage::ConcatLabel(label, "values_"), values_);
    auto bytes = set_.ComputeMetrics(KeyContext(&values_)).storage_bytes;
    mem_usage.Add(MemUsage::ConcatLabel(label, "set_"), bytes, bytes);
  }

  auto GetRawIndex(IdT id) const -> int32_t { return values_.GetRawIndex(id); }

  auto GetIdTag() const -> IdTagType { return values_.GetIdTag(); }

 private:
  class KeyContext;

  static auto GetAsKey(ConstRefType value) -> ConstRefType
    requires std::same_as<KeyT, ValueT>
  {
    return value;
  }

  template <typename T>
  static auto GetAsKey(T&& value) -> decltype(value.GetAsKey()) {
    return value.GetAsKey();
  }

  ValueStore<IdT, ValueType, TagIdT> values_;
  Set<IdT, /*SmallSize=*/0, KeyContext> set_;
};

template <typename IdT, typename KeyT, typename TagIdT, typename ValueT>
class CanonicalValueStore<IdT, KeyT, TagIdT, ValueT>::KeyContext
    : public TranslatingKeyContext<KeyContext> {
 public:
  explicit KeyContext(const ValueStore<IdT, ValueType, TagIdT>* values)
      : values_(values) {}

  // Note that it is safe to return a reference here as the underlying object's
  // lifetime is provided by the `ValueStore`.
  auto TranslateKey(IdT id) const
      -> decltype(GetAsKey(std::declval<ConstRefType>())) {
    return GetAsKey(values_->Get(id));
  }

 private:
  const ValueStore<IdT, ValueType, TagIdT>* values_;
};

template <typename IdT, typename KeyT, typename TagIdT, typename ValueT>
auto CanonicalValueStore<IdT, KeyT, TagIdT, ValueT>::Add(ValueType value)
    -> IdT {
  auto make_key = [&] { return IdT(values_.Add(std::move(value))); };
  return set_.Insert(GetAsKey(value), make_key, KeyContext(&values_)).key();
}

template <typename IdT, typename KeyT, typename TagIdT, typename ValueT>
auto CanonicalValueStore<IdT, KeyT, TagIdT, ValueT>::Lookup(KeyType key) const
    -> IdT {
  if (auto result = set_.Lookup(key, KeyContext(&values_))) {
    return result.key();
  }
  return IdT::None;
}

template <typename IdT, typename KeyT, typename TagIdT, typename ValueT>
auto CanonicalValueStore<IdT, KeyT, TagIdT, ValueT>::Reserve(size_t size)
    -> void {
  // Compute the resulting new insert count using the size of values -- the
  // set doesn't have a fast to compute current size.
  if (size > values_.size()) {
    set_.GrowForInsertCount(size - values_.size(), KeyContext(&values_));
  }
  values_.Reserve(size);
}

}  // namespace Carbon

#endif  // CARBON_TOOLCHAIN_BASE_CANONICAL_VALUE_STORE_H_