[llvm] [WIP][ADT] Avoid slow size queries on append (PR #136365)

[Jakub Kuderski via llvm-commits]

https://github.com/kuhar updated https://github.com/llvm/llvm-project/pull/136365

>From 078be5edcf584d380718f6e5b7d10325065d219a Mon Sep 17 00:00:00 2001
From: Jakub Kuderski <jakub at nod-labs.com>
Date: Fri, 18 Apr 2025 16:24:20 -0400
Subject: [PATCH 1/2] [WIP][ADT] Avoid slow size queries on append

Only calculate the size and reserve upfront when the input iterators
support cheap (O(1)) distance.

This is a follow up to:
https://github.com/llvm/llvm-project/pull/136066#issuecomment-2814895109.
---
 llvm/include/llvm/ADT/SmallVector.h | 61 +++++++++++++++++++++++------
 1 file changed, 48 insertions(+), 13 deletions(-)

diff --git a/llvm/include/llvm/ADT/SmallVector.h b/llvm/include/llvm/ADT/SmallVector.h
index bd3e887e36bce..104a3aaf79418 100644
--- a/llvm/include/llvm/ADT/SmallVector.h
+++ b/llvm/include/llvm/ADT/SmallVector.h
@@ -36,10 +36,18 @@ template <typename T> class ArrayRef;
 
 template <typename IteratorT> class iterator_range;
 
-template <class Iterator>
-using EnableIfConvertibleToInputIterator = std::enable_if_t<std::is_convertible<
+namespace detail {
+template <typename Iterator, typename IteratorCategory>
+inline constexpr bool IsOfIteratorCategory = std::is_convertible_v<
     typename std::iterator_traits<Iterator>::iterator_category,
-    std::input_iterator_tag>::value>;
+    IteratorCategory>;
+
+template <typename Iterator>
+using EnableIfConvertibleToInputIterator =
+    std::enable_if_t<std::is_convertible_v<
+        typename std::iterator_traits<Iterator>::iterator_category,
+        std::input_iterator_tag>>;
+} // namespace detail
 
 /// This is all the stuff common to all SmallVectors.
 ///
@@ -679,13 +687,37 @@ class SmallVectorImpl : public SmallVectorTemplateBase<T> {
   void swap(SmallVectorImpl &RHS);
 
   /// Add the specified range to the end of the SmallVector.
-  template <typename ItTy, typename = EnableIfConvertibleToInputIterator<ItTy>>
+  template <typename ItTy,
+            typename = detail::EnableIfConvertibleToInputIterator<ItTy>>
   void append(ItTy in_start, ItTy in_end) {
     this->assertSafeToAddRange(in_start, in_end);
-    size_type NumInputs = std::distance(in_start, in_end);
-    this->reserve(this->size() + NumInputs);
-    this->uninitialized_copy(in_start, in_end, this->end());
-    this->set_size(this->size() + NumInputs);
+    if constexpr (detail::IsOfIteratorCategory<
+                      ItTy, std::random_access_iterator_tag>) {
+      // Only reserve the required extra size upfront when the size calculation
+      // is guaranteed to be O(1).
+      size_type NumInputs = std::distance(in_start, in_end);
+      this->reserve(this->size() + NumInputs);
+      this->uninitialized_copy(in_start, in_end, this->end());
+      this->set_size(this->size() + NumInputs);
+    } else {
+      // Otherwise, append using `in_end` as the sentinel and reserve more space
+      // as necessary.
+      for (ItTy It = in_start; It != in_end;) {
+        iterator Dest = this->end();
+        size_type InitialSize = this->size();
+        size_type ExtraCap = this->capacity() - InitialSize;
+        size_type NumCopied = 0;
+        for (; NumCopied != ExtraCap && It != in_end; ++It, ++NumCopied) {
+          ::new ((void *)(Dest + NumCopied)) T(*It);
+        }
+        size_type NewSize = InitialSize + NumCopied;
+        this->set_size(NewSize);
+
+        if (It != in_end) {
+          this->reserve(NewSize + 1);
+        }
+      }
+    }
   }
 
   /// Append \p NumInputs copies of \p Elt to the end.
@@ -720,7 +752,8 @@ class SmallVectorImpl : public SmallVectorTemplateBase<T> {
   // FIXME: Consider assigning over existing elements, rather than clearing &
   // re-initializing them - for all assign(...) variants.
 
-  template <typename ItTy, typename = EnableIfConvertibleToInputIterator<ItTy>>
+  template <typename ItTy,
+            typename = detail::EnableIfConvertibleToInputIterator<ItTy>>
   void assign(ItTy in_start, ItTy in_end) {
     this->assertSafeToReferenceAfterClear(in_start, in_end);
     clear();
@@ -871,7 +904,8 @@ class SmallVectorImpl : public SmallVectorTemplateBase<T> {
     return I;
   }
 
-  template <typename ItTy, typename = EnableIfConvertibleToInputIterator<ItTy>>
+  template <typename ItTy,
+            typename = detail::EnableIfConvertibleToInputIterator<ItTy>>
   iterator insert(iterator I, ItTy From, ItTy To) {
     // Convert iterator to elt# to avoid invalidating iterator when we reserve()
     size_t InsertElt = I - this->begin();
@@ -887,8 +921,8 @@ class SmallVectorImpl : public SmallVectorTemplateBase<T> {
     this->assertSafeToAddRange(From, To);
 
     size_t NumToInsert = std::distance(From, To);
-
-    // Ensure there is enough space.
+    // Ensure there is enough space so that we do not have to re-allocate mid
+    // insertion.
     reserve(this->size() + NumToInsert);
 
     // Uninvalidate the iterator.
@@ -1212,7 +1246,8 @@ class LLVM_GSL_OWNER SmallVector : public SmallVectorImpl<T>,
     this->assign(Size, Value);
   }
 
-  template <typename ItTy, typename = EnableIfConvertibleToInputIterator<ItTy>>
+  template <typename ItTy,
+            typename = detail::EnableIfConvertibleToInputIterator<ItTy>>
   SmallVector(ItTy S, ItTy E) : SmallVectorImpl<T>(N) {
     this->append(S, E);
   }

>From 3d891d2893e978b9baf08eb2dcd2305ef5a31dc8 Mon Sep 17 00:00:00 2001
From: Jakub Kuderski <jakub at nod-labs.com>
Date: Fri, 18 Apr 2025 16:55:26 -0400
Subject: [PATCH 2/2] Address comments

---
 llvm/include/llvm/ADT/SmallVector.h | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/llvm/include/llvm/ADT/SmallVector.h b/llvm/include/llvm/ADT/SmallVector.h
index 104a3aaf79418..0f89ecd323e5a 100644
--- a/llvm/include/llvm/ADT/SmallVector.h
+++ b/llvm/include/llvm/ADT/SmallVector.h
@@ -707,9 +707,9 @@ class SmallVectorImpl : public SmallVectorTemplateBase<T> {
         size_type InitialSize = this->size();
         size_type ExtraCap = this->capacity() - InitialSize;
         size_type NumCopied = 0;
-        for (; NumCopied != ExtraCap && It != in_end; ++It, ++NumCopied) {
+        for (; NumCopied != ExtraCap && It != in_end; ++It, ++NumCopied)
           ::new ((void *)(Dest + NumCopied)) T(*It);
-        }
+
         size_type NewSize = InitialSize + NumCopied;
         this->set_size(NewSize);
 
@@ -910,7 +910,8 @@ class SmallVectorImpl : public SmallVectorTemplateBase<T> {
     // Convert iterator to elt# to avoid invalidating iterator when we reserve()
     size_t InsertElt = I - this->begin();
 
-    if (I == this->end()) {  // Important special case for empty vector.
+    // Important special case for appending to a vector, including empty vector.
+    if (I == this->end()) {
       append(From, To);
       return this->begin()+InsertElt;
     }