[libc++][format][5/6] Improve format_to_n.

Use a specialized buffer wrapper to limit the number of insertions in the
buffer. After the limit has been reached the buffer only needs to count
the number of insertions to return the buffer size required to store the
entire output.

Depends on D110498

Reviewed By: #libc, Mordante

Differential Revision: https://reviews.llvm.org/D110499

diff --git a/libcxx/include/__format/buffer.h b/libcxx/include/__format/buffer.h
index e5c3b9f..14926e3 100644 (file)
--- a/libcxx/include/__format/buffer.h
+++ b/libcxx/include/__format/buffer.h
@@ -11,12 +11,16 @@
 #define _LIBCPP___FORMAT_BUFFER_H
 
 #include <__algorithm/copy_n.h>
+#include <__algorithm/max.h>
+#include <__algorithm/min.h>
 #include <__algorithm/unwrap_iter.h>
 #include <__config>
 #include <__format/enable_insertable.h>
+#include <__format/format_to_n_result.h>
 #include <__format/formatter.h> // for __char_type TODO FMT Move the concept?
 #include <__iterator/back_insert_iterator.h>
 #include <__iterator/concepts.h>
+#include <__iterator/incrementable_traits.h>
 #include <__iterator/iterator_traits.h>
 #include <__iterator/wrap_iter.h>
 #include <__utility/move.h>
@@ -28,6 +32,9 @@
 #  pragma GCC system_header
 #endif
 
+_LIBCPP_PUSH_MACROS
+#include <__undef_macros>
+
 _LIBCPP_BEGIN_NAMESPACE_STD
 
 #if _LIBCPP_STD_VER > 17
@@ -267,10 +274,97 @@ private:
   size_t __size_{0};
 };
 
+/// The base of a buffer that counts and limits the number of insertions.
+template <class _OutIt, __formatter::__char_type _CharT, bool>
+  requires(output_iterator<_OutIt, const _CharT&>)
+struct _LIBCPP_TEMPLATE_VIS __format_to_n_buffer_base {
+  using _Size = iter_difference_t<_OutIt>;
+
+public:
+  _LIBCPP_HIDE_FROM_ABI explicit __format_to_n_buffer_base(_OutIt __out_it, _Size __n)
+      : __writer_(_VSTD::move(__out_it)), __n_(_VSTD::max(_Size(0), __n)) {}
+
+  _LIBCPP_HIDE_FROM_ABI void flush(_CharT* __ptr, size_t __size) {
+    if (_Size(__size_) <= __n_)
+      __writer_.flush(__ptr, _VSTD::min(_Size(__size), __n_ - __size_));
+    __size_ += __size;
+  }
+
+protected:
+  __internal_storage<_CharT> __storage_;
+  __output_buffer<_CharT> __output_{__storage_.begin(), __storage_.capacity(), this};
+  typename __writer_selector<_OutIt, _CharT>::type __writer_;
+
+  _Size __n_;
+  _Size __size_{0};
+};
+
+/// The base of a buffer that counts and limits the number of insertions.
+///
+/// This version is used when \c __enable_direct_output<_OutIt, _CharT> == true.
+///
+/// This class limits the size available the the direct writer so it will not
+/// exceed the maximum number of code units.
+template <class _OutIt, __formatter::__char_type _CharT>
+  requires(output_iterator<_OutIt, const _CharT&>)
+class _LIBCPP_TEMPLATE_VIS __format_to_n_buffer_base<_OutIt, _CharT, true> {
+  using _Size = iter_difference_t<_OutIt>;
+
+public:
+  _LIBCPP_HIDE_FROM_ABI explicit __format_to_n_buffer_base(_OutIt __out_it, _Size __n)
+      : __output_(_VSTD::__unwrap_iter(__out_it), __n, this), __writer_(_VSTD::move(__out_it)) {
+    if (__n <= 0) [[unlikely]]
+      __output_.reset(__storage_.begin(), __storage_.capacity());
+  }
+
+  _LIBCPP_HIDE_FROM_ABI void flush(_CharT* __ptr, size_t __size) {
+    // A flush to the direct writer happens in two occasions:
+    // - The format function has written the maximum number of allowed code
+    //   units. At this point it's no longer valid to write to this writer. So
+    //   switch to the internal storage. This internal storage doesn't need to
+    //   be written anywhere so the flush for that storage writes no output.
+    // - The format_to_n function is finished. In this case there's no need to
+    //   switch the buffer, but for simplicity the buffers are still switched.
+    // When the __n <= 0 the constructor already switched the buffers.
+    if (__size_ == 0 && __ptr != __storage_.begin()) {
+      __writer_.flush(__ptr, __size);
+      __output_.reset(__storage_.begin(), __storage_.capacity());
+    }
+
+    __size_ += __size;
+  }
+
+protected:
+  __internal_storage<_CharT> __storage_;
+  __output_buffer<_CharT> __output_;
+  __writer_direct<_OutIt, _CharT> __writer_;
+
+  _Size __size_{0};
+};
+
+/// The buffer that counts and limits the number of insertions.
+template <class _OutIt, __formatter::__char_type _CharT>
+  requires(output_iterator<_OutIt, const _CharT&>)
+struct _LIBCPP_TEMPLATE_VIS __format_to_n_buffer final
+    : public __format_to_n_buffer_base< _OutIt, _CharT, __enable_direct_output<_OutIt, _CharT>> {
+  using _Base = __format_to_n_buffer_base<_OutIt, _CharT, __enable_direct_output<_OutIt, _CharT>>;
+  using _Size = iter_difference_t<_OutIt>;
+
+public:
+  _LIBCPP_HIDE_FROM_ABI explicit __format_to_n_buffer(_OutIt __out_it, _Size __n) : _Base(_VSTD::move(__out_it), __n) {}
+  _LIBCPP_HIDE_FROM_ABI auto make_output_iterator() { return this->__output_.make_output_iterator(); }
+
+  _LIBCPP_HIDE_FROM_ABI format_to_n_result<_OutIt> result() && {
+    this->__output_.flush();
+    return {_VSTD::move(this->__writer_).out(), this->__size_};
+  }
+};
 } // namespace __format
 
 #endif //_LIBCPP_STD_VER > 17
 
 _LIBCPP_END_NAMESPACE_STD
 
