From 15f5872f8097cbdc1f154b815af897d4687ad5d4 Mon Sep 17 00:00:00 2001
From: Vladislav Shchapov <vladislav@shchapov.ru>
Date: Wed, 3 Jul 2024 21:19:33 +0500
Subject: [PATCH] Improve std::complex formatter

Signed-off-by: Vladislav Shchapov <vladislav@shchapov.ru>
---
 include/fmt/std.h | 76 ++++++++++++++++++++++++++++++++++-------------
 test/std-test.cc  | 24 +++++++++++++++
 2 files changed, 79 insertions(+), 21 deletions(-)

diff --git a/include/fmt/std.h b/include/fmt/std.h
index 528e4f1caa3ee..688e423f8e75f 100644
--- a/include/fmt/std.h
+++ b/include/fmt/std.h
@@ -631,33 +631,67 @@ struct formatter<std::atomic_flag, Char> : formatter<bool, Char> {
 #endif  // __cpp_lib_atomic_flag_test
 
 FMT_EXPORT
-template <typename F, typename Char>
-struct formatter<std::complex<F>, Char> : nested_formatter<F, Char> {
+template <typename T, typename Char> struct formatter<std::complex<T>, Char> {
  private:
-  // Functor because C++11 doesn't support generic lambdas.
-  struct writer {
-    const formatter<std::complex<F>, Char>* f;
-    const std::complex<F>& c;
-
-    template <typename OutputIt>
-    FMT_CONSTEXPR auto operator()(OutputIt out) -> OutputIt {
-      if (c.real() != 0) {
-        auto format_full = detail::string_literal<Char, '(', '{', '}', '+', '{',
-                                                  '}', 'i', ')'>{};
-        return fmt::format_to(out, basic_string_view<Char>(format_full),
-                              f->nested(c.real()), f->nested(c.imag()));
-      }
-      auto format_imag = detail::string_literal<Char, '{', '}', 'i'>{};
-      return fmt::format_to(out, basic_string_view<Char>(format_imag),
-                            f->nested(c.imag()));
+  detail::dynamic_format_specs<Char> specs_;
+
+  template <typename FormatContext, typename OutputIt>
+  FMT_CONSTEXPR auto do_format(const std::complex<T>& c,
+                               detail::dynamic_format_specs<Char>& specs,
+                               FormatContext& ctx,
+                               OutputIt out) const -> OutputIt {
+    if (c.real() != 0) {
+      *out++ = Char('(');
+      out = detail::write<Char>(out, c.real(), specs, ctx.locale());
+      specs.sign = sign::plus;
+      out = detail::write<Char>(out, c.imag(), specs, ctx.locale());
+      if (!detail::isfinite(c.imag())) *out++ = Char(' ');
+      *out++ = Char('i');
+      *out++ = Char(')');
+      return out;
     }
-  };
+    out = detail::write<Char>(out, c.imag(), specs, ctx.locale());
+    if (!detail::isfinite(c.imag())) *out++ = Char(' ');
+    *out++ = Char('i');
+    return out;
+  }
 
  public:
+  FMT_CONSTEXPR auto parse(basic_format_parse_context<Char>& ctx)
+      -> decltype(ctx.begin()) {
+    if (ctx.begin() == ctx.end() || *ctx.begin() == '}') return ctx.begin();
+    return parse_format_specs(ctx.begin(), ctx.end(), specs_, ctx,
+                              detail::type_constant<T, Char>::value);
+  }
+
   template <typename FormatContext>
-  auto format(const std::complex<F>& c, FormatContext& ctx) const
+  auto format(const std::complex<T>& c, FormatContext& ctx) const
       -> decltype(ctx.out()) {
-    return this->write_padded(ctx, writer{this, c});
+    auto specs = specs_;
+    if (specs.width_ref.kind != detail::arg_id_kind::none ||
+        specs.precision_ref.kind != detail::arg_id_kind::none) {
+      detail::handle_dynamic_spec<detail::width_checker>(specs.width,
+                                                         specs.width_ref, ctx);
+      detail::handle_dynamic_spec<detail::precision_checker>(
+          specs.precision, specs.precision_ref, ctx);
+    }
+
+    if (specs.width == 0) return do_format(c, specs, ctx, ctx.out());
+    auto buf = basic_memory_buffer<Char>();
+
+    auto outer_specs = format_specs();
+    outer_specs.width = specs.width;
+    outer_specs.fill = specs.fill;
+    outer_specs.align = specs.align;
+
+    specs.width = 0;
+    specs.fill = {};
+    specs.align = align::none;
+
+    do_format(c, specs, ctx, basic_appender<Char>(buf));
+    return detail::write<Char>(ctx.out(),
+                               basic_string_view<Char>(buf.data(), buf.size()),
+                               outer_specs);
   }
 };
 
diff --git a/test/std-test.cc b/test/std-test.cc
index 17c95cfe57ed2..1327fe8279fd2 100644
--- a/test/std-test.cc
+++ b/test/std-test.cc
@@ -66,10 +66,34 @@ TEST(std_test, thread_id) {
 }
 
 TEST(std_test, complex) {
+  using limits = std::numeric_limits<double>;
+  EXPECT_EQ(fmt::format("{}", std::complex<double>(1, limits::quiet_NaN())),
+            "(1+nan i)");
+  EXPECT_EQ(fmt::format("{}", std::complex<double>(1, -limits::infinity())),
+            "(1-inf i)");
+
+  EXPECT_EQ(fmt::format("{}", std::complex<int>(1, 2)), "(1+2i)");
+
   EXPECT_EQ(fmt::format("{}", std::complex<double>(1, 2.2)), "(1+2.2i)");
+  EXPECT_EQ(fmt::format("{}", std::complex<double>(1, -2.2)), "(1-2.2i)");
   EXPECT_EQ(fmt::format("{}", std::complex<double>(0, 2.2)), "2.2i");
+  EXPECT_EQ(fmt::format("{}", std::complex<double>(0, -2.2)), "-2.2i");
+
+  EXPECT_EQ(fmt::format("{:+}", std::complex<double>(0, 2.2)), "+2.2i");
+  EXPECT_EQ(fmt::format("{:+}", std::complex<double>(0, -2.2)), "-2.2i");
+  EXPECT_EQ(fmt::format("{:+}", std::complex<double>(1, -2.2)), "(+1-2.2i)");
+  EXPECT_EQ(fmt::format("{:+}", std::complex<double>(1, 2.2)), "(+1+2.2i)");
+  EXPECT_EQ(fmt::format("{: }", std::complex<double>(1, 2.2)), "( 1+2.2i)");
+  EXPECT_EQ(fmt::format("{: }", std::complex<double>(1, -2.2)), "( 1-2.2i)");
+
   EXPECT_EQ(fmt::format("{:>20.2f}", std::complex<double>(1, 2.2)),
             "        (1.00+2.20i)");
+  EXPECT_EQ(fmt::format("{:<20.2f}", std::complex<double>(1, 2.2)),
+            "(1.00+2.20i)        ");
+  EXPECT_EQ(fmt::format("{:<20.2f}", std::complex<double>(1, -2.2)),
+            "(1.00-2.20i)        ");
+  EXPECT_EQ(fmt::format("{:<{}.{}f}", std::complex<double>(1, -2.2), 20, 2),
+            "(1.00-2.20i)        ");
 }
 
 #ifdef __cpp_lib_source_location