[llvm-branch-commits] [openmp] [libomp] Add kmp_vector (ADT 2/2) (PR #176163)

[Michael Kruse via llvm-branch-commits]

================
@@ -154,4 +155,188 @@ class kmp_str_ref final {
   const char *end() const { return sv.data() + length(); }
 };
 
+/// kmp_vector is a vector class for managing small vectors.
+/// INLINE_THRESHOLD: Number of elements in the inline array. If exceeded, the
+/// vector will grow dynamically.
+template <typename T, size_t INLINE_THRESHOLD = 8> class kmp_vector final {
+  static_assert(std::is_copy_constructible_v<T>,
+                "T must be copy constructible");
+  static_assert(std::is_destructible_v<T>, "T must be destructible");
+
+  T inline_data[INLINE_THRESHOLD];
+  T *data = inline_data;
+  size_t count = 0;
+  size_t capacity = INLINE_THRESHOLD;
+
+  void copy_data(T *dst, const T *src, size_t num_elements) {
+    if constexpr (std::is_trivially_copyable_v<T>) {
+      memcpy(dst, src, num_elements * sizeof(T));
+    } else {
+      for (size_t i = 0; i < num_elements; i++)
+        new (&dst[i]) T(src[i]); // copy-construct to memory
+    }
+  }
+
+  /// Grow by ~1.5x
+  void grow() {
+    size_t new_capacity = capacity + (capacity / 2) + 1;
+    resize(new_capacity);
+  }
+
+  void init(size_t new_capacity, const T *init_data, size_t new_count) {
+    assert(new_capacity >= new_count &&
+           "more elements requested than capacity");
+    if (new_capacity > INLINE_THRESHOLD)
+      resize(new_capacity);
+    if (init_data)
+      copy_data(data, init_data, new_count);
+    count = new_count;
+  }
+
+  /// Move data from other vector to this vector (which must be emptied before)
+  void move_from(kmp_vector &&other) {
+    assert(empty() && "must be empty before overwriting");
+    if (other.data == other.inline_data) {
+      // Cannot move inline data, must copy.
+      init(other.capacity, other.data, other.count);
+    } else {
+      // Steal dynamic data.
+      data = other.data;
+      count = other.count;
+      capacity = other.capacity;
+    }
+    other.reset(/*free_data=*/false);
+  }
+
+  void reset(bool free_data) {
+    if (free_data && data != inline_data) {
+      clear();
+      KMP_INTERNAL_FREE(data);
+    }
+    data = inline_data;
+    count = 0;
+    capacity = INLINE_THRESHOLD;
+  }
+
+  /// resize only changes the capacity, not the size (i.e., the number of
+  /// actually used elements)
+  void resize(size_t new_capacity) {
+    // Currently only supports growing the capacity. (Consequently, doesn't need
+    // to worry about going from a dynamic array back to an inline array.)
+    assert(new_capacity > capacity && "resize() only supports growing");
+    capacity = new_capacity;
+    T *old_data = data != inline_data ? data : nullptr;
+    data =
+        static_cast<T *>(KMP_INTERNAL_REALLOC(old_data, capacity * sizeof(T)));
+    assert(data);
+    // Copy the data to the new array if we didn't use a dynamic array before.
+    if (!old_data)
+      copy_data(data, inline_data, count);
+  }
+
+public:
+  ~kmp_vector() { reset(/*free_data=*/true); }
+
+  explicit kmp_vector(size_t capacity = 0) { init(capacity, nullptr, 0); }
+
+  kmp_vector(size_t capacity, const T *init_data, size_t count) {
+    init(capacity, init_data, count);
+  }
+
+  kmp_vector(const kmp_vector &other) {
+    init(other.capacity, other.data, other.count);
+  }
+
+  kmp_vector(kmp_vector &&other) noexcept { move_from(std::move(other)); }
+
+  kmp_vector &operator=(const kmp_vector &other) {
+    if (this != &other) {
+      reset(/*free_data=*/true);
+      init(other.capacity, other.data, other.count);
+    }
+    return *this;
+  }
+
+  kmp_vector &operator=(kmp_vector &&other) noexcept {
+    if (this != &other) {
+      reset(/*free_data=*/true);
+      move_from(std::move(other));
+    }
+    return *this;
+  }
+
+  /// Destroy all elements in the vector. Doesn't free the memory.
+  void clear() {
+    if constexpr (!std::is_trivially_destructible_v<T>) {
+      for (size_t i = 0; i < count; i++)
+        data[i].~T();
+    }
+    count = 0;
+  }
+
+  /// Check if the vector contains the given value.
+  /// If a comparator is provided, it will be used to compare the values.
+  /// Otherwise, the equality operator will be used.
+  bool contains(const T &value,
+                bool (*comp)(const T &, const T &) = nullptr) const {
----------------
Meinersbur wrote:

Use template like for `kmp_str_ref`? I think this choice would be fine, but I think we also want a consistent API.

https://github.com/llvm/llvm-project/pull/176163