Optimize small_vector copy and move for inline storage

Summary:
Codegen for copy and move assignment is suboptimal when the vectors involved are inline, in particular if the destination is default-constructed (which is very common) and even more so if the value type is trivially copyable.

The main optimization is that if the storage is inline and the type is trivially-copyable, we can just copy the whole storage, regardless of the size of the container. This avoids a branchy loop, and it's usually just a handful of MOVs.

While at it, optimized all the related code paths by avoiding delegating to `swap()` and `assign()` when possible, which introduce further branches that are hard to optimize. Also move and copy in place when the vector is non-empty, avoiding unnecessary destructions.

Also fixed a couple of bugs:
- The move constructor did not clear its argument when inline, inconsistently with the move assignment and `std::vector`
- The move assignment operator passed its heap storage to the argument instead of freeing it.

Reviewed By: yfeldblum

Differential Revision: D29592519

fbshipit-source-id: 6281cdda775568d28619afea8b7ca2fb168c7e5d

folly/small_vector.h

@@ -199,6 +199,26 @@ void populateMemForward(T* mem, std::size_t n, Function const& op) {
   }
 }
 
+/*
+ * Copies `fromSize` elements from `from' to `to', where `to' is only
+ * initialized up to `toSize`, but has enough storage for `fromSize'. If
+ * `toSize' > `fromSize', the extra elements are destructed.
+ */
+template <class Iterator1, class Iterator2>
+void partiallyUninitializedCopy(
+    Iterator1 from, size_t fromSize, Iterator2 to, size_t toSize) {
+  const size_t minSize = std::min(fromSize, toSize);
+  std::copy(from, from + minSize, to);
+  if (fromSize > toSize) {
+    std::uninitialized_copy(from + minSize, from + fromSize, to + minSize);
+  } else {
+    for (auto it = to + minSize; it != to + toSize; ++it) {
+      using Value = typename std::decay<decltype(*it)>::type;
+      it->~Value();
+    }
+  }
+}
+
 template <class SizeType, bool ShouldUseHeap>
 struct IntegralSizePolicyBase {
   typedef SizeType InternalSizeType;
@@ -245,6 +265,8 @@ struct IntegralSizePolicyBase {
 
   void swapSizePolicy(IntegralSizePolicyBase& o) { std::swap(size_, o.size_); }
 
+  void resetSizePolicy() { size_ = 0; }
+
  protected:
   static bool constexpr kShouldUseHeap = ShouldUseHeap;
 
@@ -484,6 +506,11 @@ class small_vector : public detail::small_vector_base<
   small_vector(const std::allocator<Value>&) {}
 
   small_vector(small_vector const& o) {
+    if (folly::is_trivially_copyable<Value>::value && !o.isExtern()) {
+      copyInlineTrivial<Value>(o);
+      return;
+    }
+
     auto n = o.size();
     makeSize(n);
     {
@@ -497,14 +524,24 @@ class small_vector : public detail::small_vector_base<
   small_vector(small_vector&& o) noexcept(
       std::is_nothrow_move_constructible<Value>::value) {
     if (o.isExtern()) {
-      swap(o);
+      this->u.pdata_.heap_ = o.u.pdata_.heap_;
+      this->swapSizePolicy(o);
+      if (kHasInlineCapacity) {
+        this->u.setCapacity(o.u.getCapacity());
+      }
     } else {
-      auto n = o.size();
-      std::uninitialized_copy(
-          std::make_move_iterator(o.begin()),
-          std::make_move_iterator(o.end()),
-          begin());
-      this->setSize(n);
+      if (folly::is_trivially_copyable<Value>::value) {
+        copyInlineTrivial<Value>(o);
+        o.resetSizePolicy();
+      } else {
+        auto n = o.size();
+        std::uninitialized_copy(
+            std::make_move_iterator(o.begin()),
+            std::make_move_iterator(o.end()),
+            begin());
+        this->setSize(n);
+        o.clear();
+      }
     }
   }
 
@@ -537,18 +574,51 @@ class small_vector : public detail::small_vector_base<
 
   small_vector& operator=(small_vector const& o) {
     if (FOLLY_LIKELY(this != &o)) {
-      assign(o.begin(), o.end());
+      if (folly::is_trivially_copyable<Value>::value && !this->isExtern() &&
+          !o.isExtern()) {
+        copyInlineTrivial<Value>(o);
+      } else if (o.size() < capacity()) {
+        const size_t oSize = o.size();
+        detail::partiallyUninitializedCopy(o.begin(), oSize, begin(), size());
+        this->setSize(oSize);
+      } else {
+        assign(o.begin(), o.end());
+      }
     }
     return *this;
   }
 
   small_vector& operator=(small_vector&& o) noexcept(
       std::is_nothrow_move_constructible<Value>::value) {
-    // TODO: optimization:
-    // if both are internal, use move assignment where possible
     if (FOLLY_LIKELY(this != &o)) {
-      clear();
-      swap(o);
+      // If either is external, reduce to the default-constructed case for this,
+      // since there is nothing that we can move in-place.
+      if (this->isExtern() || o.isExtern()) {
+        reset();
+      }
+
+      if (!o.isExtern()) {
+        if (folly::is_trivially_copyable<Value>::value) {
+          copyInlineTrivial<Value>(o);
+          o.resetSizePolicy();
+        } else {
+          const size_t oSize = o.size();
+          detail::partiallyUninitializedCopy(
+              std::make_move_iterator(o.u.buffer()),
+              oSize,
+              this->u.buffer(),
+              size());
+          this->setSize(oSize);
+          o.clear();
+        }
+      } else {
+        this->u.pdata_.heap_ = o.u.pdata_.heap_;
+        // this was already reset above, so it's empty and internal.
+        this->swapSizePolicy(o);
+        if (kHasInlineCapacity) {
+          this->u.setCapacity(o.u.getCapacity());
+        }
+      }
     }
     return *this;
   }
@@ -605,15 +675,15 @@ class small_vector : public detail::small_vector_base<
     if (this->isExtern() && o.isExtern()) {
       this->swapSizePolicy(o);
 
-      auto thisCapacity = this->capacity();
-      auto oCapacity = o.capacity();
-
       auto* tmp = u.pdata_.heap_;
       u.pdata_.heap_ = o.u.pdata_.heap_;
       o.u.pdata_.heap_ = tmp;
 
-      this->setCapacity(oCapacity);
-      o.setCapacity(thisCapacity);
+      if (kHasInlineCapacity) {
+        const auto capacity_ = this->u.getCapacity();
+        this->setCapacity(o.u.getCapacity());
+        o.u.setCapacity(capacity_);
+      }
 
       return;
     }
@@ -665,7 +735,7 @@ class small_vector : public detail::small_vector_base<
         for (; i < oldIntern.size(); ++i) {
           oldIntern[i].~value_type();
         }
-        oldIntern.setSize(0);
+        oldIntern.resetSizePolicy();
         oldExtern.u.pdata_.heap_ = oldExternHeap;
         oldExtern.setCapacity(oldExternCapacity);
       });
@@ -958,6 +1028,27 @@ class small_vector : public detail::small_vector_base<
     this->setSize(sz);
   }
 
+  template <class T>
+  typename std::enable_if<folly::is_trivially_copyable<T>::value>::type
+  copyInlineTrivial(small_vector const& o) {
+    // Copy the entire inline storage, instead of just the size, to make the
+    // loop fixed-size and unrollable.
+    std::copy(o.u.buffer(), o.u.buffer() + MaxInline, u.buffer());
+    this->setSize(o.size());
+  }
+
+  template <class T>
+  typename std::enable_if<!folly::is_trivially_copyable<T>::value>::type
+  copyInlineTrivial(small_vector const&) {
+    assume_unreachable();
+  }
+
+  void reset() {
+    clear();
+    freeHeap();
+    this->resetSizePolicy();
+  }
+
   // The std::false_type argument is part of disambiguating the
   // iterator insert functions from integral types (see insert().)
   template <class It>

folly/test/small_vector_test.cpp

@@ -26,6 +26,7 @@
 #include <vector>
 
 #include <boost/algorithm/string.hpp>
+#include <fmt/format.h>
 
 #include <folly/Conv.h>
 #include <folly/Traits.h>
@@ -1191,3 +1192,121 @@ TEST(small_vector, erase_if) {
   EXPECT_EQ(3u, v[1]);
   EXPECT_EQ(5u, v[2]);
 }
+
+namespace {
+
+class NonTrivialInt {
+ public:
+  NonTrivialInt() {}
+  /* implicit */ NonTrivialInt(int value)
+      : value_(std::make_shared<int>(value)) {}
+
+  operator int() const { return *value_; }
+
+ private:
+  std::shared_ptr<const int> value_;
+};
+
+// Move and copy constructor and assignment have several special cases depending
+// on relative sizes, so test all combinations.
+template <class T, size_t N>
+void testMoveAndCopy() {
+  const auto fill = [](auto& v, size_t n) {
+    v.resize(n);
+    std::iota(v.begin(), v.end(), 0);
+  };
+  const auto verify = [](const auto& v, size_t n) {
+    ASSERT_EQ(v.size(), n);
+    for (size_t i = 0; i < n; ++i) {
+      EXPECT_EQ(v[i], i);
+    }
+  };
+
+  using Vec = small_vector<T, N>;
+  SCOPED_TRACE(fmt::format("N = {}", N));
+
+  const size_t kMinCapacity = Vec{}.capacity();
+
+  for (size_t from = 0; from < 16; ++from) {
+    SCOPED_TRACE(fmt::format("from = {}", from));
+
+    {
+      SCOPED_TRACE("Move-construction");
+      Vec a;
+      fill(a, from);
+      const auto aCapacity = a.capacity();
+      Vec b = std::move(a);
+      verify(b, from);
+      EXPECT_EQ(b.capacity(), aCapacity);
+      verify(a, 0);
+      EXPECT_EQ(a.capacity(), kMinCapacity);
+    }
+    {
+      SCOPED_TRACE("Copy-construction");
+      Vec a;
+      fill(a, from);
+      Vec b = a;
+      verify(b, from);
+      verify(a, from);
+    }
+
+    for (size_t to = 0; to < 16; ++to) {
+      SCOPED_TRACE(fmt::format("to = {}", to));
+      {
+        SCOPED_TRACE("Move-assignment");
+        Vec a;
+        fill(a, from);
+        Vec b;
+        fill(b, to);
+
+        const auto aCapacity = a.capacity();
+        b = std::move(a);
+        verify(b, from);
+        EXPECT_EQ(b.capacity(), aCapacity);
+        verify(a, 0);
+        EXPECT_EQ(a.capacity(), kMinCapacity);
+      }
+      {
+        SCOPED_TRACE("Copy-assignment");
+        Vec a;
+        fill(a, from);
+        Vec b;
+        fill(b, to);
+
+        b = a;
+        verify(b, from);
+        verify(a, from);
+      }
+      {
+        SCOPED_TRACE("swap");
+        Vec a;
+        fill(a, from);
+        Vec b;
+        fill(b, to);
+
+        const auto aCapacity = a.capacity();
+        const auto bCapacity = b.capacity();
+        swap(a, b);
+        verify(b, from);
+        EXPECT_EQ(b.capacity(), aCapacity);
+        verify(a, to);
+        EXPECT_EQ(a.capacity(), bCapacity);
+      }
+    }
+  }
+}
+
+} // namespace
+
+TEST(small_vector, MoveAndCopyTrivial) {
+  testMoveAndCopy<int, 0>();
+  // Capacity does not fit inline.
+  testMoveAndCopy<int, 2>();
+  // Capacity does fits inline.
+  testMoveAndCopy<int, 4>();
+}
+
+TEST(small_vector, MoveAndCopyNonTrivial) {
+  testMoveAndCopy<NonTrivialInt, 0>();
+  testMoveAndCopy<NonTrivialInt, 4>();
+}