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

[Shilei Tian via llvm-branch-commits]

================
@@ -136,4 +137,185 @@ 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
+    }
+  }
+
+  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);
+    if (new_capacity > inline_threshold)
+      resize(new_capacity);
+    if (init_data)
+      copy_data(data, init_data, new_count);
+    count = new_count;
+  }
+
+  void move_from(kmp_vector &&other) {
+    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(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(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(true);
----------------
shiltian wrote:

ditto

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