carbon-lang/common/set.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_COMMON_SET_H_
#define CARBON_COMMON_SET_H_

#include <concepts>

#include "common/check.h"
#include "common/hashtable_key_context.h"
#include "common/raw_hashtable.h"
#include "llvm/Support/Compiler.h"

namespace Carbon {

// Forward declarations to resolve cyclic references.
template <typename KeyT, typename KeyContextT>
class SetView;
template <typename KeyT, typename KeyContextT>
class SetBase;
template <typename KeyT, ssize_t SmallSize, typename KeyContextT>
class Set;

// A read-only view type for a set of keys.
//
// This view is a cheap-to-copy type that should be passed by value, but
// provides view or read-only reference semantics to the underlying set data
// structure.
//
// This should always be preferred to a `const`-ref parameter for the `SetBase`
// or `Set` type as it provides more flexibility and a cleaner API.
//
// Note that while this type is a read-only view, that applies to the underlying
// *set* data structure, not the individual entries stored within it. Those can
// be mutated freely as long as both the hashes and equality of the keys are
// preserved. If we applied a deep-`const` design here, it would prevent using
// this type in situations where the keys carry state (unhashed and not part of
// equality) that is mutated while the associative container is not. A view of
// immutable data can always be obtained by using `SetView<const T>`, and we
// enable conversions to more-const views. This mirrors the semantics of views
// like `std::span`.
//
// A specific `KeyContextT` type can optionally be provided to configure how
// keys will be hashed and compared. The default is `DefaultKeyContext` which is
// stateless and will hash using `Carbon::HashValue` and compare using
// `operator==`. Every method accepting a lookup key or operating on the keys in
// the table will also accept an instance of this type. For stateless context
// types, including the default, an instance will be default constructed if not
// provided to these methods. However, stateful contexts should be constructed
// and passed in explicitly. The context type should be small and reasonable to
// pass by value, often a wrapper or pointer to the relevant context needed for
// hashing and comparing keys. For more details about the key context, see
// `hashtable_key_context.h`.
template <typename InputKeyT, typename InputKeyContextT = DefaultKeyContext>
class SetView : RawHashtable::ViewImpl<InputKeyT, void, InputKeyContextT> {
  using ImplT = RawHashtable::ViewImpl<InputKeyT, void, InputKeyContextT>;

 public:
  using KeyT = typename ImplT::KeyT;
  using KeyContextT = typename ImplT::KeyContextT;

  // This type represents the result of lookup operations. It encodes whether
  // the lookup was a success as well as accessors for the key.
  class LookupResult {
   public:
    LookupResult() = default;
    explicit LookupResult(KeyT& key) : key_(&key) {}

    explicit operator bool() const { return key_ != nullptr; }

    auto key() const -> KeyT& { return *key_; }

   private:
    KeyT* key_ = nullptr;
  };

  // Enable implicit conversions that add `const`-ness to the key type.
  // NOLINTNEXTLINE(google-explicit-constructor)
  SetView(SetView<std::remove_const_t<KeyT>, KeyContextT> other_view)
    requires(!std::same_as<KeyT, std::remove_const_t<KeyT>>)
      : ImplT(other_view) {}

  // Tests whether a key is present in the set.
  template <typename LookupKeyT>
  auto Contains(LookupKeyT lookup_key,
                KeyContextT key_context = KeyContextT()) const -> bool;

  // Lookup a key in the set.
  template <typename LookupKeyT>
  auto Lookup(LookupKeyT lookup_key,
              KeyContextT key_context = KeyContextT()) const -> LookupResult;

  // Run the provided callback for every key in the set.
  template <typename CallbackT>
  void ForEach(CallbackT callback)
    requires(std::invocable<CallbackT, KeyT&>);

  // This routine is relatively inefficient and only intended for use in
  // benchmarking or logging of performance anomalies. The specific count
  // returned has no specific guarantees beyond being informative in benchmarks.
  // It counts how many of the keys in the hashtable have required probing
  // beyond their initial group of slots.
  //
  // TODO: Replace with a more general metrics routine that covers other
  // important aspects such as load factor, and average probe *distance*.
  auto CountProbedKeys(KeyContextT key_context = KeyContextT()) -> ssize_t {
    return ImplT::CountProbedKeys(key_context);
  }

 private:
  template <typename SetKeyT, ssize_t SmallSize, typename KeyContextT>
  friend class Set;
  friend class SetBase<KeyT, KeyContextT>;
  friend class SetView<const KeyT, KeyContextT>;

  using EntryT = typename ImplT::EntryT;

