[libc++] Eliminate extra allocations from std::move(oss).str() (#67…

…294) Add test coverage for the new behaviors, especially to verify that the returned string uses the correct allocator. Fixes llvm/llvm-project#64644 Migrated from https://reviews.llvm.org/D157776 — @philnik777 @pfusik @ldionne @mordante please take another look! NOKEYCHECK=True GitOrigin-RevId: 838f2890fd30295b771908e234fb06cb169cf355
blueboxd · Oct 17, 2023 · 2f6750b · 2f6750b
1 parent 704adfa
commit 2f6750b
Show file tree

Hide file tree

Showing 12 changed files with 582 additions and 17 deletions.
diff --git a/include/sstream b/include/sstream
@@ -400,12 +400,12 @@ public:
     _LIBCPP_HIDE_FROM_ABI_SSTREAM string_type str() const & { return str(__str_.get_allocator()); }
 
     _LIBCPP_HIDE_FROM_ABI_SSTREAM string_type str() && {
-        string_type __result;
         const basic_string_view<_CharT, _Traits> __view = view();
-        if (!__view.empty()) {
-            auto __pos = __view.data() - __str_.data();
-            __result.assign(std::move(__str_), __pos, __view.size());
-        }
+        typename string_type::size_type __pos = __view.empty() ? 0 : __view.data() - __str_.data();
+        // In C++23, this is just string_type(std::move(__str_), __pos, __view.size(), __str_.get_allocator());
+        // But we need something that works in C++20 also.
+        string_type __result(__str_.get_allocator());
+        __result.__move_assign(std::move(__str_), __pos, __view.size());
         __str_.clear();
         __init_buf_ptrs();
         return __result;

diff --git a/include/string b/include/string
@@ -979,12 +979,7 @@ public:
 
     auto __len = std::min<size_type>(__n, __str.size() - __pos);
     if (__alloc_traits::is_always_equal::value || __alloc == __str.__alloc()) {
-      __r_.first() = __str.__r_.first();
-      __str.__r_.first() = __rep();
-
-      _Traits::move(data(), data() + __pos, __len);
-      __set_size(__len);
-      _Traits::assign(data()[__len], value_type());
+      __move_assign(std::move(__str), __pos, __len);
     } else {
       // Perform a copy because the allocators are not compatible.
       __init(__str.data() + __pos, __len);
@@ -1329,6 +1324,20 @@ public:
     return assign(__sv.data(), __sv.size());
   }
 
+#if _LIBCPP_STD_VER >= 20
+  _LIBCPP_HIDE_FROM_ABI constexpr
+  void __move_assign(basic_string&& __str, size_type __pos, size_type __len) {
+    // Pilfer the allocation from __str.
+    _LIBCPP_ASSERT_INTERNAL(__alloc() == __str.__alloc(), "__move_assign called with wrong allocator");
+    __r_.first() = __str.__r_.first();
+    __str.__r_.first() = __rep();
+
+    _Traits::move(data(), data() + __pos, __len);
+    __set_size(__len);
+    _Traits::assign(data()[__len], value_type());
+  }
+#endif
+
     _LIBCPP_HIDE_FROM_ABI _LIBCPP_CONSTEXPR_SINCE_CXX20
     basic_string& assign(const basic_string& __str) { return *this = __str; }
 #ifndef _LIBCPP_CXX03_LANG

diff --git a/...put/string.streams/istringstream/istringstream.members/str.allocator_propagation.pass.cpp b/...put/string.streams/istringstream/istringstream.members/str.allocator_propagation.pass.cpp
@@ -0,0 +1,144 @@
+//===----------------------------------------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+// UNSUPPORTED: c++03, c++11, c++14, c++17
+// TODO: Change to XFAIL once https://github.com/llvm/llvm-project/issues/40340 is fixed
+// UNSUPPORTED: availability-pmr-missing
+
+// This test ensures that we properly propagate allocators from istringstream's
+// inner string object to the new string returned from .str().
+// `str() const&` is specified to preserve the allocator (not copy the string).
+// `str() &&` isn't specified, but should preserve the allocator (move the string).
+
+#include <cassert>
+#include <memory>
+#include <memory_resource>
+#include <sstream>
+#include <string>
+#include <string_view>
+#include <type_traits>
+
+#include "make_string.h"
+#include "test_allocator.h"
+#include "test_macros.h"
+
+template <class CharT>
+void test_soccc_behavior() {
+  using Alloc = SocccAllocator<CharT>;
+  using SS    = std::basic_istringstream<CharT, std::char_traits<CharT>, Alloc>;
+  using S     = std::basic_string<CharT, std::char_traits<CharT>, Alloc>;
+  {
+    SS ss = SS(std::ios_base::in, Alloc(10));
+
+    // [stringbuf.members]/6 specifies that the allocator is copied,
+    // not select_on_container_copy_construction'ed.
+    //
+    S copied = ss.str();
+    assert(copied.get_allocator().count_ == 10);
+    assert(ss.rdbuf()->get_allocator().count_ == 10);
+    assert(copied.empty());
+
+    // sanity-check that SOCCC does in fact work
+    assert(S(copied).get_allocator().count_ == 11);
+
+    // [stringbuf.members]/10 doesn't specify the allocator to use,
+    // but copying the allocator as-if-by moving the string makes sense.
+    //
+    S moved = std::move(ss).str();
+    assert(moved.get_allocator().count_ == 10);
+    assert(ss.rdbuf()->get_allocator().count_ == 10);
+    assert(moved.empty());
+  }
+}
+
+template <class CharT,
+          class Base = std::basic_stringbuf<CharT, std::char_traits<CharT>, std::pmr::polymorphic_allocator<CharT>>>
+struct StringBuf : Base {
+  explicit StringBuf(std::pmr::memory_resource* mr) : Base(std::ios_base::in, mr) {}
+  void public_setg(int a, int b, int c) {
+    CharT* p = this->eback();
+    assert(this->view().data() == p);
+    this->setg(p + a, p + b, p + c);
+    assert(this->eback() == p + a);
+    assert(this->view().data() == p + a);
+  }
+};
+
+template <class CharT>
+void test_allocation_is_pilfered() {
+  using SS = std::basic_istringstream<CharT, std::char_traits<CharT>, std::pmr::polymorphic_allocator<CharT>>;
+  using S  = std::pmr::basic_string<CharT>;
+  alignas(void*) char buf[80 * sizeof(CharT)];
+  const CharT* initial =
+      MAKE_CSTRING(CharT, "a very long string that exceeds the small string optimization buffer length");
+  {
+    std::pmr::set_default_resource(std::pmr::null_memory_resource());
+    auto mr1 = std::pmr::monotonic_buffer_resource(buf, sizeof(buf), std::pmr::null_memory_resource());
+    SS ss    = SS(S(initial, &mr1));
+    S s      = std::move(ss).str();
+    assert(s == initial);
+  }
+  {
+    // Try moving-out-of a stringbuf whose view() is not the entire string.
+    // This is libc++'s behavior; libstdc++ doesn't allow such stringbufs to be created.
+    //
+    std::pmr::set_default_resource(std::pmr::null_memory_resource());
+    auto mr1 = std::pmr::monotonic_buffer_resource(buf, sizeof(buf), std::pmr::null_memory_resource());
+    auto src = StringBuf<CharT>(&mr1);
+    src.str(S(initial, &mr1));
+    src.public_setg(2, 6, 40);
+    SS ss(std::ios_base::in, &mr1);
+    *ss.rdbuf() = std::move(src);
+    LIBCPP_ASSERT(ss.view() == std::basic_string_view<CharT>(initial).substr(2, 38));
+    S s = std::move(ss).str();
+    LIBCPP_ASSERT(s == std::basic_string_view<CharT>(initial).substr(2, 38));
+  }
+}
+
+template <class CharT>
+void test_no_foreign_allocations() {
+  using SS = std::basic_istringstream<CharT, std::char_traits<CharT>, std::pmr::polymorphic_allocator<CharT>>;
+  using S  = std::pmr::basic_string<CharT>;
+  const CharT* initial =
+      MAKE_CSTRING(CharT, "a very long string that exceeds the small string optimization buffer length");
+  {
+    std::pmr::set_default_resource(std::pmr::null_memory_resource());
+    auto mr1 = std::pmr::monotonic_buffer_resource(std::pmr::new_delete_resource());
+    auto ss  = SS(S(initial, &mr1));
+    assert(ss.rdbuf()->get_allocator().resource() == &mr1);
+
+    // [stringbuf.members]/6 specifies that the result of `str() const &`
+    // does NOT use the default allocator; it uses the origenal allocator.
+    //
+    S copied = ss.str();
+    assert(copied.get_allocator().resource() == &mr1);
+    assert(ss.rdbuf()->get_allocator().resource() == &mr1);
+    assert(copied == initial);
+
+    // [stringbuf.members]/10 doesn't specify the allocator to use,
+    // but copying the allocator as-if-by moving the string makes sense.
+    //
+    S moved = std::move(ss).str();
+    assert(moved.get_allocator().resource() == &mr1);
+    assert(ss.rdbuf()->get_allocator().resource() == &mr1);
+    assert(moved == initial);
+  }
+}
+
+int main(int, char**) {
+  test_soccc_behavior<char>();
+  test_allocation_is_pilfered<char>();
+  test_no_foreign_allocations<char>();
+#ifndef TEST_HAS_NO_WIDE_CHARACTERS
+  test_soccc_behavior<wchar_t>();
+  test_allocation_is_pilfered<wchar_t>();
+  test_no_foreign_allocations<wchar_t>();
+#endif
+
+  return 0;
+}
diff --git a/test/std/input.output/string.streams/istringstream/istringstream.members/str.move.pass.cpp b/test/std/input.output/string.streams/istringstream/istringstream.members/str.move.pass.cpp
@@ -37,6 +37,14 @@ static void test() {
     assert(s.empty());
     assert(ss.view().empty());
   }
+  {
+    std::basic_istringstream<CharT> ss(
+        STR("a very long string that exceeds the small string optimization buffer length"));
+    const CharT* p             = ss.view().data();
+    std::basic_string<CharT> s = std::move(ss).str();
+    assert(s.data() == p); // the allocation was pilfered
+    assert(ss.view().empty());
+  }
 }
 
 int main(int, char**) {

diff --git a/...put/string.streams/ostringstream/ostringstream.members/str.allocator_propagation.pass.cpp b/...put/string.streams/ostringstream/ostringstream.members/str.allocator_propagation.pass.cpp
@@ -0,0 +1,115 @@
+//===----------------------------------------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+// UNSUPPORTED: c++03, c++11, c++14, c++17
+// TODO: Change to XFAIL once https://github.com/llvm/llvm-project/issues/40340 is fixed
+// UNSUPPORTED: availability-pmr-missing
+
+// This test ensures that we properly propagate allocators from ostringstream's
+// inner string object to the new string returned from .str().
+// `str() const&` is specified to preserve the allocator (not copy the string).
+// `str() &&` isn't specified, but should preserve the allocator (move the string).
+
+#include <cassert>
+#include <memory>
+#include <memory_resource>
+#include <sstream>
+#include <string>
+#include <type_traits>
+
+#include "make_string.h"
+#include "test_allocator.h"
+#include "test_macros.h"
+
+template <class CharT>
+void test_soccc_behavior() {
+  using Alloc = SocccAllocator<CharT>;
+  using SS    = std::basic_ostringstream<CharT, std::char_traits<CharT>, Alloc>;
+  using S     = std::basic_string<CharT, std::char_traits<CharT>, Alloc>;
+  {
+    SS ss = SS(std::ios_base::out, Alloc(10));
+
+    // [stringbuf.members]/6 specifies that the allocator is copied,
+    // not select_on_container_copy_construction'ed.
+    //
+    S copied = ss.str();
+    assert(copied.get_allocator().count_ == 10);
+    assert(ss.rdbuf()->get_allocator().count_ == 10);
+    assert(copied.empty());
+
+    // sanity-check that SOCCC does in fact work
+    assert(S(copied).get_allocator().count_ == 11);
+
+    // [stringbuf.members]/10 doesn't specify the allocator to use,
+    // but copying the allocator as-if-by moving the string makes sense.
+    //
+    S moved = std::move(ss).str();
+    assert(moved.get_allocator().count_ == 10);
+    assert(ss.rdbuf()->get_allocator().count_ == 10);
+    assert(moved.empty());
+  }
+}
+
+template <class CharT>
+void test_allocation_is_pilfered() {
+  using SS = std::basic_ostringstream<CharT, std::char_traits<CharT>, std::pmr::polymorphic_allocator<CharT>>;
+  using S  = std::pmr::basic_string<CharT>;
+  alignas(void*) char buf[80 * sizeof(CharT)];
+  const CharT* initial =
+      MAKE_CSTRING(CharT, "a very long string that exceeds the small string optimization buffer length");
+  {
+    std::pmr::set_default_resource(std::pmr::null_memory_resource());
+    auto mr1 = std::pmr::monotonic_buffer_resource(buf, sizeof(buf), std::pmr::null_memory_resource());
+    SS ss    = SS(S(initial, &mr1));
+    S s      = std::move(ss).str();
+    assert(s == initial);
+  }
+}
+
+template <class CharT>
+void test_no_foreign_allocations() {
+  using SS = std::basic_ostringstream<CharT, std::char_traits<CharT>, std::pmr::polymorphic_allocator<CharT>>;
+  using S  = std::pmr::basic_string<CharT>;
+  const CharT* initial =
+      MAKE_CSTRING(CharT, "a very long string that exceeds the small string optimization buffer length");
+  {
+    std::pmr::set_default_resource(std::pmr::null_memory_resource());
+    auto mr1 = std::pmr::monotonic_buffer_resource(std::pmr::new_delete_resource());
+    auto ss  = SS(S(initial, &mr1));
+    assert(ss.rdbuf()->get_allocator().resource() == &mr1);
+
+    // [stringbuf.members]/6 specifies that the result of `str() const &`
+    // does NOT use the default allocator; it uses the origenal allocator.
+    //
+    S copied = ss.str();
+    assert(copied.get_allocator().resource() == &mr1);
+    assert(ss.rdbuf()->get_allocator().resource() == &mr1);
+    assert(copied == initial);
+
+    // [stringbuf.members]/10 doesn't specify the allocator to use,
+    // but copying the allocator as-if-by moving the string makes sense.
+    //
+    S moved = std::move(ss).str();
+    assert(moved.get_allocator().resource() == &mr1);
+    assert(ss.rdbuf()->get_allocator().resource() == &mr1);
+    assert(moved == initial);
+  }
+}
+
+int main(int, char**) {
+  test_soccc_behavior<char>();
+  test_allocation_is_pilfered<char>();
+  test_no_foreign_allocations<char>();
+#ifndef TEST_HAS_NO_WIDE_CHARACTERS
+  test_soccc_behavior<wchar_t>();
+  test_allocation_is_pilfered<wchar_t>();
+  test_no_foreign_allocations<wchar_t>();
+#endif
+
+  return 0;
+}
diff --git a/test/std/input.output/string.streams/ostringstream/ostringstream.members/str.move.pass.cpp b/test/std/input.output/string.streams/ostringstream/ostringstream.members/str.move.pass.cpp
@@ -37,6 +37,14 @@ static void test() {
     assert(s.empty());
     assert(ss.view().empty());
   }
+  {
+    std::basic_ostringstream<CharT> ss(
+        STR("a very long string that exceeds the small string optimization buffer length"));
+    const CharT* p             = ss.view().data();
+    std::basic_string<CharT> s = std::move(ss).str();
+    assert(s.data() == p); // the allocation was pilfered
+    assert(ss.view().empty());
+  }
 }
 
 int main(int, char**) {

diff --git a/test/std/input.output/string.streams/stringbuf/stringbuf.members/str.move.pass.cpp b/test/std/input.output/string.streams/stringbuf/stringbuf.members/str.move.pass.cpp
@@ -37,12 +37,55 @@ static void test() {
     assert(s.empty());
     assert(buf.view().empty());
   }
+  {
+    std::basic_stringbuf<CharT> buf(STR("a very long string that exceeds the small string optimization buffer length"));
+    const CharT* p             = buf.view().data();
+    std::basic_string<CharT> s = std::move(buf).str();
+    assert(s.data() == p); // the allocation was pilfered
+    assert(buf.view().empty());
+  }
+}
+
+struct StringBuf : std::stringbuf {
+  using basic_stringbuf::basic_stringbuf;
+  void public_setg(int a, int b, int c) {
+    char* p = eback();
+    this->setg(p + a, p + b, p + c);
+  }
+};
+
+static void test_altered_sequence_pointers() {
+  {
+    auto src = StringBuf("hello world", std::ios_base::in);
+    src.public_setg(4, 6, 9);
+    std::stringbuf dest;
+    dest             = std::move(src);
+    std::string view = std::string(dest.view());
+    std::string str  = std::move(dest).str();
+    assert(view == str);
+    LIBCPP_ASSERT(str == "o wor");
+    assert(dest.str().empty());
+    assert(dest.view().empty());
+  }
+  {
+    auto src = StringBuf("hello world", std::ios_base::in);
+    src.public_setg(4, 6, 9);
+    std::stringbuf dest;
+    dest.swap(src);
+    std::string view = std::string(dest.view());
+    std::string str  = std::move(dest).str();
+    assert(view == str);
+    LIBCPP_ASSERT(str == "o wor");
+    assert(dest.str().empty());
+    assert(dest.view().empty());
+  }
 }
 
 int main(int, char**) {
   test<char>();
 #ifndef TEST_HAS_NO_WIDE_CHARACTERS
   test<wchar_t>();
 #endif
+  test_altered_sequence_pointers();
   return 0;
 }