+_LIBCPP_POP_MACROS
+
 #endif // _LIBCPP___FORMAT_BUFFER_H

diff --git a/libcxx/include/format b/libcxx/include/format
index 55ce2b1..98fa0b0 100644 (file)
--- a/libcxx/include/format
+++ b/libcxx/include/format
@@ -309,6 +309,9 @@ requires(output_iterator<_OutIt, const _CharT&>) _LIBCPP_HIDE_FROM_ABI _OutIt
   }
 }
 
+// The function is _LIBCPP_ALWAYS_INLINE since the compiler is bad at inlining
+// https://reviews.llvm.org/D110499#inline-1180704
+// TODO FMT Evaluate whether we want to file a Clang bug report regarding this.
 template <output_iterator<const char&> _OutIt>
 _LIBCPP_ALWAYS_INLINE _LIBCPP_HIDE_FROM_ABI _LIBCPP_AVAILABILITY_FORMAT _OutIt
 vformat_to(_OutIt __out_it, string_view __fmt, format_args __args) {
@@ -369,31 +372,27 @@ format(wstring_view __fmt, const _Args&... __args) {
 }
 #endif
 
+template <class _Context, class _OutIt, class _CharT>
+_LIBCPP_HIDE_FROM_ABI format_to_n_result<_OutIt> __vformat_to_n(_OutIt __out_it, iter_difference_t<_OutIt> __n,
+                                                                basic_string_view<_CharT> __fmt,
+                                                                basic_format_args<_Context> __args) {
+  __format::__format_to_n_buffer<_OutIt, _CharT> __buffer{_VSTD::move(__out_it), __n};
+  _VSTD::__format::__vformat_to(basic_format_parse_context{__fmt, __args.__size()},
+                                _VSTD::__format_context_create(__buffer.make_output_iterator(), __args));
+  return _VSTD::move(__buffer).result();
+}
+
 template <output_iterator<const char&> _OutIt, class... _Args>
-_LIBCPP_HIDE_FROM_ABI _LIBCPP_AVAILABILITY_FORMAT format_to_n_result<_OutIt>
-format_to_n(_OutIt __out_it, iter_difference_t<_OutIt> __n, string_view __fmt,
-            const _Args&... __args) {
-  // TODO FMT Improve PoC: using std::string is inefficient.
-  string __str = _VSTD::vformat(__fmt, _VSTD::make_format_args(__args...));
-  iter_difference_t<_OutIt> __s = __str.size();
-  iter_difference_t<_OutIt> __m =
-      _VSTD::clamp(__n, iter_difference_t<_OutIt>(0), __s);
-  __out_it = _VSTD::copy_n(__str.begin(), __m, _VSTD::move(__out_it));
-  return {_VSTD::move(__out_it), __s};
+_LIBCPP_ALWAYS_INLINE _LIBCPP_HIDE_FROM_ABI _LIBCPP_AVAILABILITY_FORMAT format_to_n_result<_OutIt>
+format_to_n(_OutIt __out_it, iter_difference_t<_OutIt> __n, string_view __fmt, const _Args&... __args) {
+  return _VSTD::__vformat_to_n<format_context>(_VSTD::move(__out_it), __n, __fmt, _VSTD::make_format_args(__args...));
 }
 
 #ifndef _LIBCPP_HAS_NO_WIDE_CHARACTERS
 template <output_iterator<const wchar_t&> _OutIt, class... _Args>
-_LIBCPP_HIDE_FROM_ABI _LIBCPP_AVAILABILITY_FORMAT format_to_n_result<_OutIt>
-format_to_n(_OutIt __out_it, iter_difference_t<_OutIt> __n, wstring_view __fmt,
-            const _Args&... __args) {
-  // TODO FMT Improve PoC: using std::string is inefficient.
-  wstring __str = _VSTD::vformat(__fmt, _VSTD::make_wformat_args(__args...));
-  iter_difference_t<_OutIt> __s = __str.size();
-  iter_difference_t<_OutIt> __m =
-      _VSTD::clamp(__n, iter_difference_t<_OutIt>(0), __s);
-  __out_it = _VSTD::copy_n(__str.begin(), __m, _VSTD::move(__out_it));
-  return {_VSTD::move(__out_it), __s};
+_LIBCPP_ALWAYS_INLINE _LIBCPP_HIDE_FROM_ABI _LIBCPP_AVAILABILITY_FORMAT format_to_n_result<_OutIt>
+format_to_n(_OutIt __out_it, iter_difference_t<_OutIt> __n, wstring_view __fmt, const _Args&... __args) {
+  return _VSTD::__vformat_to_n<wformat_context>(_VSTD::move(__out_it), __n, __fmt, _VSTD::make_wformat_args(__args...));
 }
 #endif
 
@@ -507,33 +506,30 @@ format(locale __loc, wstring_view __fmt, const _Args&... __args) {
 }
 #endif
 
+template <class _Context, class _OutIt, class _CharT>
+_LIBCPP_HIDE_FROM_ABI format_to_n_result<_OutIt> __vformat_to_n(_OutIt __out_it, iter_difference_t<_OutIt> __n,
+                                                                locale __loc, basic_string_view<_CharT> __fmt,
+                                                                basic_format_args<_Context> __args) {
+  __format::__format_to_n_buffer<_OutIt, _CharT> __buffer{_VSTD::move(__out_it), __n};
+  _VSTD::__format::__vformat_to(
+      basic_format_parse_context{__fmt, __args.__size()},
+      _VSTD::__format_context_create(__buffer.make_output_iterator(), __args, _VSTD::move(__loc)));
+  return _VSTD::move(__buffer).result();
+}
+
 template <output_iterator<const char&> _OutIt, class... _Args>
-_LIBCPP_HIDE_FROM_ABI _LIBCPP_AVAILABILITY_FORMAT format_to_n_result<_OutIt>
-format_to_n(_OutIt __out_it, iter_difference_t<_OutIt> __n, locale __loc,
-            string_view __fmt, const _Args&... __args) {
-  // TODO FMT Improve PoC: using std::string is inefficient.
-  string __str = _VSTD::vformat(_VSTD::move(__loc), __fmt,
-                                _VSTD::make_format_args(__args...));
-  iter_difference_t<_OutIt> __s = __str.size();
-  iter_difference_t<_OutIt> __m =
-      _VSTD::clamp(__n, iter_difference_t<_OutIt>(0), __s);
-  __out_it = _VSTD::copy_n(__str.begin(), __m, _VSTD::move(__out_it));
-  return {_VSTD::move(__out_it), __s};
+_LIBCPP_ALWAYS_INLINE _LIBCPP_HIDE_FROM_ABI _LIBCPP_AVAILABILITY_FORMAT format_to_n_result<_OutIt>
+format_to_n(_OutIt __out_it, iter_difference_t<_OutIt> __n, locale __loc, string_view __fmt, const _Args&... __args) {
+  return _VSTD::__vformat_to_n<format_context>(_VSTD::move(__out_it), __n, _VSTD::move(__loc), __fmt,
+                                               _VSTD::make_format_args(__args...));
 }
 
 #ifndef _LIBCPP_HAS_NO_WIDE_CHARACTERS
 template <output_iterator<const wchar_t&> _OutIt, class... _Args>
-_LIBCPP_HIDE_FROM_ABI _LIBCPP_AVAILABILITY_FORMAT format_to_n_result<_OutIt>
-format_to_n(_OutIt __out_it, iter_difference_t<_OutIt> __n, locale __loc,
-            wstring_view __fmt, const _Args&... __args) {
-  // TODO FMT Improve PoC: using std::string is inefficient.
-  wstring __str = _VSTD::vformat(_VSTD::move(__loc), __fmt,
-                                 _VSTD::make_wformat_args(__args...));
-  iter_difference_t<_OutIt> __s = __str.size();
-  iter_difference_t<_OutIt> __m =
-      _VSTD::clamp(__n, iter_difference_t<_OutIt>(0), __s);
-  __out_it = _VSTD::copy_n(__str.begin(), __m, _VSTD::move(__out_it));
-  return {_VSTD::move(__out_it), __s};
+_LIBCPP_ALWAYS_INLINE _LIBCPP_HIDE_FROM_ABI _LIBCPP_AVAILABILITY_FORMAT format_to_n_result<_OutIt>
+format_to_n(_OutIt __out_it, iter_difference_t<_OutIt> __n, locale __loc, wstring_view __fmt, const _Args&... __args) {
+  return _VSTD::__vformat_to_n<wformat_context>(_VSTD::move(__out_it), __n, _VSTD::move(__loc), __fmt,
+                                                _VSTD::make_wformat_args(__args...));
 }
 #endif