[llvm-branch-commits] [libcxx] [libc++][format][3/3] Improves formatting performance. (PR #108990)

[Louis Dionne via llvm-branch-commits]

================
@@ -319,188 +353,222 @@ struct _LIBCPP_TEMPLATE_VIS __back_insert_iterator_container<back_insert_iterato
   using type = _Container;
 };
 
-/// Write policy for inserting the buffer in a container.
-template <class _Container>
-class _LIBCPP_TEMPLATE_VIS __writer_container {
+// A dynamically growing buffer.
+template <__fmt_char_type _CharT>
+class _LIBCPP_TEMPLATE_VIS __allocating_buffer : public __output_buffer<_CharT> {
 public:
-  using _CharT = typename _Container::value_type;
+  __allocating_buffer(const __allocating_buffer&)            = delete;
+  __allocating_buffer& operator=(const __allocating_buffer&) = delete;
 
-  _LIBCPP_HIDE_FROM_ABI explicit __writer_container(back_insert_iterator<_Container> __out_it)
-      : __container_{__out_it.__get_container()} {}
+  [[nodiscard]] _LIBCPP_HIDE_FROM_ABI __allocating_buffer() : __allocating_buffer{nullptr} {}
 
-  _LIBCPP_HIDE_FROM_ABI auto __out_it() { return std::back_inserter(*__container_); }
+  [[nodiscard]]
+  _LIBCPP_HIDE_FROM_ABI explicit __allocating_buffer(__max_output_size* __max_output_size)
+      : __output_buffer<_CharT>{__buffer_, __buffer_size_, __prepare_write, __max_output_size} {}
 
-  _LIBCPP_HIDE_FROM_ABI void __flush(_CharT* __ptr, size_t __n) {
-    __container_->insert(__container_->end(), __ptr, __ptr + __n);
+  _LIBCPP_HIDE_FROM_ABI ~__allocating_buffer() {
+    if (__ptr_ != __buffer_) {
+      ranges::destroy_n(__ptr_, this->__size());
+      allocator_traits<_Alloc>::deallocate(__alloc_, __ptr_, this->__capacity());
+    }
   }
 
+  [[nodiscard]] _LIBCPP_HIDE_FROM_ABI basic_string_view<_CharT> __view() { return {__ptr_, this->__size()}; }
+
 private:
-  _Container* __container_;
-};
+  // At the moment the allocator is hard-code. There might be reasons to have
+  // an allocator trait in the future. This ensures forward compatibility.
+  using _Alloc = allocator<_CharT>;
+  _LIBCPP_NO_UNIQUE_ADDRESS _Alloc __alloc_;
 
-/// Selects the type of the writer used for the output iterator.
-template <class _OutIt, class _CharT>
-class _LIBCPP_TEMPLATE_VIS __writer_selector {
-  using _Container = typename __back_insert_iterator_container<_OutIt>::type;
+  // Since allocating is expensive the class has a small internal buffer. When
+  // its capacity is exceeded a dynamic buffer will be allocated.
+  static constexpr size_t __buffer_size_ = 256;
+  _CharT __buffer_[__buffer_size_];
 
-public:
-  using type =
-      conditional_t<!same_as<_Container, void>,
-                    __writer_container<_Container>,
-                    conditional_t<__enable_direct_output<_OutIt, _CharT>,
-                                  __writer_direct<_OutIt, _CharT>,
-                                  __writer_iterator<_OutIt, _CharT>>>;
+  _CharT* __ptr_{__buffer_};
+
+  _LIBCPP_HIDE_FROM_ABI void __grow_buffer(size_t __capacity) {
+    if (__capacity < __buffer_size_)
+      return;
+
+    _LIBCPP_ASSERT_INTERNAL(__capacity > this->__capacity(), "the buffer must grow");
+    auto __result = std::__allocate_at_least(__alloc_, __capacity);
+    auto __guard  = std::__make_exception_guard([&] {
+      allocator_traits<_Alloc>::deallocate(__alloc_, __result.ptr, __result.count);
+    });
+    // This shouldn't throw, but just to be safe. Note that at -O1 this
+    // guard is optimized away so there is no runtime overhead.
+    new (__result.ptr) _CharT[__result.count];
+    std::copy_n(__ptr_, this->__size(), __result.ptr);
+    __guard.__complete();
+    if (__ptr_ != __buffer_) {
+      ranges::destroy_n(__ptr_, this->__capacity());
+      allocator_traits<_Alloc>::deallocate(__alloc_, __ptr_, this->__capacity());
+    }
+
+    __ptr_ = __result.ptr;
+    this->__buffer_moved(__ptr_, __result.count);
+  }
+
+  _LIBCPP_HIDE_FROM_ABI void __prepare_write(size_t __size_hint) {
+    __grow_buffer(std::max<size_t>(this->__capacity() + __size_hint, this->__capacity() * 1.6));
+  }
+
+  _LIBCPP_HIDE_FROM_ABI static void __prepare_write(__output_buffer<_CharT>& __buffer, size_t __size_hint) {
+    static_cast<__allocating_buffer<_CharT>&>(__buffer).__prepare_write(__size_hint);
+  }
 };
 
-/// The generic formatting buffer.
+// A buffer that directly writes to the underlying buffer.
 template <class _OutIt, __fmt_char_type _CharT>
-  requires(output_iterator<_OutIt, const _CharT&>)
-class _LIBCPP_TEMPLATE_VIS __format_buffer {
-  using _Storage =
-      conditional_t<__enable_direct_output<_OutIt, _CharT>, __direct_storage<_CharT>, __internal_storage<_CharT>>;
-
+class _LIBCPP_TEMPLATE_VIS __direct_iterator_buffer : public __output_buffer<_CharT> {
 public:
-  _LIBCPP_HIDE_FROM_ABI explicit __format_buffer(_OutIt __out_it)
-    requires(same_as<_Storage, __internal_storage<_CharT>>)
-      : __output_(__storage_.__begin(), __storage_.__buffer_size, this), __writer_(std::move(__out_it)) {}
+  [[nodiscard]] _LIBCPP_HIDE_FROM_ABI explicit __direct_iterator_buffer(_OutIt __out_it)
+      : __direct_iterator_buffer{__out_it, nullptr} {}
 
-  _LIBCPP_HIDE_FROM_ABI explicit __format_buffer(_OutIt __out_it)
-    requires(same_as<_Storage, __direct_storage<_CharT>>)
-      : __output_(std::__unwrap_iter(__out_it), size_t(-1), this), __writer_(std::move(__out_it)) {}
+  [[nodiscard]]
+  _LIBCPP_HIDE_FROM_ABI explicit __direct_iterator_buffer(_OutIt __out_it, __max_output_size* __max_output_size)
+      : __output_buffer<_CharT>{std::__unwrap_iter(__out_it), __buffer_size, __prepare_write, __max_output_size},
----------------
ldionne wrote:

For a follow-up patch, you can avoid unwrapping so we catch OOB accesses.

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