  SetView() = default;
  // NOLINTNEXTLINE(google-explicit-constructor): Implicit by design.
  SetView(ImplT base) : ImplT(base) {}
  SetView(ssize_t size, RawHashtable::Storage* storage)
      : ImplT(size, storage) {}
};

// A base class for a `Set` type that remains mutable while type-erasing the
// `SmallSize` (SSO) template parameter.
//
// A pointer or reference to this type is the preferred way to pass a mutable
// handle to a `Set` type across API boundaries as it avoids encoding specific
// SSO sizing information while providing a near-complete mutable API.
template <typename InputKeyT, typename InputKeyContextT>
class SetBase
    : protected RawHashtable::BaseImpl<InputKeyT, void, InputKeyContextT> {
 protected:
  using ImplT = RawHashtable::BaseImpl<InputKeyT, void, InputKeyContextT>;

 public:
  using KeyT = typename ImplT::KeyT;
  using KeyContextT = typename ImplT::KeyContextT;
  using ViewT = SetView<KeyT, KeyContextT>;
  using LookupResult = typename ViewT::LookupResult;

  // The result type for insertion operations both indicates whether an insert
  // was needed (as opposed to the key already being in the set), and provides
  // access to the key.
  class InsertResult {
   public:
    InsertResult() = default;
    explicit InsertResult(bool inserted, KeyT& key)
        : key_(&key), inserted_(inserted) {}

    auto is_inserted() const -> bool { return inserted_; }

    auto key() const -> KeyT& { return *key_; }

   private:
    KeyT* key_;
    bool inserted_;
  };

  // Implicitly convertible to the relevant view type.
  //
  // NOLINTNEXTLINE(google-explicit-constructor): Designed to implicitly decay.
  operator ViewT() const { return this->view_impl(); }

  // We can't chain the above conversion with the conversions on `ViewT` to add
  // const, so explicitly support adding const to produce a view here.
  //
  // NOLINTNEXTLINE(google-explicit-constructor): Designed to implicitly decay.
  operator SetView<const KeyT, KeyContextT>() const { return ViewT(*this); }

  // Convenience forwarder to the view type.
  template <typename LookupKeyT>
  auto Contains(LookupKeyT lookup_key,
                KeyContextT key_context = KeyContextT()) const -> bool {
    return ViewT(*this).Contains(lookup_key, key_context);
  }

  // Convenience forwarder to the view type.
  template <typename LookupKeyT>
  auto Lookup(LookupKeyT lookup_key,
              KeyContextT key_context = KeyContextT()) const -> LookupResult {
    return ViewT(*this).Lookup(lookup_key, key_context);
  }

  // Convenience forwarder to the view type.
  template <typename CallbackT>
  void ForEach(CallbackT callback)
    requires(std::invocable<CallbackT, KeyT&>)
  {
    return ViewT(*this).ForEach(callback);
  }

  // Convenience forwarder to the view type.
  auto CountProbedKeys(KeyContextT key_context = KeyContextT()) const
      -> ssize_t {
    return ViewT(*this).CountProbedKeys(key_context);
  }

  // Insert a key into the set. If the key is already present, no insertion is
  // performed and that present key is available in the result. Otherwise a new
  // key is inserted and constructed from the argument and available in the
  // result.
  template <typename LookupKeyT>
  auto Insert(LookupKeyT lookup_key, KeyContextT key_context = KeyContextT())
      -> InsertResult;

  // Insert a key into the set and call the provided callback to allow in-place
  // construction of the key if not already present. The lookup key is passed
  // through to the callback so it needn't be captured and can be kept in a
  // register argument throughout.
  //
  // Example:
  // ```cpp
  //   m.Insert("widget", [](MyStringViewType lookup_key, void* key_storage) {
  //     new (key_storage) MyStringType(lookup_key);
  //   });
  // ```
  template <typename LookupKeyT, typename InsertCallbackT>
  auto Insert(LookupKeyT lookup_key, InsertCallbackT insert_cb,
              KeyContextT key_context = KeyContextT()) -> InsertResult
    requires std::invocable<InsertCallbackT, LookupKeyT, void*>;

  // Erase a key from the set.
  template <typename LookupKeyT>
  auto Erase(LookupKeyT lookup_key, KeyContextT key_context = KeyContextT())
      -> bool;

  // Clear all key/value pairs from the set but leave the underlying hashtable
  // allocated and in place.
  void Clear();

 protected:
  using ImplT::ImplT;
};

// A data structure for a set of keys.
//
// This set supports small size optimization (or "SSO"). The provided
// `SmallSize` type parameter indicates the size of an embedded buffer for
// storing sets small enough to fit. The default is zero, which always allocates
// a heap buffer on construction. When non-zero, must be a multiple of the
// `MaxGroupSize` which is currently 16. The library will check that the size is
// valid and provide an error at compile time if not. We don't automatically
// select the next multiple or otherwise fit the size to the constraints to make
// it clear in the code how much memory is used by the SSO buffer.
//
// This data structure optimizes heavily for small key types that are cheap to
// move and even copy. Using types with large keys or expensive to copy keys may
// create surprising performance bottlenecks. A `std::string` key should be fine
// with generally small strings, but if some or many strings are large heap
// allocations the performance of hashtable routines may be unacceptably bad and
// another data structure or key design is likely preferable.
//
// Note that this type should typically not appear on API boundaries; either
// `SetBase` or `SetView` should be used instead.
template <typename InputKeyT, ssize_t SmallSize = 0,
          typename InputKeyContextT = DefaultKeyContext>
class Set : public RawHashtable::TableImpl<SetBase<InputKeyT, InputKeyContextT>,
                                           SmallSize> {
  using BaseT = SetBase<InputKeyT, InputKeyContextT>;
  using ImplT = RawHashtable::TableImpl<BaseT, SmallSize>;

 public:
  using KeyT = typename BaseT::KeyT;

  Set() = default;
  Set(const Set& arg) = default;
  Set(Set&& arg) noexcept = default;

  // Reset the entire state of the hashtable to as it was when constructed,
  // throwing away any intervening allocations.
  void Reset();
};

template <typename InputKeyT, typename InputKeyContextT>
template <typename LookupKeyT>
auto SetView<InputKeyT, InputKeyContextT>::Contains(
    LookupKeyT lookup_key, KeyContextT key_context) const -> bool {
  return this->LookupEntry(lookup_key, key_context) != nullptr;
}

template <typename InputKeyT, typename InputKeyContextT>
template <typename LookupKeyT>
auto SetView<InputKeyT, InputKeyContextT>::Lookup(LookupKeyT lookup_key,
                                                  KeyContextT key_context) const
    -> LookupResult {
  EntryT* entry = this->LookupEntry(lookup_key, key_context);
  if (!entry) {
    return LookupResult();
  }

  return LookupResult(entry->key());
}

template <typename InputKeyT, typename InputKeyContextT>
template <typename CallbackT>
void SetView<InputKeyT, InputKeyContextT>::ForEach(CallbackT callback)
  requires(std::invocable<CallbackT, KeyT&>)
{
  this->ForEachEntry([callback](EntryT& entry) { callback(entry.key()); },
                     [](auto...) {});
}

template <typename InputKeyT, typename InputKeyContextT>
template <typename LookupKeyT>
auto SetBase<InputKeyT, InputKeyContextT>::Insert(LookupKeyT lookup_key,
                                                  KeyContextT key_context)
    -> InsertResult {
  return Insert(
      lookup_key,
      [](LookupKeyT lookup_key, void* key_storage) {
        new (key_storage) KeyT(std::move(lookup_key));
      },
      key_context);
}

template <typename InputKeyT, typename InputKeyContextT>
template <typename LookupKeyT, typename InsertCallbackT>
auto SetBase<InputKeyT, InputKeyContextT>::Insert(LookupKeyT lookup_key,
                                                  InsertCallbackT insert_cb,
                                                  KeyContextT key_context)
    -> InsertResult
  requires std::invocable<InsertCallbackT, LookupKeyT, void*>
{
  auto [entry, inserted] = this->InsertImpl(lookup_key, key_context);
  CARBON_DCHECK(entry) << "Should always result in a valid index.";

  if (LLVM_LIKELY(!inserted)) {
    return InsertResult(false, entry->key());
  }

  insert_cb(lookup_key, static_cast<void*>(&entry->key_storage));
  return InsertResult(true, entry->key());
}

template <typename InputKeyT, typename InputKeyContextT>
template <typename LookupKeyT>
auto SetBase<InputKeyT, InputKeyContextT>::Erase(LookupKeyT lookup_key,
                                                 KeyContextT key_context)
    -> bool {
  return this->EraseImpl(lookup_key, key_context);
}

template <typename InputKeyT, typename InputKeyContextT>
void SetBase<InputKeyT, InputKeyContextT>::Clear() {
  this->ClearImpl();
}

template <typename InputKeyT, ssize_t SmallSize, typename InputKeyContextT>
void Set<InputKeyT, SmallSize, InputKeyContextT>::Reset() {
  this->ResetImpl();
}

}  // namespace Carbon

#endif  // CARBON_COMMON_SET_H_