YT-20462: Introduce TPolymorphicYsonStruct for serialization of polymorphic types

19597a5cbe777b391118536405f7900b779a29c5

6 files changed, 709 insertions, 12 deletions

diff --git a/yt/yt/core/ytree/polymorphic_yson_struct-inl.h b/yt/yt/core/ytree/polymorphic_yson_struct-inl.h
new file mode 100644
index 0000000000..ef8625df47
--- /dev/null
+++ b/yt/yt/core/ytree/polymorphic_yson_struct-inl.h
@@ -0,0 +1,175 @@
+#ifndef POLYMORPHIC_YSON_STRUCT_INL_H_
+#error "Direct inclusion of this file is not allowed, include polymorphic_yson_struct.h"
+// For the sake of sane code completion.
+#include "polymorphic_yson_struct.h"
+#endif
+
+namespace NYT::NYTree {
+
+////////////////////////////////////////////////////////////////////////////////
+
+namespace NDetail {
+
+template <class TEnum>
+    requires TEnumTraits<TEnum>::IsEnum
+template <class TBase, class TDerived>
+TIntrusivePtr<TBase> TEnumTraitsExt<TEnum>::ConcreteFactory()
+{
+    return New<TDerived>();
+}
+
+template <class TEnum>
+    requires TEnumTraits<TEnum>::IsEnum
+template <class TBase, class... TDerived>
+TInstanceFactory<TEnum, TBase>  TEnumTraitsExt<TEnum>::MakeFactory()
+{
+    static constexpr auto keys = TTraits::GetDomainValues();
+    using TTuple = std::tuple<TBase, TDerived...>;
+
+    TInstanceFactory<TEnum, TBase> mapping;
+
+    [&] <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>());
+
+    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]();
+}
+
+} // namespace NDetail
+
+////////////////////////////////////////////////////////////////////////////////
+
+template <CPolymorphicEnumMapping TMapping>
+TPolymorphicYsonStruct<TMapping>::TPolymorphicYsonStruct(TKey key, TIntrusivePtr<TBase> ptr) noexcept
+    : Storage_(std::move(ptr))
+    , HeldType_(key)
+{
+    YT_VERIFY(Storage_);
+}
+
+template <CPolymorphicEnumMapping TMapping>
+template <CYsonStructSource TSource>
+void TPolymorphicYsonStruct<TMapping>::Load(
+    TSource source,
+    bool postprocess,
+    bool setDefaults,
+    const NYPath::TYPath& path,
+    std::optional<EUnrecognizedStrategy> recursiveUnrecognizedStrategy)
+{
+    using TTraits = NPrivate::TYsonSourceTraits<TSource>;
+
+    // TODO(arkady-e1ppa): Support parsing without node
+    // conversion smth? FillMap wont work since we
+    // do not know the value_type of a map we would
+    // parse (unless we want to slice which we don't).
+    IMapNodePtr map = TTraits::AsNode(source)->AsMap();
+
+    auto key = map->FindChildValue<TKey>("type");
+    THROW_ERROR_EXCEPTION_UNLESS(
+        key.has_value(),
+        "Concrete type must be specified! Use \"type\": \"concrete_type\"");
+
+    auto type = *key;
+    if (!Storage_ || HeldType_ != type) {
+        // NB: We will try to merge configs if types match.
+        HeldType_ = type;
+        Storage_ = TMapping::MakeInstance(HeldType_);
+    }
+
+    if (recursiveUnrecognizedStrategy) {
+        Storage_->SetUnrecognizedStrategy(*recursiveUnrecognizedStrategy);
+    }
+
+    // "type" must be unrecognized for the original struct
+    // therefore we must delete it prior to |Load| call.
+    map->RemoveChild("type");
+    Storage_->Load(map, postprocess, setDefaults, path);
+}
+
+template <CPolymorphicEnumMapping TMapping>
+void TPolymorphicYsonStruct<TMapping>::Save(NYson::IYsonConsumer* consumer) const
+{
+    consumer->OnBeginMap();
+
+    consumer->OnKeyedItem("type");
+    consumer->OnStringScalar(FormatEnum(HeldType_));
+
+    Storage_->SaveAsMapFragment(consumer);
+    consumer->OnEndMap();
+}
+
+template <CPolymorphicEnumMapping TMapping>
+template <std::derived_from<typename TMapping::TBase> TConcrete>
+TIntrusivePtr<TConcrete> TPolymorphicYsonStruct<TMapping>::TryGetConcrete() const
+{
+    return DynamicPointerCast<TConcrete>(Storage_);
+}
+
+template <CPolymorphicEnumMapping TMapping>
+typename TPolymorphicYsonStruct<TMapping>::TKey TPolymorphicYsonStruct<TMapping>::GetCurrentType() const
+{
+    return HeldType_;
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+template <CPolymorphicEnumMapping TMapping>
+void Serialize(const TPolymorphicYsonStruct<TMapping>& value, NYson::IYsonConsumer* consumer)
+{
+    value.Save(consumer);
+}
+
+template <CPolymorphicEnumMapping TMapping, CYsonStructSource TSource>
+void Deserialize(TPolymorphicYsonStruct<TMapping>& value, TSource source)
+{
+    value.Load(std::move(source));
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+#undef DEFINE_POLYMORPHIC_YSON_STRUCT
+
+#define POLYMORPHIC_YSON_STRUCT_IMPL__GET_ENUM_SEQ_ELEM(item) \
+    PP_LEFT_PARENTHESIS PP_ELEMENT(item, 0) PP_RIGHT_PARENTHESIS
+
+#define POLYMORPHIC_YSON_STRUCT_IMPL__GET_ENUM_SEQ(seq) \
+    PP_FOR_EACH(POLYMORPHIC_YSON_STRUCT_IMPL__GET_ENUM_SEQ_ELEM, seq)
+
+#define POLYMORPHIC_YSON_STRUCT_IMPL__ENUM_NAME(Struct) \
+    E##Struct##Type
+
+#define POLYMORPHIC_YSON_STRUCT_IMPL__MAKE_MAPPING_ENUM(Struct, seq) \
+    DEFINE_ENUM(EType, POLYMORPHIC_YSON_STRUCT_IMPL__GET_ENUM_SEQ(seq))
+
+#define POLYMORPHIC_YSON_STRUCT_IMPL__GET_CLASS_ELEM(item) \
+    PP_COMMA() 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)>
+
+#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_CLASS(name, seq); \
+} /*NPolymorphicYsonStructFor##name*/ \
+using POLYMORPHIC_YSON_STRUCT_IMPL__ENUM_NAME(name) = NPolymorphicYsonStructFor##name::EType; \
+using T##name = TPolymorphicYsonStruct<NPolymorphicYsonStructFor##name::TMapping>; \
+static_assert(true)
+
+////////////////////////////////////////////////////////////////////////////////
+
+} // namespace NYT::NYTree
diff --git a/yt/yt/core/ytree/polymorphic_yson_struct.h b/yt/yt/core/ytree/polymorphic_yson_struct.h
new file mode 100644
index 0000000000..0f347abb4b
--- /dev/null
+++ b/yt/yt/core/ytree/polymorphic_yson_struct.h
@@ -0,0 +1,205 @@
+#pragma once
+
+#include "yson_struct.h"
+
+namespace NYT::NYTree {
+
+////////////////////////////////////////////////////////////////////////////////
+
+//! Code for mapping generation.
+namespace NDetail {
+
+template <class T, class... TArgs>
+constexpr bool DifferentFrom = (!std::same_as<T, TArgs> && ...);
+
+template <class TTuple>
+struct TAllDifferentImpl;
+
+template <>
+struct TAllDifferentImpl<std::tuple<>>
+{
+    static constexpr bool Value = true;
+};
+
+template <class TCurrent, class... TSuffix>
+struct TAllDifferentImpl<std::tuple<TCurrent, TSuffix...>>
+{
+    static constexpr bool Value =
+        DifferentFrom<TCurrent, TSuffix...> &&
+        TAllDifferentImpl<std::tuple<TSuffix...>>::Value;
+};
+
+template <class... TArgs>
+constexpr bool AllDifferent = TAllDifferentImpl<std::tuple<TArgs...>>::Value;
+
+template <class TBase, class... TArgs>
+constexpr bool AllDerived = (std::derived_from<TArgs, TBase> && ...);
+
+template <class TBase, class... TArgs>
+constexpr bool CHierarchy =
+    AllDerived<TBase, TArgs...> &&
+    AllDifferent<TBase, TArgs...>;
+
+////////////////////////////////////////////////////////////////////////////////
+
+template <CYsonStructDerived TBase>
+using TInstanceFactoryLeaf =
+    TIntrusivePtr<TBase>(*)();
+
+template <class TEnum, CYsonStructDerived TBase>
+using TInstanceFactory = TEnumIndexedArray<TEnum, TInstanceFactoryLeaf<TBase>>;
+
+template <class TEnum>
+    requires TEnumTraits<TEnum>::IsEnum
+struct TEnumTraitsExt
+{
+    using TTraits = TEnumTraits<TEnum>;
+
+    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();
+};
+
+////////////////////////////////////////////////////////////////////////////////
+
+template <class TEnum, class TB, class... TD>
+struct TPolymorphicEnumMapping
+{
+    using TKey = TEnum;
+    using TBase = TB;
+    using TDerived = std::tuple<TD...>;
+    using THierarchy = std::tuple<TB, TD...>;
+
+    static TIntrusivePtr<TB> MakeInstance(TEnum e);
+};
+
+template <class T>
+constexpr bool IsMapping = requires (T t) {
+    [] <class E, class B, class... D> (TPolymorphicEnumMapping<E, B, D...>) {
+    } (t);
+};
+
+////////////////////////////////////////////////////////////////////////////////
+
+} // namespace NDetail
+
+////////////////////////////////////////////////////////////////////////////////
+
+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 T>
+concept CPolymorphicEnumMapping = NDetail::IsMapping<T>;
+
+////////////////////////////////////////////////////////////////////////////////
+
+//! Wrapper class which allows automatic (de-)serialization
+//! of polymorphic types. Without it, one would have
+//! to keep type serialized (e.g. in INodePtr form)
+//! manually add some type marker like enum inside
+//! the base object layout and then do some switches
+//! over enum to properly deserialize the struct.
+//! TPolymorphicYsonStruct does so almost automatically.
+//! "Almost" because it still requires you to define some
+//! enum marker and manually describe the hierarchy during
+//! the type declaration. For most of the cases one can
+//! use DEFINE_POLYMORPHIC_YSON_STRUCT to do that with
+//! minimal effort. Yson layout expected/generated by this struct:
+/*
+    "type"   = "derived2";
+    "field1" = ...;
+    "field2" = ...;
+*/
+// NB(arkady-e1ppa): Word "type" is reserved
+// and must not be used as a field name of any
+// class in the hierarchy.
+// TODO(arkady-e1ppa): Add customisation for reserved name.
+// Would require constexpr strings and thus certain limitations.
+template <CPolymorphicEnumMapping TMapping>
+class TPolymorphicYsonStruct
+{
+    // TODO(arkady-e1ppa): Support non refcounted hierarchies
+    // e.g. shared_ptr<TYsonStructLite> or shared_ptr<TExternalizedYsonStruct>.
+    // TODO(arkady-e1ppa): Support lookup by enum instead of concrete type.
+    // would probably need enum->Type static mapping which can be done
+    // with some enum indexed tuple (But one has to implement it first)
+    // TODO(arkady-e1ppa): Support ctor from anyone from the
+    // hierarchy.
+    using TKey = typename TMapping::TKey;
+    using TBase = typename TMapping::TBase;
+
+public:
+    using TImplementsYsonStructField = void;
+
+    TPolymorphicYsonStruct() = default;
+
+    TPolymorphicYsonStruct(TKey key, TIntrusivePtr<TBase> ptr) noexcept;
+
+    template <CYsonStructSource TSource>
+    void Load(
+        TSource source,
+        bool postprocess = true,
+        bool setDefaults = true,
+        const NYPath::TYPath& path = {},
+        std::optional<EUnrecognizedStrategy> recursiveUnrecognizedStrategy = {});
+
+    void Save(NYson::IYsonConsumer* consumer) const;
+
+    //! Empty if empty or the type is wrong.
+    template <std::derived_from<TBase> TConcrete>
+    TIntrusivePtr<TConcrete> TryGetConcrete() const;
+
+    TKey GetCurrentType() const;
+
+private:
+    TIntrusivePtr<TBase> Storage_;
+    TKey HeldType_;
+
+    void PrepareInstance(INodePtr& node);
+};
+
+////////////////////////////////////////////////////////////////////////////////
+
+template <CPolymorphicEnumMapping TMapping>
+void Serialize(const TPolymorphicYsonStruct<TMapping>& value, NYson::IYsonConsumer* consumer);
+
+template <CPolymorphicEnumMapping TMapping, CYsonStructSource TSource>
+void Deserialize(TPolymorphicYsonStruct<TMapping>& value, TSource source);
+
+////////////////////////////////////////////////////////////////////////////////
+
+//! Usage:
+/*
+    DEFINE_POLYMORPHIC_YSON_STRUCT(Struct,
+        ((Base)     (TBaseStruct))
+        ((Derived1) (TDerivedStruct1))
+        ((Derived2) (TDerivedStruct2))
+    );
+    .....TypeName.....ActualTypeName
+
+    Macro generates two names:
+    1) TStruct -- polymorphic struct which can be
+    serialized.
+    2) EStructType -- enum which holds short names for
+    hierarchy members. They are keys in serialization.
+*/
+#define DEFINE_POLYMORPHIC_YSON_STRUCT(name, seq)
+
+////////////////////////////////////////////////////////////////////////////////
+
+} // namespace NYT::NYTree
+
+#define POLYMORPHIC_YSON_STRUCT_INL_H_
+#include "polymorphic_yson_struct-inl.h"
+#undef POLYMORPHIC_YSON_STRUCT_INL_H_
diff --git a/yt/yt/core/ytree/unittests/yson_struct_ut.cpp b/yt/yt/core/ytree/unittests/yson_struct_ut.cpp
index 9636396ac5..dd8b43d002 100644
--- a/yt/yt/core/ytree/unittests/yson_struct_ut.cpp
+++ b/yt/yt/core/ytree/unittests/yson_struct_ut.cpp
@@ -4,6 +4,7 @@
 
 #include <yt/yt/core/ytree/ephemeral_node_factory.h>
 #include <yt/yt/core/ytree/fluent.h>
