SPIRV-Tools/source/enum_set.h
Nathan Gauër 17d9669d51
enumset: add iterator based constructor/insert (#5344)
Expanding a bit the EnumSet API to have iterator-based
insert and constructors (like the STL).
This is also a pre-requisite from the capability-trimming pass as
it allows to build a const set from a constexpr std::array easily.

Signed-off-by: Nathan Gauër <brioche@google.com>
2023-07-20 17:54:50 +00:00

470 lines
15 KiB
C++

// Copyright (c) 2023 Google Inc.
//
// 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
//
// http://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.
#include <algorithm>
#include <cassert>
#include <cstdint>
#include <functional>
#include <initializer_list>
#include <limits>
#include <type_traits>
#include <vector>
#ifndef SOURCE_ENUM_SET_H_
#define SOURCE_ENUM_SET_H_
#include "source/latest_version_spirv_header.h"
namespace spvtools {
// This container is optimized to store and retrieve unsigned enum values.
// The base model for this implementation is an open-addressing hashtable with
// linear probing. For small enums (max index < 64), all operations are O(1).
//
// - Enums are stored in buckets (64 contiguous values max per bucket)
// - Buckets ranges don't overlap, but don't have to be contiguous.
// - Enums are packed into 64-bits buckets, using 1 bit per enum value.
//
// Example:
// - MyEnum { A = 0, B = 1, C = 64, D = 65 }
// - 2 buckets are required:
// - bucket 0, storing values in the range [ 0; 64[
// - bucket 1, storing values in the range [64; 128[
//
// - Buckets are stored in a sorted vector (sorted by bucket range).
// - Retrieval is done by computing the theoretical bucket index using the enum
// value, and
// doing a linear scan from this position.
// - Insertion is done by retrieving the bucket and either:
// - inserting a new bucket in the sorted vector when no buckets has a
// compatible range.
// - setting the corresponding bit in the bucket.
// This means insertion in the middle/beginning can cause a memmove when no
// bucket is available. In our case, this happens at most 23 times for the
// largest enum we have (Opcodes).
template <typename T>
class EnumSet {
private:
using BucketType = uint64_t;
using ElementType = std::underlying_type_t<T>;
static_assert(std::is_enum_v<T>, "EnumSets only works with enums.");
static_assert(std::is_signed_v<ElementType> == false,
"EnumSet doesn't supports signed enums.");
// Each bucket can hold up to `kBucketSize` distinct, contiguous enum values.
// The first value a bucket can hold must be aligned on `kBucketSize`.
struct Bucket {
// bit mask to store `kBucketSize` enums.
BucketType data;
// 1st enum this bucket can represent.
T start;
friend bool operator==(const Bucket& lhs, const Bucket& rhs) {
return lhs.start == rhs.start && lhs.data == rhs.data;
}
};
// How many distinct values can a bucket hold? 1 bit per value.
static constexpr size_t kBucketSize = sizeof(BucketType) * 8ULL;
public:
class Iterator {
public:
typedef Iterator self_type;
typedef T value_type;
typedef T& reference;
typedef T* pointer;
typedef std::forward_iterator_tag iterator_category;
typedef size_t difference_type;
Iterator(const Iterator& other)
: set_(other.set_),
bucketIndex_(other.bucketIndex_),
bucketOffset_(other.bucketOffset_) {}
Iterator& operator++() {
do {
if (bucketIndex_ >= set_->buckets_.size()) {
bucketIndex_ = set_->buckets_.size();
bucketOffset_ = 0;
break;
}
if (bucketOffset_ + 1 == kBucketSize) {
bucketOffset_ = 0;
++bucketIndex_;
} else {
++bucketOffset_;
}
} while (bucketIndex_ < set_->buckets_.size() &&
!set_->HasEnumAt(bucketIndex_, bucketOffset_));
return *this;
}
Iterator operator++(int) {
Iterator old = *this;
operator++();
return old;
}
T operator*() const {
assert(set_->HasEnumAt(bucketIndex_, bucketOffset_) &&
"operator*() called on an invalid iterator.");
return GetValueFromBucket(set_->buckets_[bucketIndex_], bucketOffset_);
}
bool operator!=(const Iterator& other) const {
return set_ != other.set_ || bucketOffset_ != other.bucketOffset_ ||
bucketIndex_ != other.bucketIndex_;
}
bool operator==(const Iterator& other) const {
return !(operator!=(other));
}
Iterator& operator=(const Iterator& other) {
set_ = other.set_;
bucketIndex_ = other.bucketIndex_;
bucketOffset_ = other.bucketOffset_;
return *this;
}
private:
Iterator(const EnumSet* set, size_t bucketIndex, ElementType bucketOffset)
: set_(set), bucketIndex_(bucketIndex), bucketOffset_(bucketOffset) {}
private:
const EnumSet* set_ = nullptr;
// Index of the bucket in the vector.
size_t bucketIndex_ = 0;
// Offset in bits in the current bucket.
ElementType bucketOffset_ = 0;
friend class EnumSet;
};
// Required to allow the use of std::inserter.
using value_type = T;
using const_iterator = Iterator;
using iterator = Iterator;
public:
iterator cbegin() const noexcept {
auto it = iterator(this, /* bucketIndex= */ 0, /* bucketOffset= */ 0);
if (buckets_.size() == 0) {
return it;
}
// The iterator has the logic to find the next valid bit. If the value 0
// is not stored, use it to find the next valid bit.
if (!HasEnumAt(it.bucketIndex_, it.bucketOffset_)) {
++it;
}
return it;
}
iterator begin() const noexcept { return cbegin(); }
iterator cend() const noexcept {
return iterator(this, buckets_.size(), /* bucketOffset= */ 0);
}
iterator end() const noexcept { return cend(); }
// Creates an empty set.
EnumSet() : buckets_(0), size_(0) {}
// Creates a set and store `value` in it.
EnumSet(T value) : EnumSet() { insert(value); }
// Creates a set and stores each `values` in it.
EnumSet(std::initializer_list<T> values) : EnumSet() {
for (auto item : values) {
insert(item);
}
}
// Creates a set, and insert `count` enum values pointed by `array` in it.
EnumSet(ElementType count, const T* array) : EnumSet() {
for (ElementType i = 0; i < count; i++) {
insert(array[i]);
}
}
// Creates a set initialized with the content of the range [begin; end[.
template <class InputIt>
EnumSet(InputIt begin, InputIt end) : EnumSet() {
for (; begin != end; ++begin) {
insert(*begin);
}
}
// Copies the EnumSet `other` into a new EnumSet.
EnumSet(const EnumSet& other)
: buckets_(other.buckets_), size_(other.size_) {}
// Moves the EnumSet `other` into a new EnumSet.
EnumSet(EnumSet&& other)
: buckets_(std::move(other.buckets_)), size_(other.size_) {}
// Deep-copies the EnumSet `other` into this EnumSet.
EnumSet& operator=(const EnumSet& other) {
buckets_ = other.buckets_;
size_ = other.size_;
return *this;
}
// Matches std::unordered_set::insert behavior.
std::pair<iterator, bool> insert(const T& value) {
const size_t index = FindBucketForValue(value);
const ElementType offset = ComputeBucketOffset(value);
if (index >= buckets_.size() ||
buckets_[index].start != ComputeBucketStart(value)) {
size_ += 1;
InsertBucketFor(index, value);
return std::make_pair(Iterator(this, index, offset), true);
}
auto& bucket = buckets_[index];
const auto mask = ComputeMaskForValue(value);
if (bucket.data & mask) {
return std::make_pair(Iterator(this, index, offset), false);
}
size_ += 1;
bucket.data |= ComputeMaskForValue(value);
return std::make_pair(Iterator(this, index, offset), true);
}
// Inserts `value` in the set if possible.
// Similar to `std::unordered_set::insert`, except the hint is ignored.
// Returns an iterator to the inserted element, or the element preventing
// insertion.
iterator insert(const_iterator, const T& value) {
return insert(value).first;
}
// Inserts `value` in the set if possible.
// Similar to `std::unordered_set::insert`, except the hint is ignored.
// Returns an iterator to the inserted element, or the element preventing
// insertion.
iterator insert(const_iterator, T&& value) { return insert(value).first; }
// Inserts all the values in the range [`first`; `last[.
// Similar to `std::unordered_set::insert`.
template <class InputIt>
void insert(InputIt first, InputIt last) {
for (auto it = first; it != last; ++it) {
insert(*it);
}
}
// Removes the value `value` into the set.
// Similar to `std::unordered_set::erase`.
// Returns the number of erased elements.
size_t erase(const T& value) {
const size_t index = FindBucketForValue(value);
if (index >= buckets_.size() ||
buckets_[index].start != ComputeBucketStart(value)) {
return 0;
}
auto& bucket = buckets_[index];
const auto mask = ComputeMaskForValue(value);
if (!(bucket.data & mask)) {
return 0;
}
size_ -= 1;
bucket.data &= ~mask;
if (bucket.data == 0) {
buckets_.erase(buckets_.cbegin() + index);
}
return 1;
}
// Returns true if `value` is present in the set.
bool contains(T value) const {
const size_t index = FindBucketForValue(value);
if (index >= buckets_.size() ||
buckets_[index].start != ComputeBucketStart(value)) {
return false;
}
auto& bucket = buckets_[index];
return bucket.data & ComputeMaskForValue(value);
}
// Returns the 1 if `value` is present in the set, `0` otherwise.
inline size_t count(T value) const { return contains(value) ? 1 : 0; }
// Returns true if the set is holds no values.
inline bool empty() const { return size_ == 0; }
// Returns the number of enums stored in this set.
size_t size() const { return size_; }
// Returns true if this set contains at least one value contained in `in_set`.
// Note: If `in_set` is empty, this function returns true.
bool HasAnyOf(const EnumSet<T>& in_set) const {
if (in_set.empty()) {
return true;
}
auto lhs = buckets_.cbegin();
auto rhs = in_set.buckets_.cbegin();
while (lhs != buckets_.cend() && rhs != in_set.buckets_.cend()) {
if (lhs->start == rhs->start) {
if (lhs->data & rhs->data) {
// At least 1 bit is shared. Early return.
return true;
}
lhs++;
rhs++;
continue;
}
// LHS bucket is smaller than the current RHS bucket. Catching up on RHS.
if (lhs->start < rhs->start) {
lhs++;
continue;
}
// Otherwise, RHS needs to catch up on LHS.
rhs++;
}
return false;
}
private:
// Returns the index of the last bucket in which `value` could be stored.
static constexpr inline size_t ComputeLargestPossibleBucketIndexFor(T value) {
return static_cast<size_t>(value) / kBucketSize;
}
// Returns the smallest enum value that could be contained in the same bucket
// as `value`.
static constexpr inline T ComputeBucketStart(T value) {
return static_cast<T>(kBucketSize *
ComputeLargestPossibleBucketIndexFor(value));
}
// Returns the index of the bit that corresponds to `value` in the bucket.
static constexpr inline ElementType ComputeBucketOffset(T value) {
return static_cast<ElementType>(value) % kBucketSize;
}
// Returns the bitmask used to represent the enum `value` in its bucket.
static constexpr inline BucketType ComputeMaskForValue(T value) {
return 1ULL << ComputeBucketOffset(value);
}
// Returns the `enum` stored in `bucket` at `offset`.
// `offset` is the bit-offset in the bucket storage.
static constexpr inline T GetValueFromBucket(const Bucket& bucket,
BucketType offset) {
return static_cast<T>(static_cast<ElementType>(bucket.start) + offset);
}
// For a given enum `value`, finds the bucket index that could contain this
// value. If no such bucket is found, the index at which the new bucket should
// be inserted is returned.
size_t FindBucketForValue(T value) const {
// Set is empty, insert at 0.
if (buckets_.size() == 0) {
return 0;
}
const T wanted_start = ComputeBucketStart(value);
assert(buckets_.size() > 0 &&
"Size must not be 0 here. Has the code above changed?");
size_t index = std::min(buckets_.size() - 1,
ComputeLargestPossibleBucketIndexFor(value));
// This loops behaves like std::upper_bound with a reverse iterator.
// Buckets are sorted. 3 main cases:
// - The bucket matches
// => returns the bucket index.
// - The found bucket is larger
// => scans left until it finds the correct bucket, or insertion point.
// - The found bucket is smaller
// => We are at the end, so we return past-end index for insertion.
for (; buckets_[index].start >= wanted_start; index--) {
if (index == 0) {
return 0;
}
}
return index + 1;
}
// Creates a new bucket to store `value` and inserts it at `index`.
// If the `index` is past the end, the bucket is inserted at the end of the
// vector.
void InsertBucketFor(size_t index, T value) {
const T bucket_start = ComputeBucketStart(value);
Bucket bucket = {1ULL << ComputeBucketOffset(value), bucket_start};
auto it = buckets_.emplace(buckets_.begin() + index, std::move(bucket));
#if defined(NDEBUG)
(void)it; // Silencing unused variable warning.
#else
assert(std::next(it) == buckets_.end() ||
std::next(it)->start > bucket_start);
assert(it == buckets_.begin() || std::prev(it)->start < bucket_start);
#endif
}
// Returns true if the bucket at `bucketIndex/ stores the enum at
// `bucketOffset`, false otherwise.
bool HasEnumAt(size_t bucketIndex, BucketType bucketOffset) const {
assert(bucketIndex < buckets_.size());
assert(bucketOffset < kBucketSize);
return buckets_[bucketIndex].data & (1ULL << bucketOffset);
}
// Returns true if `lhs` and `rhs` hold the exact same values.
friend bool operator==(const EnumSet& lhs, const EnumSet& rhs) {
if (lhs.size_ != rhs.size_) {
return false;
}
if (lhs.buckets_.size() != rhs.buckets_.size()) {
return false;
}
return lhs.buckets_ == rhs.buckets_;
}
// Returns true if `lhs` and `rhs` hold at least 1 different value.
friend bool operator!=(const EnumSet& lhs, const EnumSet& rhs) {
return !(lhs == rhs);
}
// Storage for the buckets.
std::vector<Bucket> buckets_;
// How many enums is this set storing.
size_t size_ = 0;
};
// A set of spv::Capability.
using CapabilitySet = EnumSet<spv::Capability>;
} // namespace spvtools
#endif // SOURCE_ENUM_SET_H_