YT-22473: Enable for-each formatting in known ranges

5a3405e64b730f0056e381af07658d6c2edcb92b

5 files changed, 197 insertions, 17 deletions

diff --git a/library/cpp/yt/string/format-inl.h b/library/cpp/yt/string/format-inl.h
index 4eaa70bef2..67c8affa6e 100644
--- a/library/cpp/yt/string/format-inl.h
+++ b/library/cpp/yt/string/format-inl.h
@@ -222,24 +222,67 @@ void FormatPointerValue(
 // so that someone doesn't accidentally implement the
 // "SimpleRange" concept and have a non-trivial
 // formatting procedure at the same time.
+// Sadly, clang is bugged and thus we must do implementation by hand
+// if we want to use this concept in class specializations.
 
 template <class R>
-concept CKnownRange =
-    requires (R r) { [] <class... Ts> (std::vector<Ts...>) { } (r); } ||
-    requires (R r) { [] <class T, size_t E> (std::span<T, E>) { } (r); } ||
-    requires (R r) { [] <class T, size_t N> (TCompactVector<T, N>) { } (r); } ||
-    requires (R r) { [] <class... Ts> (std::set<Ts...>) { } (r); } ||
-    requires (R r) { [] <class... Ts> (THashSet<Ts...>) { } (r); } ||
-    requires (R r) { [] <class... Ts> (THashMultiSet<Ts...>) { } (r); };
+constexpr bool CKnownRange = false;
+
+template <class T>
+    requires requires (T* t) { [] <class... Ts> (std::vector<Ts...>*) {} (t); }
+constexpr bool CKnownRange<T> = true;
+template <class T, size_t E>
+constexpr bool CKnownRange<std::span<T, E>> = true;
+template <class T, size_t N>
+constexpr bool CKnownRange<std::array<T, N>> = true;
+template <class T, size_t N>
+constexpr bool CKnownRange<TCompactVector<T, N>> = true;
+template <class T>
+    requires requires (T* t) { [] <class... Ts> (std::set<Ts...>*) {} (t); }
+constexpr bool CKnownRange<T> = true;
+template <class T>
+    requires requires (T* t) { [] <class... Ts> (std::multiset<Ts...>*) {} (t); }
+constexpr bool CKnownRange<T> = true;
+template <class... Ts>
+constexpr bool CKnownRange<THashSet<Ts...>> = true;
+template <class... Ts>
+constexpr bool CKnownRange<THashMultiSet<Ts...>> = true;
 
 ////////////////////////////////////////////////////////////////////////////////
 
 template <class R>
-concept CKnownKVRange =
-    requires (R r) { [] <class... Ts> (std::map<Ts...>) { } (r); } ||
-    requires (R r) { [] <class... Ts> (std::multimap<Ts...>) { } (r); } ||
-    requires (R r) { [] <class... Ts> (THashMap<Ts...>) { } (r); } ||
-    requires (R r) { [] <class... Ts> (THashMultiMap<Ts...>) { } (r); };
+constexpr bool CKnownKVRange = false;
+
+template <class T>
+    requires requires (T* t) { [] <class... Ts> (std::map<Ts...>*) {} (t); }
+constexpr bool CKnownKVRange<T> = true;
+template <class T>
+    requires requires (T* t) { [] <class... Ts> (std::multimap<Ts...>*) {} (t); }
+constexpr bool CKnownKVRange<T> = true;
+template <class... Ts>
+constexpr bool CKnownKVRange<THashMap<Ts...>> = true;
+template <class... Ts>
+constexpr bool CKnownKVRange<THashMultiMap<Ts...>> = true;
+
+// TODO(arkady-e1ppa): Uncomment me when
+// https://github.com/llvm/llvm-project/issues/58534 is shipped.
+// template <class R>
+// concept CKnownRange =
+//     requires (R r) { [] <class... Ts> (std::vector<Ts...>) { } (r); } ||
+//     requires (R r) { [] <class T, size_t E> (std::span<T, E>) { } (r); } ||
+//     requires (R r) { [] <class T, size_t N> (TCompactVector<T, N>) { } (r); } ||
+//     requires (R r) { [] <class... Ts> (std::set<Ts...>) { } (r); } ||
+//     requires (R r) { [] <class... Ts> (THashSet<Ts...>) { } (r); } ||
+//     requires (R r) { [] <class... Ts> (THashMultiSet<Ts...>) { } (r); };
+
+// ////////////////////////////////////////////////////////////////////////////////
+
+// template <class R>
+// concept CKnownKVRange =
+//     requires (R r) { [] <class... Ts> (std::map<Ts...>) { } (r); } ||
+//     requires (R r) { [] <class... Ts> (std::multimap<Ts...>) { } (r); } ||
+//     requires (R r) { [] <class... Ts> (THashMap<Ts...>) { } (r); } ||
+//     requires (R r) { [] <class... Ts> (THashMultiMap<Ts...>) { } (r); };
 
 } // namespace NDetail
 
