Enhance safety for vector assignment operations

Author: winner245

libcxx/include/__vector/vector.h

@@ -1025,9 +1025,14 @@ vector<_Tp, _Allocator>::__assign_with_size(_ForwardIterator __first, _Sentinel
       this->__destruct_at_end(__m);
     }
   } else {
+    __split_buffer<value_type, allocator_type&> __v(__recommend(__new_size), 0, __alloc_);
+    __v.__construct_at_end_with_size(__first, __new_size);
     __vdeallocate();
-    __vallocate(__recommend(__new_size));
-    __construct_at_end(__first, __last, __new_size);
+    this->__begin_ = __v.__begin_;
+    this->__end_   = __v.__end_;
+    this->__cap_   = __v.__cap_;
+    __v.__first_ = __v.__begin_ = __v.__end_ = __v.__cap_ = nullptr;
+    __annotate_new(__new_size);
   }
 }
 
@@ -1041,9 +1046,14 @@ _LIBCPP_CONSTEXPR_SINCE_CXX20 void vector<_Tp, _Allocator>::assign(size_type __n
     else
       this->__destruct_at_end(this->__begin_ + __n);
   } else {
+    __split_buffer<value_type, allocator_type&> __v(__recommend(__n), 0, __alloc_);
+    __v.__construct_at_end(__n, __u);
     __vdeallocate();
-    __vallocate(__recommend(static_cast<size_type>(__n)));
-    __construct_at_end(__n, __u);
+    this->__begin_ = __v.__begin_;
+    this->__end_   = __v.__end_;
+    this->__cap_   = __v.__cap_;
+    __v.__first_ = __v.__begin_ = __v.__end_ = __v.__cap_ = nullptr;
+    __annotate_new(__n);
   }
 }
 

libcxx/test/std/containers/sequences/vector/vector.cons/assign_exceptions.pass.cpp

@@ -0,0 +1,199 @@
+//===----------------------------------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+// <vector>
+
+// Check that vector assignments don't change the rhs vector when an operation throws an exception during rallocations triggered by assignments
+
+#include <cassert>
+#include <ranges>
+#include <vector>
+
+#include "allocators.h"
+#include "exception_test_helpers.h"
+#include "test_allocator.h"
+#include "test_macros.h"
+
+void test_allocation_exception() {
+  {
+    limited_alloc_wrapper<int> alloc1 = limited_allocator<int, 100>();
+    limited_alloc_wrapper<int> alloc2 = limited_allocator<int, 200>();
+    std::vector<int, limited_alloc_wrapper<int>> v(100, alloc1);
+    std::vector<int, limited_alloc_wrapper<int>> in(200, alloc2);
+    try { // Throw in copy-assignment operator during allocation
+      v = in;
+    } catch (const std::exception&) {
+    }
+    assert(v.size() == 100);
+  }
+
+  {
+    limited_alloc_wrapper<int> alloc1 = limited_allocator<int, 100>();
+    limited_alloc_wrapper<int> alloc2 = limited_allocator<int, 200>();
+    std::vector<int, limited_alloc_wrapper<int>> v(100, alloc1);
+    std::vector<int, limited_alloc_wrapper<int>> in(200, alloc2);
+    try { // Throw in move-assignment operator during allocation
+      v = std::move(in);
+    } catch (const std::exception&) {
+    }
+    assert(v.size() == 100);
+  }
+
+  {
+    std::vector<int, limited_allocator<int, 5>> v(5);
+    std::initializer_list<int> in{1, 2, 3, 4, 5, 6};
+    try { // Throw in operator=(initializer_list<value_type>) during allocation
+      v = in;
+    } catch (const std::exception&) {
+    }
+    assert(v.size() == 5);
+  }
+
+  {
+    std::vector<int, limited_allocator<int, 100>> v(100);
+    std::vector<int> in(101, 1);
+    try { // Throw in assign(_ForwardIterator, _ForwardIterator) during allocation
+      v.assign(in.begin(), in.end());
+    } catch (const std::exception&) {
+    }
+    assert(v.size() == 100);
+  }
+
+  {
+    std::vector<int, limited_allocator<int, 100>> v(100);
+    try { // Throw in assign(size_type, const_reference) during allocation
+      v.assign(101, 1);
+    } catch (const std::exception&) {
+    }
+    assert(v.size() == 100);
+  }
+
+  {
+    std::vector<int, limited_allocator<int, 5>> v(5);
+    std::initializer_list<int> in{1, 2, 3, 4, 5, 6};
+    try { // Throw in assign(initializer_list<value_type>) during allocation
+      v.assign(in);
+    } catch (const std::exception&) {
+    }
+    assert(v.size() == 5);
+  }
+
+#if TEST_STD_VER >= 23
+  {
+    std::vector<int, limited_allocator<int, 100>> v(100);
+    std::vector<int> in(101, 1);
+    try { // Throw in assign(_ForwardIterator, _ForwardIterator) during allocation
+      v.assign_range(in);
+    } catch (const std::exception&) {
+    }
+    assert(v.size() == 100);
+  }
+#endif
+}
+
+void test_construction_exception() {
+  {
+    int throw_after = 10;
+    throwing_t t    = throw_after;
+    std::vector<throwing_t> in(6, t);
+    std::vector<throwing_t> v(3, t);
+    try { // Throw in copy-assignment operator from element type during construction
+      v = in;
+    } catch (int) {
+    }
+    assert(v.size() == 3);
+  }
+
+  {
+    int throw_after = 10;
+    throwing_t t    = throw_after;
+    NONPOCMAAllocator<throwing_t> alloc1(1);
+    NONPOCMAAllocator<throwing_t> alloc2(2);
+    std::vector<throwing_t, NONPOCMAAllocator<throwing_t>> in(6, t, alloc1);
+    std::vector<throwing_t, NONPOCMAAllocator<throwing_t>> v(3, t, alloc2);
+    try { // Throw in move-assignment operator from element type during construction
+      v = std::move(in);
+    } catch (int) {
+    }
+    assert(v.size() == 3);
+  }
+
+  {
+    int throw_after = 10;
+    throwing_t t    = throw_after;
+    std::initializer_list<throwing_t> in{t, t, t, t, t, t};
+    std::vector<throwing_t> v(3, t);
+    try { // Throw in operator=(initializer_list<value_type>) from element type during construction
+      v = in;
+    } catch (int) {
+    }
+    assert(v.size() == 3);
+  }
+
+  {
+    std::vector<int> v(3);
+    try { // Throw in assign(_ForwardIterator, _ForwardIterator) from forward iterator during construction
+      v.assign(throwing_iterator<int, std::forward_iterator_tag>(),
+               throwing_iterator<int, std::forward_iterator_tag>(6));
+    } catch (int) {
+    }
+    assert(v.size() == 3);
+  }
+
+  {
+    int throw_after = 10;
+    throwing_t t    = throw_after;
+    std::vector<throwing_t> in(6, t);
+    std::vector<throwing_t> v(3, t);
+    try { // Throw in assign(_ForwardIterator, _ForwardIterator) from element type during construction
+      v.assign(in.begin(), in.end());
+    } catch (int) {
+    }
+    assert(v.size() == 3);
+  }
+#if TEST_STD_VER >= 23
+  {
+    int throw_after = 10;
+    throwing_t t    = throw_after;
+    std::vector<throwing_t> in(6, t);
+    std::vector<throwing_t> v(3, t);
+    try { // Throw in assign_range(_Range&&) from element type during construction
+      v.assign_range(in);
+    } catch (int) {
+    }
+    assert(v.size() == 3);
+  }
+#endif
+  {
+    int throw_after = 4;
+    throwing_t t    = throw_after;
+    std::vector<throwing_t> v(3, t);
+    try { // Throw in assign(size_type, const_reference) from element type during construction
+      v.assign(6, t);
+    } catch (int) {
+    }
+    assert(v.size() == 3);
+  }
+
+  {
+    int throw_after = 10;
+    throwing_t t    = throw_after;
+    std::initializer_list<throwing_t> in{t, t, t, t, t, t};
+    std::vector<throwing_t> v(3, t);
+    try { // Throw in assign(initializer_list<value_type>) from element type during construction
+      v.assign(in);
+    } catch (int) {
+    }
+    assert(v.size() == 3);
+  }
+}
+
+int main() {
+  test_allocation_exception();
+  test_construction_exception();
+}