+#include <yt/yt/core/ytree/polymorphic_yson_struct.h>
 #include <yt/yt/core/ytree/tree_builder.h>
 #include <yt/yt/core/ytree/tree_visitor.h>
 #include <yt/yt/core/ytree/ypath_client.h>
@@ -2575,5 +2576,265 @@ TEST(TYsonStructTest, OuterYsonStructWithValidation)
 
 ////////////////////////////////////////////////////////////////////////////////
 
+struct TPolyBase
+    : public TYsonStruct
+{
+    int BaseField;
+
+    REGISTER_YSON_STRUCT(TPolyBase);
+
+    static void Register(TRegistrar registrar)
+    {
+        registrar.Parameter("base_field", &TThis::BaseField)
+            .Default(42);
+    }
+};
+
+struct TPolyDerived1
+    : public TPolyBase
+{
+    int Field1;
+
+    REGISTER_YSON_STRUCT(TPolyDerived1);
+
+    static void Register(TRegistrar registrar)
+    {
+        registrar.Parameter("field1", &TThis::Field1)
+            .Default(33);
+    }
+};
+
+struct TPolyDerived2
+    : public TPolyBase
+{
+    int Field2;
+
+    REGISTER_YSON_STRUCT(TPolyDerived2);
+
+    static void Register(TRegistrar registrar)
+    {
+        registrar.Parameter("field2", &TThis::Field2)
+            .Default(11);
+
+        registrar.UnrecognizedStrategy(EUnrecognizedStrategy::Throw);
+    }
+};
+
+TEST(TYsonStructTest, TestSlicing1)
+{
+    TIntrusivePtr<TPolyBase> sliced = New<TPolyDerived1>();
+
+    auto node = BuildYsonNodeFluently()
+        .BeginMap()
+            .Item("base_field").Value(11)
+            .Item("field1").Value(123)
+        .EndMap();
+
+    Deserialize(sliced, node->AsMap());
+
+    auto concrete = DynamicPointerCast<TPolyDerived1>(sliced);
+
+    EXPECT_TRUE(concrete.operator bool());
+    EXPECT_EQ(concrete->BaseField, 11);
+    EXPECT_EQ(concrete->Field1, 123);
+}
+
+TEST(TYsonStructTest, TestSlicing2)
+{
+    TIntrusivePtr<TPolyBase> sliced = New<TPolyBase>();
+
+    auto node = BuildYsonNodeFluently()
+        .BeginMap()
+            .Item("base_field").Value(11)
+            .Item("field1").Value(123)
+        .EndMap();
+
+    Deserialize(sliced, node->AsMap());
+
+    auto concrete = DynamicPointerCast<TPolyDerived1>(sliced);
+
+    EXPECT_FALSE(concrete.operator bool());
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+DEFINE_POLYMORPHIC_YSON_STRUCT(MyPoly,
+    ((Base) (TPolyBase))
+    ((Drv1) (TPolyDerived1))
+    ((Drv2) (TPolyDerived2))
+);
+
+TEST(TYsonStructTest, TestPolymorphicYsonStruct)
+{
+    TMyPoly poly;
+
+    auto node = BuildYsonNodeFluently()
+        .BeginMap()
+            .Item("type").Value("base")
+            .Item("base_field").Value(11)
+            .Item("field1").Value(123)
+        .EndMap();
+
+    Deserialize(poly, node->AsMap());
+    EXPECT_EQ(poly.GetCurrentType(), EMyPolyType::Base);
+
+    auto basePtr = poly.TryGetConcrete<TPolyBase>();
+    EXPECT_TRUE(basePtr.operator bool());
+    EXPECT_EQ(basePtr->BaseField, 11);
+
+    node = BuildYsonNodeFluently()
+        .BeginMap()
+            .Item("type").Value("drv1")
+            .Item("base_field").Value(14)
+            .Item("field1").Value(111)
+            .Item("field2").Value(1337) // Unrecognized but will be dropped.
+        .EndMap();
+
+    Deserialize(poly, node->AsMap());
+    EXPECT_EQ(poly.GetCurrentType(), EMyPolyType::Drv1);
+
+    auto drv1Ptr = poly.TryGetConcrete<TPolyDerived1>();
+    EXPECT_TRUE(drv1Ptr.operator bool());
+    EXPECT_EQ(drv1Ptr->BaseField, 14);
+    EXPECT_EQ(drv1Ptr->Field1, 111);
+
+    node = BuildYsonNodeFluently()
+        .BeginMap()
+            .Item("type").Value("drv2")
+            .Item("base_field").Value(7)
+            .Item("field1").Value(188)
+        .EndMap();
+
+    // Now field1 is unrecognized.
+    EXPECT_THROW(Deserialize(poly, node->AsMap()), std::exception);
+
+    node = BuildYsonNodeFluently()
+        .BeginMap()
+            .Item("type").Value("drv2")
+            .Item("base_field").Value(17)
+            .Item("field2").Value(144)
+        .EndMap();
+
+    Deserialize(poly, node->AsMap());
+    EXPECT_EQ(poly.GetCurrentType(), EMyPolyType::Drv2);
+
+    auto drv2Ptr = poly.TryGetConcrete<TPolyDerived2>();
+    EXPECT_TRUE(drv2Ptr.operator bool());
+    EXPECT_EQ(drv2Ptr->BaseField, 17);
+    EXPECT_EQ(drv2Ptr->Field2, 144);
+}
+
+TEST(TYsonStructTest, TestPolymorphicYsonStructSaveLoad)
+{
+    auto drv = New<TPolyDerived2>();
+    drv->Field2 = 5;
+    drv->BaseField = 0;
+
+    auto poly = TMyPoly{EMyPolyType::Drv2, std::move(drv)};
+
+    auto serialized = ConvertToYsonString(poly);
+    auto deserialized = ConvertTo<TMyPoly>(serialized);
+
+    EXPECT_EQ(deserialized.GetCurrentType(), EMyPolyType::Drv2);
+
+    drv = poly.TryGetConcrete<TPolyDerived2>();
+
+    EXPECT_TRUE(drv.operator bool());
+    EXPECT_EQ(drv->BaseField, 0);
+    EXPECT_EQ(drv->Field2, 5);
+}
+
+TEST(TYsonStructTest, TestPolymorphicYsonStructMergeIfPossible)
+{
+    TMyPoly poly;
+
+    auto node = BuildYsonNodeFluently()
+        .BeginMap()
+            .Item("type").Value("drv1")
+            .Item("base_field").Value(14)
+            // Field1 is missing -- default is 33
+        .EndMap();
+
+    Deserialize(poly, node->AsMap());
+    EXPECT_EQ(poly.GetCurrentType(), EMyPolyType::Drv1);
+
+    auto drv1Ptr = poly.TryGetConcrete<TPolyDerived1>();
+    EXPECT_TRUE(drv1Ptr.operator bool());
+    EXPECT_EQ(drv1Ptr->BaseField, 14);
+    EXPECT_EQ(drv1Ptr->Field1, 33);
+
+    node = BuildYsonNodeFluently()
+        .BeginMap()
+            .Item("type").Value("drv1")
+            .Item("field1").Value(18)
+            // BaseField is missing -- default is 42
+        .EndMap();
+
+    poly.Load(node->AsMap(), /*postprocess*/false, /*setDefaults*/false);
+    EXPECT_EQ(poly.GetCurrentType(), EMyPolyType::Drv1);
+
+    drv1Ptr = poly.TryGetConcrete<TPolyDerived1>();
+    EXPECT_TRUE(drv1Ptr.operator bool());
+    EXPECT_EQ(drv1Ptr->BaseField, 14); // <- Field must remain the same.
+    EXPECT_EQ(drv1Ptr->Field1, 18);
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+struct TPolyHolder
+    : public TYsonStructLite
+{
+    TMyPoly PolyField;
+
+    REGISTER_YSON_STRUCT_LITE(TPolyHolder);
+
+    static void Register(TRegistrar registrar)
+    {
+        registrar.Parameter("poly_field", &TThis::PolyField)
+            .Default();
+
+        registrar.UnrecognizedStrategy(EUnrecognizedStrategy::ThrowRecursive);
+    }
+};
+
+TEST(TYsonStructTest, TestPolymorphicYsonStructAsField)
+{
+    TPolyHolder holder;
+
+    auto node = BuildYsonNodeFluently()
+        .BeginMap()
+            .Item("poly_field").BeginMap()
+                .Item("type").Value("drv2")
+                .Item("base_field").Value(17)
+                .Item("field2").Value(144)
+            .EndMap()
+        .EndMap();
+
+    Deserialize(holder, node->AsMap());
+
+    EXPECT_EQ(holder.PolyField.GetCurrentType(), EMyPolyType::Drv2);
+    auto drv = holder.PolyField.TryGetConcrete<TPolyDerived2>();
+
+    EXPECT_TRUE(drv);
+    EXPECT_EQ(drv->BaseField, 17);
+    EXPECT_EQ(drv->Field2, 144);
+
+    node = BuildYsonNodeFluently()
+        .BeginMap()
+            .Item("poly_field").BeginMap()
+                .Item("type").Value("drv1")
+                .Item("base_field").Value(17)
+                .Item("field2").Value(144)
+            .EndMap()
+        .EndMap();
+
+    // field2 is unrecognized for drv1. Its unrecognized strategy is
+    // Drop by default but holder has recursive throw so it must
+    // throw.
+    EXPECT_THROW(Deserialize(holder, node->AsMap()), std::exception);
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
 } // namespace
 } // namespace NYT::NYTree
diff --git a/yt/yt/core/ytree/yson_struct.cpp b/yt/yt/core/ytree/yson_struct.cpp
index fbdad0be6b..bb7a7834af 100644
--- a/yt/yt/core/ytree/yson_struct.cpp
+++ b/yt/yt/core/ytree/yson_struct.cpp
@@ -74,7 +74,12 @@ void TYsonStructBase::Load(IInputStream* input)
 void TYsonStructBase::Save(IYsonConsumer* consumer) const
 {
     consumer->OnBeginMap();
+    SaveAsMapFragment(consumer);
+    consumer->OnEndMap();
+}
 
+void TYsonStructBase::SaveAsMapFragment(NYson::IYsonConsumer* consumer) const
+{
     for (const auto& [name, parameter] : Meta_->GetParameterSortedList()) {
         if (!parameter->CanOmitValue(this)) {
             consumer->OnKeyedItem(name);
@@ -90,8 +95,6 @@ void TYsonStructBase::Save(IYsonConsumer* consumer) const
             Serialize(child, consumer);
         }
     }
-
-    consumer->OnEndMap();
 }
 
 void TYsonStructBase::Save(IOutputStream* output) const
diff --git a/yt/yt/core/ytree/yson_struct.h b/yt/yt/core/ytree/yson_struct.h
index 5d384840a8..e0f9b81346 100644
--- a/yt/yt/core/ytree/yson_struct.h
+++ b/yt/yt/core/ytree/yson_struct.h
@@ -82,6 +82,12 @@ public:
 
     void Save(NYson::IYsonConsumer* consumer) const;
 
+    // Doesn't call OnBeginMap/OnEndMap.
+    // Can be used to inject extra data into the serialized format
+    // by some kind of wrapper to be parsed in the wrapper
+    // of the |Load| call.
+    void SaveAsMapFragment(NYson::IYsonConsumer* consumer) const;
+
     void Save(IOutputStream* output) const;
 
     IMapNodePtr GetLocalUnrecognized() const;
@@ -181,6 +187,32 @@ protected:
 
 ////////////////////////////////////////////////////////////////////////////////
 
+template <class T, class S>
+concept CYsonStructFieldFor =
+    CYsonStructSource<S> &&
+    requires (
+        T& parameter,
+        S source,
+        bool postprocess,
+        bool setDefaults,
+        const NYPath::TYPath& path,
+        std::optional<EUnrecognizedStrategy> recursiveUnrecognizedStrategy)
+    {
+        // NB(arkady-e1ppa): This alias serves no purpose other
+        // than an easy way to grep for every implementation.
+        typename T::TImplementsYsonStructField;
+
+        // For YsonStruct.
+        parameter.Load(
+            source,
+            postprocess,
+            setDefaults,
+            path,
+            recursiveUnrecognizedStrategy);
+    };
+
+////////////////////////////////////////////////////////////////////////////////
+
 YT_DECLARE_THREAD_LOCAL(IYsonStructMeta*, CurrentlyInitializingYsonMeta);
 YT_DECLARE_THREAD_LOCAL(i64, YsonMetaRegistryDepth);
 
diff --git a/yt/yt/core/ytree/yson_struct_detail-inl.h b/yt/yt/core/ytree/yson_struct_detail-inl.h
index 21b18e4bbb..1deebab6ca 100644
--- a/yt/yt/core/ytree/yson_struct_detail-inl.h
+++ b/yt/yt/core/ytree/yson_struct_detail-inl.h
@@ -182,7 +182,7 @@ struct TYsonSourceTraits<NYson::TYsonPullParserCursor*>
 // e.g. std::optional<std::vector<T>>.
 
 // std::optional
-template <class T, CYsonStructSource TSource>
+template <CYsonStructSource TSource, class T>
 void LoadFromSource(
     std::optional<T>& parameter,
     TSource source,
@@ -190,7 +190,7 @@ void LoadFromSource(
     std::optional<EUnrecognizedStrategy> recursiveUnrecognizedStrategy);
 
 // std::vector
-template <CStdVector TVector, CYsonStructSource TSource>
+template <CYsonStructSource TSource, CStdVector TVector>
 void LoadFromSource(
     TVector& parameter,
     TSource source,
@@ -198,7 +198,7 @@ void LoadFromSource(
     std::optional<EUnrecognizedStrategy> recursiveUnrecognizedStrategy);
 
 // any map.
-template <CAnyMap TMap, CYsonStructSource TSource>
+template <CYsonStructSource TSource, CAnyMap TMap>
 void LoadFromSource(
     TMap& parameter,
     TSource source,
@@ -208,7 +208,7 @@ void LoadFromSource(
 ////////////////////////////////////////////////////////////////////////////////
 
 // Primitive type
-template <class T, CYsonStructSource TSource>
+template <CYsonStructSource TSource, class T>
 void LoadFromSource(
     T& parameter,
     TSource source,
@@ -249,7 +249,7 @@ void LoadFromSource(
 }
 
 // TYsonStruct
-template <CYsonStructDerived T, CYsonStructSource TSource>
+template <CYsonStructSource TSource, CYsonStructDerived T>
 void LoadFromSource(
     TIntrusivePtr<T>& parameter,
     TSource source,
@@ -268,7 +268,7 @@ void LoadFromSource(
 }
 
 // YsonStructLite
-template <std::derived_from<TYsonStructLite> T, CYsonStructSource TSource>
+template <CYsonStructSource TSource, std::derived_from<TYsonStructLite> T>
 void LoadFromSource(
     T& parameter,
     TSource source,
@@ -284,7 +284,7 @@ void LoadFromSource(
 }
 
 // ExternalizedYsonStruct
-template <CExternallySerializable T, CYsonStructSource TSource>
+template <CYsonStructSource TSource, CExternallySerializable T>
 void LoadFromSource(
     T& parameter,
     TSource source,
@@ -299,8 +299,29 @@ void LoadFromSource(
     }
 }
 
+// CYsonStructExtension
+template <CYsonStructSource TSource, CYsonStructFieldFor<TSource> TExtension>
+void LoadFromSource(
+    TExtension& parameter,
+    TSource source,
+    const NYPath::TYPath& path,
+    std::optional<EUnrecognizedStrategy> recursiveUnrecognizedStrategy)
+{
+    try {
+        parameter.Load(
+            std::move(source),
+            /*postprocess*/ false,
+            /*setDefaults*/ false,
+            path,
+            recursiveUnrecognizedStrategy);
+    } catch (const std::exception& ex) {
+        THROW_ERROR_EXCEPTION("Error loading parameter %v", path)
+            << ex;
+    }
+}
+
 // std::optional
-template <class T, CYsonStructSource TSource>
+template <CYsonStructSource TSource, class T>
 void LoadFromSource(
     std::optional<T>& parameter,
     TSource source,
@@ -332,7 +353,7 @@ void LoadFromSource(
 }
 
 // std::vector
-template <CStdVector TVector, CYsonStructSource TSource>
+template <CYsonStructSource TSource, CStdVector TVector>
 void LoadFromSource(
     TVector& parameter,
     TSource source,
@@ -360,7 +381,7 @@ void LoadFromSource(
 }
 
 // any map.
-template <CAnyMap TMap, CYsonStructSource TSource>
+template <CYsonStructSource TSource, CAnyMap TMap>
 void LoadFromSource(
     TMap& parameter,
     TSource source,