Merge pull request #44637 from ydb-platform/merge-rightlib-260626-1120

7 files changed, 199 insertions, 0 deletions

diff --git a/library/cpp/yt/memory/exact_ref_counted_cast-inl.h b/library/cpp/yt/memory/exact_ref_counted_cast-inl.h
new file mode 100644
index 00000000000..71163e26e0a
--- /dev/null
+++ b/library/cpp/yt/memory/exact_ref_counted_cast-inl.h
@@ -0,0 +1,46 @@
+#ifndef EXACT_REF_COUNTED_CAST_INL_H_
+#error "Direct inclusion of this file is not allowed, include exact_ref_counted_cast.h"
+// For the sake of sane code completion.
+#include "exact_ref_counted_cast.h"
+#endif
+
+#include "new.h"
+
+#include <typeinfo>
+
+namespace NYT {
+
+////////////////////////////////////////////////////////////////////////////////
+
+template <class T, class U>
+std::conditional_t<std::is_const_v<U>, const T, T>* ExactRefCountedCast(U* ptr) noexcept
+{
+    static_assert(std::is_polymorphic_v<U>,
+        "ExactRefCountedCast requires a polymorphic operand type");
+    static_assert(std::is_base_of_v<TRefCountedBase, T>,
+        "ExactRefCountedCast target must derive from TRefCountedBase");
+    static_assert(std::is_final_v<TRefCountedWrapper<T>>,
+        "TRefCountedWrapper must stay final for this cast to remain O(1)");
+
+    using TWrapper = std::conditional_t<std::is_const_v<U>, const TRefCountedWrapper<T>, TRefCountedWrapper<T>>;
+    // Wrapper is final, so dynamic_cast lowers to one type_info compare.
+    return dynamic_cast<TWrapper*>(ptr);
+}
+
+template <class T, class U>
+std::conditional_t<std::is_const_v<U>, const T, T>& ExactRefCountedCast(U& ref)
+{
+    static_assert(std::is_polymorphic_v<U>,
+        "ExactRefCountedCast requires a polymorphic operand type");
+    static_assert(std::is_base_of_v<TRefCountedBase, T>,
+        "ExactRefCountedCast target must derive from TRefCountedBase");
+    static_assert(std::is_final_v<TRefCountedWrapper<T>>,
+        "TRefCountedWrapper must stay final for this cast to remain O(1)");
+
+    using TWrapper = std::conditional_t<std::is_const_v<U>, const TRefCountedWrapper<T>, TRefCountedWrapper<T>>;
+    return dynamic_cast<TWrapper&>(ref);
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+} // namespace NYT
diff --git a/library/cpp/yt/memory/exact_ref_counted_cast.h b/library/cpp/yt/memory/exact_ref_counted_cast.h
new file mode 100644
index 00000000000..3b186103f3a
--- /dev/null
+++ b/library/cpp/yt/memory/exact_ref_counted_cast.h
@@ -0,0 +1,36 @@
+#pragma once
+
+#include "ref_counted.h"
+
+#include <type_traits>
+
+namespace NYT {
+
+////////////////////////////////////////////////////////////////////////////////
+
+//! A downcast that succeeds iff the object was allocated as *exactly* |New<T>()|
+//! (not a subclass, not some other allocation path).
+/*!
+ *  |New<T>()| does not construct a |T| but a final |TRefCountedWrapper<T>|
+ *  deriving from |T|, so |T| is never a leaf and |dynamic_cast<T*>| takes the
+ *  slow is-a path. This casts to the wrapper instead and upcasts back to |T*|;
+ *  since the wrapper is final, dynamic_cast lowers to one type_info compare --
+ *  O(1), as cheap on a miss as on a hit.
+ *
+ *  Pointer form returns |nullptr| on a miss (or null operand); reference form
+ *  throws |std::bad_cast|. Constness is preserved. |T| must derive from
+ *  |TRefCountedBase| and |U| must be polymorphic. Requires RTTI.
+ */
+template <class T, class U>
+[[nodiscard]] std::conditional_t<std::is_const_v<U>, const T, T>* ExactRefCountedCast(U* ptr) noexcept;
+
+template <class T, class U>
+[[nodiscard]] std::conditional_t<std::is_const_v<U>, const T, T>& ExactRefCountedCast(U& ref);
+
+////////////////////////////////////////////////////////////////////////////////
+
+} // namespace NYT
+
+#define EXACT_REF_COUNTED_CAST_INL_H_
+#include "exact_ref_counted_cast-inl.h"
+#undef EXACT_REF_COUNTED_CAST_INL_H_
diff --git a/library/cpp/yt/memory/ref-inl.h b/library/cpp/yt/memory/ref-inl.h
index 632a98f130d..5c6dc11e0a2 100644
--- a/library/cpp/yt/memory/ref-inl.h
+++ b/library/cpp/yt/memory/ref-inl.h
@@ -65,6 +65,11 @@ Y_FORCE_INLINE TRef TRef::Slice(size_t startOffset, size_t endOffset) const
     return TRef(Begin() + startOffset, endOffset - startOffset);
 }
 
+Y_FORCE_INLINE bool TRef::Contains(TRef other) const
+{
+    return other.Begin() >= Begin() && other.End() <= End();
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 
 Y_FORCE_INLINE TMutableRef::TMutableRef(void* data, size_t size)
diff --git a/library/cpp/yt/memory/ref.h b/library/cpp/yt/memory/ref.h
index a6c6c4f2260..1fc451f50aa 100644
--- a/library/cpp/yt/memory/ref.h
+++ b/library/cpp/yt/memory/ref.h
@@ -54,6 +54,9 @@ public:
     //! Creates a TRef for a part of existing range.
     TRef Slice(size_t startOffset, size_t endOffset) const;
 
+    //! Returns |true| if #other's range lies entirely within this range.
+    bool Contains(TRef other) const;
+
     //! Compares the content for bitwise equality.
     static bool AreBitwiseEqual(TRef lhs, TRef rhs);
 };
