YT-20732: Allow yson_struct_lite fields

e950a68ff542bfb63149f0b398a0c6fa85adffbd

4 files changed, 304 insertions, 15 deletions

diff --git a/yt/yt/core/ytree/unittests/yson_struct_ut.cpp b/yt/yt/core/ytree/unittests/yson_struct_ut.cpp
index 1ecb485db7..b86a966d95 100644
--- a/yt/yt/core/ytree/unittests/yson_struct_ut.cpp
+++ b/yt/yt/core/ytree/unittests/yson_struct_ut.cpp
@@ -731,6 +731,28 @@ TEST(TYsonStructTest, Reconfigure)
     EXPECT_EQ(95, subconfig->MyInt);
 }
 
+struct TTestYsonStructWithFieldInitializer
+    : public TYsonStruct
+{
+    TTestSubconfigPtr Sub = New<TTestSubconfig>();
+
+    REGISTER_YSON_STRUCT(TTestYsonStructWithFieldInitializer);
+
+    static void Register(TRegistrar registrar)
+    {
+        registrar.Parameter("sub", &TThis::Sub)
+            .DefaultNew();
+    }
+};
+
+TEST(TYsonStructTest, TestNestedWithFieldInitializer)
+{
+    using TConfig = TTestYsonStructWithFieldInitializer;
+    using TPtr = TIntrusivePtr<TConfig>;
+
+    auto yson = ConvertTo<TPtr>(TYsonString(TStringBuf("{}")));
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 
 class TTestConfigLite
@@ -816,6 +838,242 @@ TEST(TYsonStructTest, TestConvertToLite)
 
 ////////////////////////////////////////////////////////////////////////////////
 
+struct TTestLiteFieldNormalYson
+    : public virtual TYsonStruct
+{
+    TTestLiteWithDefaults SubLite;
+
+    REGISTER_YSON_STRUCT(TTestLiteFieldNormalYson);
+
+    static void Register(TRegistrar registrar)
+    {
+        registrar.Parameter("sub_lite", &TThis::SubLite)
+            .Default();
+    }
+};
+
+struct TTestLiteFieldNormalYsonSecondBase
+    : public virtual TYsonStruct
+{
+    TTestLiteWithDefaults SubLite2;
+
+    REGISTER_YSON_STRUCT(TTestLiteFieldNormalYsonSecondBase);
+
+    static void Register(TRegistrar registrar)
+    {
+        registrar.Parameter("sub_lite_2", &TThis::SubLite2)
+            .Default();
+    }
+};
+
+struct TTestLiteFieldNormalYsonDoubleDerived
+    : public TTestLiteFieldNormalYson
+    , public TTestLiteFieldNormalYsonSecondBase
+{
+    TTestLiteWithDefaults SubLite3;
+
+    REGISTER_YSON_STRUCT(TTestLiteFieldNormalYsonDoubleDerived);
+
+    static void Register(TRegistrar registrar)
+    {
+        registrar.Parameter("sub_lite_3", &TThis::SubLite3)
+            .Default();
+    }
+};
+
+TEST(TYsonStructTest, YsonStructWithLiteField)
+{
+    using TConfig = TTestLiteFieldNormalYson;
+    using TConfigPtr = TIntrusivePtr<TConfig>;
+
+    {
+        auto yson = ConvertTo<TConfigPtr>(TYsonString(TStringBuf("{}")));
+        auto& sub = yson->SubLite;
+        EXPECT_EQ(sub.MyString, "y");
+        EXPECT_EQ(sub.MyInt, 10);
+        EXPECT_TRUE(sub.Subconfig);
+    }
+
+    {
+        auto yson = New<TConfig>();
+        auto& sub = yson->SubLite;
+        EXPECT_EQ(sub.MyString, "y");
+        EXPECT_EQ(sub.MyInt, 10);
+        EXPECT_TRUE(sub.Subconfig);
+    }
+}
+
+TEST(TYsonStructTest, DoubleDerivedYsonStructWithLiteFields)
+{
+    using TConfig = TTestLiteFieldNormalYsonDoubleDerived;
+    using TConfigPtr = TIntrusivePtr<TConfig>;
+
+    {
+        auto yson = ConvertTo<TConfigPtr>(TYsonString(TStringBuf("{}")));
+
+        {
+            auto& sub = yson->SubLite;
+            EXPECT_EQ(sub.MyString, "y");
+            EXPECT_EQ(sub.MyInt, 10);
+            EXPECT_TRUE(sub.Subconfig);
+        }
+        {
+            auto& sub = yson->SubLite2;
+            EXPECT_EQ(sub.MyString, "y");
+            EXPECT_EQ(sub.MyInt, 10);
+            EXPECT_TRUE(sub.Subconfig);
+        }
+        {
+            auto& sub = yson->SubLite3;
+            EXPECT_EQ(sub.MyString, "y");
+            EXPECT_EQ(sub.MyInt, 10);
+            EXPECT_TRUE(sub.Subconfig);
+        }
+    }
+
+    {
+        auto yson = New<TConfig>();
+
+        {
+            auto& sub = yson->SubLite;
+            EXPECT_EQ(sub.MyString, "y");
+            EXPECT_EQ(sub.MyInt, 10);
+            EXPECT_TRUE(sub.Subconfig);
+        }
+        {
+            auto& sub = yson->SubLite2;
+            EXPECT_EQ(sub.MyString, "y");
+            EXPECT_EQ(sub.MyInt, 10);
+            EXPECT_TRUE(sub.Subconfig);
+        }
+        {
+            auto& sub = yson->SubLite3;
+            EXPECT_EQ(sub.MyString, "y");
+            EXPECT_EQ(sub.MyInt, 10);
+            EXPECT_TRUE(sub.Subconfig);
+        }
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+struct TTestLiteFieldLite
+    : public virtual TYsonStructLite
+{
+    TTestLiteWithDefaults SubLite;
+
+    REGISTER_YSON_STRUCT_LITE(TTestLiteFieldLite);
+
+    static void Register(TRegistrar registrar)
+    {
+        registrar.Parameter("sub_lite", &TThis::SubLite)
+            .Default();
+    }
+};
+
+struct TTestLiteFieldLiteSecondBase
+    : public virtual TYsonStructLite
+{
+    TTestLiteWithDefaults SubLite2;
+
+    REGISTER_YSON_STRUCT_LITE(TTestLiteFieldLiteSecondBase);
+
+    static void Register(TRegistrar registrar)
+    {
+        registrar.Parameter("sub_lite_2", &TThis::SubLite2)
+            .Default();
+    }
+};
+
+struct TTestLiteFieldLiteDoubleDerived
+    : public TTestLiteFieldLite
+    , public TTestLiteFieldLiteSecondBase
+{
+    TTestLiteWithDefaults SubLite3;
+
+    REGISTER_YSON_STRUCT_LITE(TTestLiteFieldLiteDoubleDerived);
+
+    static void Register(TRegistrar registrar)
+    {
+        registrar.Parameter("sub_lite_3", &TThis::SubLite3)
+            .Default();
+    }
+};
+
+TEST(TYsonStructTest, LiteWithLiteField)
+{
+    using TConfig = TTestLiteFieldLite;
+
+    {
+        auto yson = ConvertTo<TTestLiteFieldLite>(TYsonString(TStringBuf("{}")));
+        auto& sub = yson.SubLite;
+        EXPECT_EQ(sub.MyString, "y");
+        EXPECT_EQ(sub.MyInt, 10);
+        EXPECT_TRUE(sub.Subconfig);
+    }
+
+    {
+        TConfig yson;
+        auto& sub = yson.SubLite;
+        EXPECT_EQ(sub.MyString, "y");
+        EXPECT_EQ(sub.MyInt, 10);
+        EXPECT_TRUE(sub.Subconfig);
+    }
+}
+
+TEST(TYsonStructTest, DoubleDerivedLiteWithLiteFields)
+{
+    using TConfig = TTestLiteFieldLiteDoubleDerived;
+
+    {
+        auto yson = ConvertTo<TConfig>(TYsonString(TStringBuf("{}")));
+
+        {
+            auto& sub = yson.SubLite;
+            EXPECT_EQ(sub.MyString, "y");
+            EXPECT_EQ(sub.MyInt, 10);
+            EXPECT_TRUE(sub.Subconfig);
+        }
+        {
+            auto& sub = yson.SubLite2;
+            EXPECT_EQ(sub.MyString, "y");
+            EXPECT_EQ(sub.MyInt, 10);
+            EXPECT_TRUE(sub.Subconfig);
+        }
+        {
+            auto& sub = yson.SubLite3;
+            EXPECT_EQ(sub.MyString, "y");
+            EXPECT_EQ(sub.MyInt, 10);
+            EXPECT_TRUE(sub.Subconfig);
+        }
+    }
+
+    {
+        TConfig yson;
+
+        {
+            auto& sub = yson.SubLite;
+            EXPECT_EQ(sub.MyString, "y");
+            EXPECT_EQ(sub.MyInt, 10);
+            EXPECT_TRUE(sub.Subconfig);
+        }
+        {
+            auto& sub = yson.SubLite2;
+            EXPECT_EQ(sub.MyString, "y");
+            EXPECT_EQ(sub.MyInt, 10);
+            EXPECT_TRUE(sub.Subconfig);
+        }
+        {
+            auto& sub = yson.SubLite3;
+            EXPECT_EQ(sub.MyString, "y");
+            EXPECT_EQ(sub.MyInt, 10);
+            EXPECT_TRUE(sub.Subconfig);
+        }
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
 class TTestConfigWithAliases
     : public TYsonStruct
 {
@@ -1980,7 +2238,7 @@ class TTestConfigWithSubStructLite
     : public TYsonStructLite
 {
 public:
-    std::optional<TTestSubConfigLiteWithDefaults> Sub;
+    TTestSubConfigLiteWithDefaults Sub;
 
     REGISTER_YSON_STRUCT_LITE(TTestConfigWithSubStructLite);
 
@@ -2004,8 +2262,8 @@ TEST(TYsonStructTest, CustomSubStructLite)
         .BeginMap()
         .EndMap();
     testStruct.Load(testNode->AsMap());
-    EXPECT_EQ(testStruct.Sub->MyInt, 11);
-    EXPECT_EQ(testStruct.Sub->MyString, "x");
+    EXPECT_EQ(testStruct.Sub.MyInt, 11);
+    EXPECT_EQ(testStruct.Sub.MyString, "x");
 
     testNode = BuildYsonNodeFluently()
         .BeginMap()
@@ -2014,8 +2272,8 @@ TEST(TYsonStructTest, CustomSubStructLite)
                 .EndMap()
         .EndMap();
     testStruct.Load(testNode->AsMap());
-    EXPECT_EQ(testStruct.Sub->MyInt, 11);
-    EXPECT_EQ(testStruct.Sub->MyString, "x");
+    EXPECT_EQ(testStruct.Sub.MyInt, 11);
+    EXPECT_EQ(testStruct.Sub.MyString, "x");
 
     testNode = BuildYsonNodeFluently()
         .BeginMap()
@@ -2025,8 +2283,8 @@ TEST(TYsonStructTest, CustomSubStructLite)
                 .EndMap()
         .EndMap();
     testStruct.Load(testNode->AsMap());
-    EXPECT_EQ(testStruct.Sub->MyInt, 11);
-    EXPECT_EQ(testStruct.Sub->MyString, "C");
+    EXPECT_EQ(testStruct.Sub.MyInt, 11);
+    EXPECT_EQ(testStruct.Sub.MyString, "C");
 }
 
 ////////////////////////////////////////////////////////////////////////////////
diff --git a/yt/yt/core/ytree/yson_struct-inl.h b/yt/yt/core/ytree/yson_struct-inl.h
index be12349395..88fafd913d 100644
--- a/yt/yt/core/ytree/yson_struct-inl.h
+++ b/yt/yt/core/ytree/yson_struct-inl.h
@@ -104,7 +104,9 @@ void TYsonStructRegistry::InitializeStruct(TStruct* target)
     if (CurrentlyInitializingMeta_) {
         // TODO(renadeen): assert target is from the same type hierarchy.
         // Call initialization method that is provided by user.
-        TStruct::Register(TYsonStructRegistrar<TStruct>(CurrentlyInitializingMeta_));
+        if (RegistryDepth_ <= 1) {
+            TStruct::Register(TYsonStructRegistrar<TStruct>(CurrentlyInitializingMeta_));
+        }
         return;
     }
 
@@ -428,13 +430,16 @@ private: \
     using TThis = TStruct; \
     friend class ::NYT::NYTree::TYsonStructRegistry;
 
-#define YSON_STRUCT_IMPL__CTOR_BODY(TStruct) \
+#define YSON_STRUCT_IMPL__CTOR_BODY \
     ::NYT::NYTree::TYsonStructRegistry::Get()->InitializeStruct(this);
 
 #define YSON_STRUCT_LITE_IMPL__CTOR_BODY(TStruct) \
-    YSON_STRUCT_IMPL__CTOR_BODY(TStruct) \
-    if (std::type_index(typeid(TStruct)) == this->FinalType_ && !::NYT::NYTree::TYsonStructRegistry::Get()->InitializationInProgress()) { \
-        this->SetDefaults(); \
+    YSON_STRUCT_IMPL__CTOR_BODY \
+    if (std::type_index(typeid(TStruct)) == this->FinalType_) { \
+        ::NYT::NYTree::TYsonStructRegistry::Get()->OnFinalCtorCalled(); \
+        if (!::NYT::NYTree::TYsonStructRegistry::Get()->InitializationInProgress()) { \
+            this->SetDefaults(); \
+        } \
     } \
 
 //! NB(arkady-e1ppa): Alias is used by registrar postprocessors
@@ -462,11 +467,10 @@ public: \
     TStruct() \
     { \
         static_assert(std::derived_from<TStruct, ::NYT::NYTree::TYsonStruct>, "Class must inherit from TYsonStruct"); \
-        YSON_STRUCT_IMPL__CTOR_BODY(TStruct) \
+        YSON_STRUCT_IMPL__CTOR_BODY \
     } \
     YSON_STRUCT_IMPL__DECLARE_ALIASES(TStruct)
 
-
 #define DECLARE_YSON_STRUCT_LITE(TStruct) \
 public: \
     TStruct(); \
@@ -494,7 +498,7 @@ TStruct::TStruct() \
 TStruct::TStruct() \
 { \
     static_assert(std::derived_from<TStruct, ::NYT::NYTree::TYsonStruct>, "Class must inherit from TYsonStruct"); \
-    YSON_STRUCT_IMPL__CTOR_BODY(TStruct) \
+    YSON_STRUCT_IMPL__CTOR_BODY \
 }
 
 //! NB(arkady-e1ppa): These constructors are only used internally.
diff --git a/yt/yt/core/ytree/yson_struct.cpp b/yt/yt/core/ytree/yson_struct.cpp
index 95714319c8..2afabecb0b 100644
--- a/yt/yt/core/ytree/yson_struct.cpp
+++ b/yt/yt/core/ytree/yson_struct.cpp
@@ -22,6 +22,11 @@ TYsonStructFinalClassHolder::TYsonStructFinalClassHolder(std::type_index typeInd
 
 ////////////////////////////////////////////////////////////////////////////////
 
+TYsonStructBase::TYsonStructBase()
+{
+    TYsonStructRegistry::Get()->OnBaseCtorCalled();
+}
+
 IMapNodePtr TYsonStructBase::GetLocalUnrecognized() const
 {
     return LocalUnrecognized_;
@@ -144,6 +149,7 @@ void TYsonStructBase::WriteSchema(IYsonConsumer* consumer) const
 
 void TYsonStruct::InitializeRefCounted()
 {
+    TYsonStructRegistry::Get()->OnFinalCtorCalled();
     if (!TYsonStructRegistry::InitializationInProgress()) {
         SetDefaults();
     }
@@ -161,6 +167,20 @@ bool TYsonStructRegistry::InitializationInProgress()
     return CurrentlyInitializingMeta_ != nullptr;
 }
 
+void TYsonStructRegistry::OnBaseCtorCalled()
+{
+    if (CurrentlyInitializingMeta_ != nullptr) {
+        ++RegistryDepth_;
+    }
+}
+
+void TYsonStructRegistry::OnFinalCtorCalled()
+{
+    if (CurrentlyInitializingMeta_ != nullptr) {
+        --RegistryDepth_;
+    }
+}
+
 TYsonStructRegistry::TForbidCachedDynamicCastGuard::TForbidCachedDynamicCastGuard(TYsonStructBase* target)
     : Target_(target)
 {
diff --git a/yt/yt/core/ytree/yson_struct.h b/yt/yt/core/ytree/yson_struct.h
index fd66023a47..e6efb571a2 100644
--- a/yt/yt/core/ytree/yson_struct.h
+++ b/yt/yt/core/ytree/yson_struct.h
@@ -56,6 +56,8 @@ public:
     using TPostprocessor = std::function<void()>;
     using TPreprocessor = std::function<void()>;
 
+    TYsonStructBase();
+
     virtual ~TYsonStructBase() = default;
 
     void Load(
@@ -187,8 +189,13 @@ public:
     template <class TStruct>
     void InitializeStruct(TStruct* target);
 
+    void OnBaseCtorCalled();
+
+    void OnFinalCtorCalled();
+
 private:
     static inline YT_THREAD_LOCAL(IYsonStructMeta*) CurrentlyInitializingMeta_ = nullptr;
+    static inline YT_THREAD_LOCAL(i64) RegistryDepth_ = 0;
 
     template <class TStruct>
     friend class TYsonStructRegistrar;