@@ -300,6 +343,22 @@ concept CFormattableKVRange =
     CFormattable<typename R::key_type> &&
     CFormattable<typename R::value_type>;
 
+
+////////////////////////////////////////////////////////////////////////////////
+
+// Specializations of TFormatArg for ranges
+template <class R>
+    requires CFormattableRange<std::remove_cvref_t<R>>
+struct TFormatArg<R>
+    : public TFormatArgBase
+{
+    using TUnderlying = typename std::remove_cvref_t<R>::value_type;
+
+    static constexpr auto ConversionSpecifiers = TFormatArg<TUnderlying>::ConversionSpecifiers;
+
+    static constexpr auto FlagSpecifiers = TFormatArg<TUnderlying>::FlagSpecifiers;
+};
+
 ////////////////////////////////////////////////////////////////////////////////
 
 template <class TRange, class TFormatter>
@@ -740,9 +799,9 @@ void FormatValue(TStringBuilderBase* builder, const TEnumIndexedArray<E, T>& col
 
 // One-valued ranges
 template <CFormattableRange TRange>
-void FormatValue(TStringBuilderBase* builder, const TRange& collection, TStringBuf /*spec*/)
+void FormatValue(TStringBuilderBase* builder, const TRange& collection, TStringBuf spec)
 {
-    NYT::FormatRange(builder, collection, TDefaultFormatter());
+    NYT::FormatRange(builder, collection, TSpecBoundFormatter(spec));
 }
 
 // Two-valued ranges
diff --git a/library/cpp/yt/string/format_analyser.h b/library/cpp/yt/string/format_analyser.h
index 9a05e3eb8c..3dd1ba4cd3 100644
--- a/library/cpp/yt/string/format_analyser.h
+++ b/library/cpp/yt/string/format_analyser.h
@@ -15,7 +15,7 @@ public:
     template <class... TArgs>
     static consteval void ValidateFormat(std::string_view fmt);
 
-private:
+public:
     // Non-constexpr function call will terminate compilation.
     // Purposefully undefined and non-constexpr/consteval
     template <class T>
diff --git a/library/cpp/yt/string/format_string.h b/library/cpp/yt/string/format_string.h
index 0c34763773..ebef9d7042 100644
--- a/library/cpp/yt/string/format_string.h
+++ b/library/cpp/yt/string/format_string.h
@@ -27,7 +27,7 @@ private:
 
 // This class used to properly bind to
 // string literals and allow compile-time parsing/checking
-// of those. If you need a runtime format, use TRuntimeFormat
+// of those. If you need a runtime format, use TRuntimeFormat.
 template <class... TArgs>
 class TBasicFormatString
 {
@@ -50,7 +50,7 @@ private:
     static void CrashCompilerClassIsNotFormattable();
 };
 
-// Used to properly infer template arguments if Format.
+// Used to properly infer template arguments in Format.
 template <class... TArgs>
 using TFormatString = TBasicFormatString<std::type_identity_t<TArgs>...>;
 
diff --git a/library/cpp/yt/string/string.h b/library/cpp/yt/string/string.h
index 8257c1a5ea..b91d1b3805 100644
--- a/library/cpp/yt/string/string.h
+++ b/library/cpp/yt/string/string.h
@@ -27,6 +27,25 @@ struct TDefaultFormatter
     }
 };
 
+//! Bind spec to a formatter.
+//! Used in ranges processing.
+class TSpecBoundFormatter
+{
+public:
+    explicit TSpecBoundFormatter(TStringBuf spec)
+        : Spec_(spec)
+    { }
+
+    template <class T>
+    void operator()(TStringBuilderBase* builder, const T& obj) const
+    {
+        FormatValue(builder, obj, Spec_);
+    }
+
+private:
+    TStringBuf Spec_;
+};
+
 static constexpr TStringBuf DefaultJoinToStringDelimiter = ", ";
 static constexpr TStringBuf DefaultKeyValueDelimiter = ": ";
 static constexpr TStringBuf DefaultRangeEllipsisFormat = "...";