diff --git a/library/cpp/yt/memory/unittests/exact_ref_counted_cast_ut.cpp b/library/cpp/yt/memory/unittests/exact_ref_counted_cast_ut.cpp
new file mode 100644
index 00000000000..75cce328dc0
--- /dev/null
+++ b/library/cpp/yt/memory/unittests/exact_ref_counted_cast_ut.cpp
@@ -0,0 +1,102 @@
+#include <library/cpp/testing/gtest/gtest.h>
+
+#include <library/cpp/yt/memory/exact_ref_counted_cast.h>
+#include <library/cpp/yt/memory/new.h>
+#include <library/cpp/yt/memory/ref_counted.h>
+
+#include <type_traits>
+#include <typeinfo>
+
+namespace NYT {
+namespace {
+
+////////////////////////////////////////////////////////////////////////////////
+
+struct TBase
+    : public TRefCounted
+{ };
+
+struct TDerived
+    : public TBase
+{
+    int Tag = 42;
+};
+
+struct TMoreDerived
+    : public TDerived
+{ };
+
+////////////////////////////////////////////////////////////////////////////////
+// Pointer form.
+
+TEST(TExactRefCountedCastTest, ExactMatch)
+{
+    auto object = New<TDerived>();
+    TBase* base = object.Get();
+    EXPECT_EQ(ExactRefCountedCast<TDerived>(base), object.Get());
+    EXPECT_EQ(ExactRefCountedCast<TDerived>(base)->Tag, 42);
+}
+
+TEST(TExactRefCountedCastTest, SubclassDoesNotMatch)
+{
+    auto object = New<TMoreDerived>();
+    TBase* base = object.Get();
+    // Object was New<TMoreDerived>(); an exact probe for TDerived must miss.
+    EXPECT_EQ(ExactRefCountedCast<TDerived>(base), nullptr);
+    // ...but the exact type matches.
+    EXPECT_EQ(ExactRefCountedCast<TMoreDerived>(base), object.Get());
+}
+
+TEST(TExactRefCountedCastTest, BaseDoesNotMatch)
+{
+    auto object = New<TDerived>();
+    TBase* base = object.Get();
+    // It is exactly a TDerived, not a TBase.
+    EXPECT_EQ(ExactRefCountedCast<TBase>(base), nullptr);
+}
+
+TEST(TExactRefCountedCastTest, Null)
+{
+    TBase* base = nullptr;
+    EXPECT_EQ(ExactRefCountedCast<TDerived>(base), nullptr);
+}
+
+TEST(TExactRefCountedCastTest, ConstnessPreserved)
+{
+    auto object = New<TDerived>();
+    const TBase* base = object.Get();
+    auto* result = ExactRefCountedCast<TDerived>(base);
+    static_assert(std::is_same_v<decltype(result), const TDerived*>);
+    EXPECT_EQ(result, object.Get());
+}
+
+////////////////////////////////////////////////////////////////////////////////
+// Reference form.
+
+TEST(TExactRefCountedCastTest, RefExactMatch)
+{
+    auto object = New<TDerived>();
+    TBase& base = *object;
+    EXPECT_EQ(&ExactRefCountedCast<TDerived>(base), object.Get());
+}
+
+TEST(TExactRefCountedCastTest, RefSubclassThrows)
+{
+    auto object = New<TMoreDerived>();
+    TBase& base = *object;
+    EXPECT_THROW((void)ExactRefCountedCast<TDerived>(base), std::bad_cast);
+    EXPECT_EQ(&ExactRefCountedCast<TMoreDerived>(base), object.Get());
+}
+
+TEST(TExactRefCountedCastTest, RefConstnessPreserved)
+{
+    auto object = New<TDerived>();
+    const TBase& base = *object;
+    static_assert(std::is_same_v<decltype(ExactRefCountedCast<TDerived>(base)), const TDerived&>);
+    EXPECT_EQ(&ExactRefCountedCast<TDerived>(base), object.Get());
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+} // namespace
+} // namespace NYT
diff --git a/library/cpp/yt/memory/unittests/ya.make b/library/cpp/yt/memory/unittests/ya.make
index c3abcbefbe5..c545bc0d41e 100644
--- a/library/cpp/yt/memory/unittests/ya.make
+++ b/library/cpp/yt/memory/unittests/ya.make
@@ -9,6 +9,7 @@ SRCS(
     chunked_memory_pool_allocator_ut.cpp
     chunked_memory_pool_output_ut.cpp
     erased_storage_ut.cpp
+    exact_ref_counted_cast_ut.cpp
     free_list_ut.cpp
     function_view_ut.cpp
     intrusive_ptr_ut.cpp
diff --git a/library/cpp/yt/memory/ya.make b/library/cpp/yt/memory/ya.make
index 37b05048b76..de4cc766061 100644
--- a/library/cpp/yt/memory/ya.make
+++ b/library/cpp/yt/memory/ya.make
@@ -47,6 +47,12 @@ CHECK_DEPENDENT_DIRS(
 
 END()
 
+IF (NOT OPENSOURCE)
+    RECURSE(
+        benchmark
+    )
+ENDIF()
+
 RECURSE_FOR_TESTS(
     unittests
 )