Intermediate changes

5 files changed, 159 insertions, 78 deletions

diff --git a/yt/yt/core/concurrency/unittests/profiled_fair_share_invoker_pool_ut.cpp b/yt/yt/core/concurrency/unittests/profiled_fair_share_invoker_pool_ut.cpp
index d1b3ab09d5..6aab403a3e 100644
--- a/yt/yt/core/concurrency/unittests/profiled_fair_share_invoker_pool_ut.cpp
+++ b/yt/yt/core/concurrency/unittests/profiled_fair_share_invoker_pool_ut.cpp
@@ -14,6 +14,8 @@
 
 #include <yt/yt/library/profiling/solomon/exporter.h>
 
+#include <library/cpp/json/yson/json2yson.h>
+
 #include <util/datetime/base.h>
 
 #include <algorithm>
@@ -666,15 +668,7 @@ public:
 
     auto GetSensors(TString json)
     {
-        for (auto& c : json) {
-            if (c == ':') {
-                c = '=';
-            } else if (c == ',') {
-                c = ';';
-            }
-        }
-
-        auto yson = NYson::TYsonString(json);
+        auto yson = NYson::TYsonString(NJson2Yson::SerializeJsonValueAsYson(NJson::ReadJsonFastTree(json)));
 
         auto list = NYTree::ConvertToNode(yson)->AsMap()->FindChild("sensors");
 
@@ -730,7 +724,7 @@ public:
 
         THashMap<TString, int> invokerNameToDequeued = invokerNameToEnqueued;
 
-        for (const auto& entry : GetSensors(json)) {
+        for (const auto& entry : GetSensors(std::move(json))) {
             auto mapEntry = entry->AsMap();
             auto labels = mapEntry->FindChild("labels")->AsMap();
 
diff --git a/yt/yt/core/concurrency/unittests/ya.make b/yt/yt/core/concurrency/unittests/ya.make
index ae8a5c6f99..b8b94dcfc9 100644
--- a/yt/yt/core/concurrency/unittests/ya.make
+++ b/yt/yt/core/concurrency/unittests/ya.make
@@ -48,6 +48,8 @@ INCLUDE(${ARCADIA_ROOT}/yt/opensource.inc)
 PEERDIR(
     yt/yt/core
     yt/yt/core/test_framework
+
+    library/cpp/json/yson
 )
 
 REQUIREMENTS(
diff --git a/yt/yt/core/ytree/polymorphic_yson_struct-inl.h b/yt/yt/core/ytree/polymorphic_yson_struct-inl.h
index ef8625df47..de5b3ab7ef 100644
--- a/yt/yt/core/ytree/polymorphic_yson_struct-inl.h
+++ b/yt/yt/core/ytree/polymorphic_yson_struct-inl.h
@@ -10,42 +10,33 @@ namespace NYT::NYTree {
 
 namespace NDetail {
 
-template <class TEnum>
-    requires TEnumTraits<TEnum>::IsEnum
-template <class TBase, class TDerived>
-TIntrusivePtr<TBase> TEnumTraitsExt<TEnum>::ConcreteFactory()
+template <class TEnum, TEnum Value, class TBase, class TDerived>
+TIntrusivePtr<TBase> TMappingLeaf<TEnum, Value, TBase, TDerived>::CreateInstance()
 {
     return New<TDerived>();
 }
 
-template <class TEnum>
-    requires TEnumTraits<TEnum>::IsEnum
-template <class TBase, class... TDerived>
-TInstanceFactory<TEnum, TBase>  TEnumTraitsExt<TEnum>::MakeFactory()
+template <class TEnum, TEnum BaseValue, CYsonStructDerived TBase, TEnum... Values, class... TDerived>
+    requires (CHierarchy<TBase, TDerived...>)
+TIntrusivePtr<TBase>
+TPolymorphicMapping<TEnum, TLeafTag<BaseValue, TBase>, TLeafTag<Values, TDerived>...>::
+CreateInstance(TEnum value)
 {
-    static constexpr auto keys = TTraits::GetDomainValues();
-    using TTuple = std::tuple<TBase, TDerived...>;
+    if (value == BaseValue) {
+        return TLeaf<BaseValue, TBase>::CreateInstance();
+    }
 
-    TInstanceFactory<TEnum, TBase> mapping;
+    TIntrusivePtr<TBase> ret;
 
-    [&] <size_t... Idx> (std::index_sequence<Idx...>) {
-        ([&] {
-            mapping[keys[Idx]] = &TEnumTraitsExt<TEnum>::ConcreteFactory<TBase, std::tuple_element_t<Idx, TTuple>>;
-        } (), ...);
-    } (std::make_index_sequence<sizeof...(TDerived) + 1>());
+    ([&ret, value] {
+        if (value == Values) {
+            ret = TLeaf<Values, TDerived>::CreateInstance();
+            return false;
+        }
+        return true;
+    } () && ...);
 
-    return mapping;
-}
-
-////////////////////////////////////////////////////////////////////////////////
-
-template <class TEnum, class TB, class... TD>
-TIntrusivePtr<TB> TPolymorphicEnumMapping<TEnum, TB, TD...>::MakeInstance(TEnum e)
-{
-    static auto factory =
-        TEnumTraitsExt<TEnum>::template MakeFactory<TB, TD...>();
-
-    return factory[e]();
+    return ret;
 }
 
 } // namespace NDetail
@@ -53,6 +44,11 @@ TIntrusivePtr<TB> TPolymorphicEnumMapping<TEnum, TB, TD...>::MakeInstance(TEnum
 ////////////////////////////////////////////////////////////////////////////////
 
 template <CPolymorphicEnumMapping TMapping>
+TPolymorphicYsonStruct<TMapping>::TPolymorphicYsonStruct(TKey key)
+    : TPolymorphicYsonStruct(key, TMapping::CreateInstance(key))
+{ }
+
+template <CPolymorphicEnumMapping TMapping>
 TPolymorphicYsonStruct<TMapping>::TPolymorphicYsonStruct(TKey key, TIntrusivePtr<TBase> ptr) noexcept
     : Storage_(std::move(ptr))
     , HeldType_(key)
@@ -77,6 +73,11 @@ void TPolymorphicYsonStruct<TMapping>::Load(
     // parse (unless we want to slice which we don't).
     IMapNodePtr map = TTraits::AsNode(source)->AsMap();
 
+    if (!map || map->GetChildCount() == 0) {
+        // Empty struct.
+        return;
+    }
+
     auto key = map->FindChildValue<TKey>("type");
     THROW_ERROR_EXCEPTION_UNLESS(
         key.has_value(),
@@ -86,7 +87,7 @@ void TPolymorphicYsonStruct<TMapping>::Load(
     if (!Storage_ || HeldType_ != type) {
         // NB: We will try to merge configs if types match.
         HeldType_ = type;
-        Storage_ = TMapping::MakeInstance(HeldType_);
+        Storage_ = TMapping::CreateInstance(HeldType_);
     }
 
     if (recursiveUnrecognizedStrategy) {
@@ -97,6 +98,10 @@ void TPolymorphicYsonStruct<TMapping>::Load(
     // therefore we must delete it prior to |Load| call.
     map->RemoveChild("type");
     Storage_->Load(map, postprocess, setDefaults, path);
+
+    // NB(arkady-e1ppa): We must not actually remove contents of the node as a postcondition
+    // since it mutates serialized data which might be used for config validation.
+    map->AddChild("type", ConvertToNode(HeldType_));
 }
 
 template <CPolymorphicEnumMapping TMapping>
@@ -104,26 +109,51 @@ void TPolymorphicYsonStruct<TMapping>::Save(NYson::IYsonConsumer* consumer) cons
 {
     consumer->OnBeginMap();
 
-    consumer->OnKeyedItem("type");
-    consumer->OnStringScalar(FormatEnum(HeldType_));
+    if (Storage_) {
+        consumer->OnKeyedItem("type");
+        consumer->OnStringScalar(FormatEnum(HeldType_));
+
+        Storage_->SaveAsMapFragment(consumer);
+    }
 
-    Storage_->SaveAsMapFragment(consumer);
     consumer->OnEndMap();
 }
 
 template <CPolymorphicEnumMapping TMapping>
-template <std::derived_from<typename TMapping::TBase> TConcrete>
+template <std::derived_from<typename TMapping::TBaseClass> TConcrete>
 TIntrusivePtr<TConcrete> TPolymorphicYsonStruct<TMapping>::TryGetConcrete() const
 {
     return DynamicPointerCast<TConcrete>(Storage_);
 }
 
 template <CPolymorphicEnumMapping TMapping>
+template <typename TMapping::TKey Value>
+TIntrusivePtr<typename TMapping::template TDerivedToEnum<Value>> TPolymorphicYsonStruct<TMapping>::TryGetConcrete() const
+{
+    if (Value != HeldType_) {
+        return {};
+    }
+    return TryGetConcrete<typename TMapping::template TDerivedToEnum<Value>>();
+}
+
+template <CPolymorphicEnumMapping TMapping>
 typename TPolymorphicYsonStruct<TMapping>::TKey TPolymorphicYsonStruct<TMapping>::GetCurrentType() const
 {
     return HeldType_;
 }
 
+template <CPolymorphicEnumMapping TMapping>
+typename TPolymorphicYsonStruct<TMapping>::TBase* TPolymorphicYsonStruct<TMapping>::operator->()
+{
+    return Storage_.Get();
+}
+
+template <CPolymorphicEnumMapping TMapping>
+const typename TPolymorphicYsonStruct<TMapping>::TBase* TPolymorphicYsonStruct<TMapping>::operator->() const
+{
+    return Storage_.Get();
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 
 template <CPolymorphicEnumMapping TMapping>
@@ -141,6 +171,7 @@ void Deserialize(TPolymorphicYsonStruct<TMapping>& value, TSource source)
 ////////////////////////////////////////////////////////////////////////////////
 
 #undef DEFINE_POLYMORPHIC_YSON_STRUCT
+#undef DEFINE_POLYMORPHIC_YSON_STRUCT_FOR_ENUM
 
 #define POLYMORPHIC_YSON_STRUCT_IMPL__GET_ENUM_SEQ_ELEM(item) \
     PP_LEFT_PARENTHESIS PP_ELEMENT(item, 0) PP_RIGHT_PARENTHESIS
@@ -151,23 +182,37 @@ void Deserialize(TPolymorphicYsonStruct<TMapping>& value, TSource source)
 #define POLYMORPHIC_YSON_STRUCT_IMPL__ENUM_NAME(Struct) \
     E##Struct##Type
 
-#define POLYMORPHIC_YSON_STRUCT_IMPL__MAKE_MAPPING_ENUM(Struct, seq) \
+#define POLYMORPHIC_YSON_STRUCT_IMPL__MAKE_MAPPING_ENUM(seq) \
     DEFINE_ENUM(EType, POLYMORPHIC_YSON_STRUCT_IMPL__GET_ENUM_SEQ(seq))
 
+#define POLYMORPHIC_YSON_STRUCT_IMPL__MAKE_MAPPING_ENUM_ALIAS(EnumName) \
+    using EType = EnumName;
+
 #define POLYMORPHIC_YSON_STRUCT_IMPL__GET_CLASS_ELEM(item) \
     PP_COMMA() PP_ELEMENT(item, 1)
 
+#define POLYMORPHIC_YSON_STRUCT_IMPL__MAKE_LEAF_FROM_ETYPE(item) \
+    PP_COMMA() ::NYT::NYTree::NDetail::TLeafTag<EType:: PP_ELEMENT(item, 0), PP_ELEMENT(item, 1)>
+
 #define POLYMORPHIC_YSON_STRUCT_IMPL__MAKE_MAPPING_CLASS(Struct, seq) \
-    using TMapping = TPolymorphicEnumMapping<EType PP_FOR_EACH(POLYMORPHIC_YSON_STRUCT_IMPL__GET_CLASS_ELEM, seq)>
+    using TMapping = ::NYT::NYTree::TPolymorphicEnumMapping<EType PP_FOR_EACH(POLYMORPHIC_YSON_STRUCT_IMPL__MAKE_LEAF_FROM_ETYPE, seq)>
 
 #define DEFINE_POLYMORPHIC_YSON_STRUCT(name, seq) \
 namespace NPolymorphicYsonStructFor##name { \
  \
-    POLYMORPHIC_YSON_STRUCT_IMPL__MAKE_MAPPING_ENUM(name, seq); \
+    POLYMORPHIC_YSON_STRUCT_IMPL__MAKE_MAPPING_ENUM(seq); \
     POLYMORPHIC_YSON_STRUCT_IMPL__MAKE_MAPPING_CLASS(name, seq); \
 } /*NPolymorphicYsonStructFor##name*/ \
 using POLYMORPHIC_YSON_STRUCT_IMPL__ENUM_NAME(name) = NPolymorphicYsonStructFor##name::EType; \
-using T##name = TPolymorphicYsonStruct<NPolymorphicYsonStructFor##name::TMapping>; \
+using T##name = ::NYT::NYTree::TPolymorphicYsonStruct<NPolymorphicYsonStructFor##name::TMapping>; \
+static_assert(true)
+
+#define DEFINE_POLYMORPHIC_YSON_STRUCT_FOR_ENUM(name, enum, seq) \
+namespace NPolymorphicYsonStructFor##name { \
+    POLYMORPHIC_YSON_STRUCT_IMPL__MAKE_MAPPING_ENUM_ALIAS(enum); \
+    POLYMORPHIC_YSON_STRUCT_IMPL__MAKE_MAPPING_CLASS(name, seq); \
+} /*NPolymorphicYsonStructFor##name*/ \
+using T##name = ::NYT::NYTree::TPolymorphicYsonStruct<NPolymorphicYsonStructFor##name::TMapping>; \
 static_assert(true)
 
 ////////////////////////////////////////////////////////////////////////////////
diff --git a/yt/yt/core/ytree/polymorphic_yson_struct.h b/yt/yt/core/ytree/polymorphic_yson_struct.h
index 0f347abb4b..64566fce4a 100644
--- a/yt/yt/core/ytree/polymorphic_yson_struct.h
+++ b/yt/yt/core/ytree/polymorphic_yson_struct.h
@@ -42,47 +42,49 @@ constexpr bool CHierarchy =
 
 ////////////////////////////////////////////////////////////////////////////////
 
-template <CYsonStructDerived TBase>
-using TInstanceFactoryLeaf =
-    TIntrusivePtr<TBase>(*)();
+template <auto Value>
+struct TFactoryTag
+{ };
 
-template <class TEnum, CYsonStructDerived TBase>
-using TInstanceFactory = TEnumIndexedArray<TEnum, TInstanceFactoryLeaf<TBase>>;
+template <auto Value, class TDerived>
+struct TLeafTag
+{ };
 
-template <class TEnum>
-    requires TEnumTraits<TEnum>::IsEnum
-struct TEnumTraitsExt
+template <class TEnum, TEnum Value, class TBase, class TDerived>
+struct TMappingLeaf
 {
-    using TTraits = TEnumTraits<TEnum>;
+    static TIntrusivePtr<TBase> CreateInstance();
 
-    template <class TBase, class... TDerived>
-    static constexpr bool CompatibleHierarchy =
-        CHierarchy<TBase, TDerived...> &&
-        (TTraits::GetDomainSize() == sizeof...(TDerived) + 1);
-
-    template <class TBase, class TDerived>
-    static TIntrusivePtr<TBase> ConcreteFactory();
-
-    template <class TBase, class... TDerived>
-    static TInstanceFactory<TEnum, TBase> MakeFactory();
+    // This is intentionally undefined to be used in invoke_result.
+    friend TDerived EnumToDerivedMethod(TFactoryTag<Value>, TMappingLeaf*);
 };
 
-////////////////////////////////////////////////////////////////////////////////
+template <class TEnum, class... TLeafTags>
+struct TPolymorphicMapping;
 
-template <class TEnum, class TB, class... TD>
-struct TPolymorphicEnumMapping
+template <class TEnum, TEnum BaseValue, CYsonStructDerived TBase, TEnum... Values, class... TDerived>
+    requires (CHierarchy<TBase, TDerived...>)
+struct TPolymorphicMapping<TEnum, TLeafTag<BaseValue, TBase>, TLeafTag<Values, TDerived>...>
+    : public TMappingLeaf<TEnum, BaseValue, TBase, TBase>
+    , public TMappingLeaf<TEnum, Values, TBase, TDerived>...
 {
+    template <TEnum Value, class TConcrete>
+    using TLeaf = TMappingLeaf<TEnum, Value, TBase, TConcrete>;
+
+    template <TEnum Value>
+    using TDerivedToEnum = decltype(EnumToDerivedMethod(TFactoryTag<Value>{}, std::declval<TPolymorphicMapping*>()));
+
     using TKey = TEnum;
-    using TBase = TB;
-    using TDerived = std::tuple<TD...>;
-    using THierarchy = std::tuple<TB, TD...>;
+    using TBaseClass = TBase;
 
-    static TIntrusivePtr<TB> MakeInstance(TEnum e);
+    static TIntrusivePtr<TBase> CreateInstance(TEnum value);
 };
 
+////////////////////////////////////////////////////////////////////////////////
+
 template <class T>
 constexpr bool IsMapping = requires (T t) {
-    [] <class E, class B, class... D> (TPolymorphicEnumMapping<E, B, D...>) {
+    [] <class TEnum, class... TLeafTags> (TPolymorphicMapping<TEnum, TLeafTags...>) {
     } (t);
 };
 
@@ -95,9 +97,8 @@ constexpr bool IsMapping = requires (T t) {
 template <class TBase, class... TDerived>
 concept CHierarchy = NDetail::CHierarchy<TBase, TDerived...>;
 
-template <class TEnum, CYsonStructDerived TBase, class... TDerived>
-    requires CHierarchy<TBase, TDerived...>
-using TPolymorphicEnumMapping = NDetail::TPolymorphicEnumMapping<TEnum, TBase, TDerived...>;
+template <class TEnum, class... TLeafTags>
+using TPolymorphicEnumMapping = NDetail::TPolymorphicMapping<TEnum, TLeafTags...>;
 
 template <class T>
 concept CPolymorphicEnumMapping = NDetail::IsMapping<T>;
@@ -137,13 +138,14 @@ class TPolymorphicYsonStruct
     // TODO(arkady-e1ppa): Support ctor from anyone from the
     // hierarchy.
     using TKey = typename TMapping::TKey;
-    using TBase = typename TMapping::TBase;
+    using TBase = typename TMapping::TBaseClass;
 
 public:
     using TImplementsYsonStructField = void;
 
     TPolymorphicYsonStruct() = default;
 
+    explicit TPolymorphicYsonStruct(TKey key);
     TPolymorphicYsonStruct(TKey key, TIntrusivePtr<TBase> ptr) noexcept;
 
     template <CYsonStructSource TSource>
@@ -160,8 +162,14 @@ public:
     template <std::derived_from<TBase> TConcrete>
     TIntrusivePtr<TConcrete> TryGetConcrete() const;
 
+    template <TKey Value>
+    TIntrusivePtr<typename TMapping::template TDerivedToEnum<Value>> TryGetConcrete() const;
+
     TKey GetCurrentType() const;
 
+    TBase* operator->();
+    const TBase* operator->() const;
+
 private:
     TIntrusivePtr<TBase> Storage_;
     TKey HeldType_;
@@ -196,6 +204,29 @@ void Deserialize(TPolymorphicYsonStruct<TMapping>& value, TSource source);
 */
 #define DEFINE_POLYMORPHIC_YSON_STRUCT(name, seq)
 
+//! Usage:
+/*
+    DEFINE_ENUM(EMyEnum
+        (Pear)
+        (Apple)
+    );
+
+    DEFINE_POLYMORPHIC_YSON_STRUCT_FOR_ENUM(Struct, EMyEnum,
+        ((Pear)  (TPearClass))
+        ((Apple) (TAppleClass))
+    )
+
+    // NB(arkady-e1ppa): enum names in the list must be unqualified! E.g.
+
+    DEFINE_POLYMORPHIC_YSON_STRUCT_FOR_ENUM(Struct, EMyEnum,
+        ((EMyEnum::Pear)  (TPearClass))
+        ((EMyEnum::Apple) (TAppleClass))
+    )
+
+    Will not compile
+*/
+#define DEFINE_POLYMORPHIC_YSON_STRUCT_FOR_ENUM(name, enum, seq)
+
 ////////////////////////////////////////////////////////////////////////////////
 
 } // namespace NYT::NYTree
diff --git a/yt/yt/core/ytree/unittests/yson_struct_ut.cpp b/yt/yt/core/ytree/unittests/yson_struct_ut.cpp
index 59d1bd7c10..95ec07ab8c 100644
--- a/yt/yt/core/ytree/unittests/yson_struct_ut.cpp
+++ b/yt/yt/core/ytree/unittests/yson_struct_ut.cpp
@@ -2877,6 +2877,15 @@ struct TPolyHolder
     }
 };
 
+TEST(TYsonStructTest, TestPolymorphicYsonStructSerializeEmpty)
+{
+    TPolyHolder holder;
+
+    auto node = ConvertToNode(holder);
+
+    Deserialize(holder, node->AsMap());
+}
+
 TEST(TYsonStructTest, TestPolymorphicYsonStructAsField)
 {
     TPolyHolder holder;