diff --git a/library/cpp/yt/string/unittests/format_ut.cpp b/library/cpp/yt/string/unittests/format_ut.cpp
index 4083338c44..aa820fd5af 100644
--- a/library/cpp/yt/string/unittests/format_ut.cpp
+++ b/library/cpp/yt/string/unittests/format_ut.cpp
@@ -13,6 +13,44 @@ namespace {
 
 ////////////////////////////////////////////////////////////////////////////////
 
+struct TWithCustomFlags
+{
+    [[maybe_unused]]
+    friend void FormatValue(TStringBuilderBase* builder, const TWithCustomFlags&, TStringBuf spec)
+    {
+        if (spec.Contains('R')) {
+            builder->AppendString("R");
+        }
+        if (spec.Contains('N')) {
+            builder->AppendString("N");
+        }
+
+        builder->AppendString("P");
+    }
+};
+
+////////////////////////////////////////////////////////////////////////////////
+
+} // namespace
+
+////////////////////////////////////////////////////////////////////////////////
+
+template <>
+struct TFormatArg<TWithCustomFlags>
+{
+    [[maybe_unused]] static constexpr std::array ConversionSpecifiers = {
+        'v',
+    };
+
+    [[maybe_unused]] static constexpr std::array FlagSpecifiers = {
+        'R', 'N',
+    };
+};
+
+////////////////////////////////////////////////////////////////////////////////
+
+namespace {
+
 // Some compile-time sanity checks.
 static_assert(CFormattable<int>);
 static_assert(CFormattable<double>);
@@ -22,6 +60,7 @@ static_assert(CFormattable<TStringBuf>);
 static_assert(CFormattable<TString>);
 static_assert(CFormattable<std::span<int>>);
 static_assert(CFormattable<std::vector<int>>);
+static_assert(CFormattable<std::array<int, 5>>);
 
 // N.B. TCompactVector<int, 1> is not buildable on Windows
 static_assert(CFormattable<TCompactVector<int, 2>>);
@@ -40,6 +79,8 @@ struct TUnformattable
 { };
 static_assert(!CFormattable<TUnformattable>);
 
+static_assert(CFormattable<TWithCustomFlags>);
+
 ////////////////////////////////////////////////////////////////////////////////
 
 TEST(TFormatTest, Nothing)
@@ -223,6 +264,67 @@ TEST(TFormatTest, RuntimeFormat)
     EXPECT_EQ(Format(TRuntimeFormat(format), "World"), "Hello World");
 }
 
+TEST(TFormatTest, CustomFlagsSimple)
+{
+    EXPECT_EQ(Format("%v", TWithCustomFlags{}), TString("P"));
+    EXPECT_EQ(Format("%Rv", TWithCustomFlags{}), TString("RP"));
+    EXPECT_EQ(Format("%Nv", TWithCustomFlags{}), TString("NP"));
+    EXPECT_EQ(Format("%RNv", TWithCustomFlags{}), TString("RNP"));
+    EXPECT_EQ(Format("%NRv", TWithCustomFlags{}), TString("RNP"));
+}
+
+TEST(TFormatTest, CustomFlagsCollection)
+{
+    constexpr int elementCount = 5;
+    auto toCollection = [] (TString pattern) {
+        TString ret = "[";
+        for (int i = 0; i < elementCount - 1; ++i) {
+            ret += pattern + ", ";
+        }
+
+        ret += pattern + "]";
+
+        return ret;
+    };
+
+    std::vector vec(elementCount, TWithCustomFlags{});
+
+    EXPECT_EQ(Format("%v", vec), toCollection("P"));
+    EXPECT_EQ(Format("%Rv", vec), toCollection("RP"));
+    EXPECT_EQ(Format("%Nv", vec), toCollection("NP"));
+    EXPECT_EQ(Format("%RNv", vec), toCollection("RNP"));
+    EXPECT_EQ(Format("%NRv", vec), toCollection("RNP"));
+}
+
+TEST(TFormatTest, CustomFlagsCollectionTwoLevels)
+{
+    constexpr int elementCount1 = 5;
+    constexpr int elementCount2 = 3;
+    auto toCollection = [] (int count, TString pattern) {
+        TString ret = "[";
+        for (int i = 0; i < count - 1; ++i) {
+            ret += pattern + ", ";
+        }
+
+        ret += pattern + "]";
+
+        return ret;
+    };
+    auto toCollectionD2 = [&] (TString pattern) {
+        return toCollection(elementCount2, toCollection(elementCount1, std::move(pattern)));
+    };
+
+    std::vector vec(elementCount1, TWithCustomFlags{});
+    std::array<decltype(vec), elementCount2> arr;
+    std::ranges::fill(arr, vec);
+
+    EXPECT_EQ(Format("%v", arr), toCollectionD2("P"));
+    EXPECT_EQ(Format("%Rv", arr), toCollectionD2("RP"));
+    EXPECT_EQ(Format("%Nv", arr), toCollectionD2("NP"));
+    EXPECT_EQ(Format("%RNv", arr), toCollectionD2("RNP"));
+    EXPECT_EQ(Format("%NRv", arr), toCollectionD2("RNP"));
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 
 } // namespace