[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);
+      init(other.capacity, other.data, other.count);
+    }
+    return *this;
+  }
+
+  kmp_vector &operator=(kmp_vector &&other) noexcept {
+    if (this != &other) {
+      reset(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 {
+    for (size_t i = 0; i < count; i++) {
+      if (comp ? comp(data[i], value) : data[i] == value)
+        return true;
+    }
+    return false;
+  }
+
+  bool empty() const { return !count; }
+
+  // Check if the two vectors are equal with set semantics.
+  // Current implementation is naive O(n^2) and not optimized for performance.
+  // Handles duplicates correctly.
+  bool is_set_equal(const kmp_vector &other,
+                    bool (*comp)(const T &, const T &) = nullptr) const {
+    for (const T &val : *this) {
+      if (!other.contains(val, comp))
+        return false;
+    }
+    for (const T &val : other) {
+      if (!contains(val, comp))
+        return false;
+    }
+    return true;
+  }
+
+  // Add a new element to the end of the vector.
+  void push_back(const T &value) {
+    if (count == capacity)
+      grow();
+    if constexpr (std::is_trivially_copyable_v<T>) {
+      data[count++] = value;
+    } else {
+      new (&data[count++]) T(value);
+    }
+  }
+
+  // Reserve space for the given number of elements.
+  // (Note: does not shrink the vector.)
----------------
shiltian wrote:

is that what LLVM's ADT does as well?

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