libcxx/test/support/allocators.h

@@ -251,6 +251,50 @@ using POCCAAllocator = MaybePOCCAAllocator<T, /*POCCAValue = */true>;
 template <class T>
 using NonPOCCAAllocator = MaybePOCCAAllocator<T, /*POCCAValue = */false>;
 
+template <class T, bool POCMAValue>
+class MaybePOCMAAllocator {
+  template <class, bool>
+  friend class MaybePOCMAAllocator;
+
+public:
+  using propagate_on_container_move_assignment = std::integral_constant<bool, POCMAValue>;
+  using value_type                             = T;
+
+  template <class U>
+  struct rebind {
+    using other = MaybePOCMAAllocator<U, POCMAValue>;
+  };
+
+  constexpr MaybePOCMAAllocator(int id) : id_(id) {}
+
+  template <class U>
+  constexpr MaybePOCMAAllocator(const MaybePOCMAAllocator<U, POCMAValue>& other) : id_(other.id_) {}
+
+  constexpr T* allocate(std::size_t n) { return std::allocator<T>().allocate(n); }
+
+  constexpr void deallocate(T* p, std::size_t n) { std::allocator<T>().deallocate(p, n); }
+
+  constexpr int id() const { return id_; }
+
+  template <class U>
+  constexpr friend bool operator==(const MaybePOCMAAllocator& lhs, const MaybePOCMAAllocator<U, POCMAValue>& rhs) {
+    return lhs.id() == rhs.id();
+  }
+
+  template <class U>
+  constexpr friend bool operator!=(const MaybePOCMAAllocator& lhs, const MaybePOCMAAllocator<U, POCMAValue>& rhs) {
+    return !(lhs == rhs);
+  }
+
+private:
+  int id_;
+};
+
+template <class T>
+using POCMAAllocator = MaybePOCMAAllocator<T, /*POCMAValue = */ true>;
+template <class T>
+using NONPOCMAAllocator = MaybePOCMAAllocator<T, /*POCMAValue = */ false>;
+
 #endif // TEST_STD_VER >= 11
 
 #endif // ALLOCATORS_H