YT-22885: Handle unknown values in (Try)CheckedEnumCast

commit_hash:5ce8019253cdb971d1af36350e3efa3a4ec8545c

3 files changed, 61 insertions, 3 deletions

diff --git a/library/cpp/yt/misc/cast-inl.h b/library/cpp/yt/misc/cast-inl.h
index 50920f0193..5e5a0d90ee 100644
--- a/library/cpp/yt/misc/cast-inl.h
+++ b/library/cpp/yt/misc/cast-inl.h
@@ -119,7 +119,7 @@ T CheckedIntegralCast(S value)
 
 template <class T, class S>
     requires TEnumTraits<T>::IsEnum
-constexpr std::optional<T> TryCheckedEnumCast(S value)
+constexpr std::optional<T> TryCheckedEnumCast(S value, bool enableUnknown)
 {
     auto underlying = TryCheckedIntegralCast<std::underlying_type_t<T>>(value);
     [[unlikely]] if (!underlying) {
@@ -127,6 +127,15 @@ constexpr std::optional<T> TryCheckedEnumCast(S value)
     }
     auto candidate = static_cast<T>(*underlying);
     [[unlikely]] if (!TEnumTraits<T>::IsValidValue(candidate)) {
+        if (enableUnknown) {
+            if constexpr (constexpr auto optionalUnknownValue = TEnumTraits<T>::TryGetUnknownValue()) {
+                if constexpr (TEnumTraits<T>::IsBitEnum) {
+                    return static_cast<T>(*underlying & ToUnderlying(TEnumTraits<T>::GetAllSetValue())) | *optionalUnknownValue;
+                } else {
+                    return *optionalUnknownValue;
+                }
+            }
+        }
         return std::nullopt;
     }
     return candidate;
@@ -136,7 +145,7 @@ template <class T, class S>
     requires TEnumTraits<T>::IsEnum
 T CheckedEnumCast(S value)
 {
-    auto result = TryCheckedEnumCast<T>(value);
+    auto result = TryCheckedEnumCast<T>(value, /*enableUnknown*/ true);
     [[unlikely]] if (!result) {
         if constexpr (std::is_signed_v<S>) {
             throw TSimpleException(Sprintf("Error casting %s value \"%" PRIi64 "\" to enum %s",
diff --git a/library/cpp/yt/misc/cast.h b/library/cpp/yt/misc/cast.h
index a4e32e0287..ccf1eb21a6 100644
--- a/library/cpp/yt/misc/cast.h
+++ b/library/cpp/yt/misc/cast.h
@@ -26,7 +26,7 @@ T CheckedIntegralCast(S value);
 
 template <class T, class S>
     requires TEnumTraits<T>::IsEnum
-constexpr std::optional<T> TryCheckedEnumCast(S value);
+constexpr std::optional<T> TryCheckedEnumCast(S value, bool enableUnknown = false);
 
 template <class T, class S>
     requires TEnumTraits<T>::IsEnum
diff --git a/library/cpp/yt/misc/unittests/cast_ut.cpp b/library/cpp/yt/misc/unittests/cast_ut.cpp
index 8662575d6c..e6f5409f2f 100644
--- a/library/cpp/yt/misc/unittests/cast_ut.cpp
+++ b/library/cpp/yt/misc/unittests/cast_ut.cpp
@@ -21,6 +21,26 @@ DEFINE_BIT_ENUM_WITH_UNDERLYING_TYPE(EFeatures, ui8,
     ((Second)(0x0002))
 );
 
+DEFINE_BIT_ENUM(ELangsWithUnknown,
+    ((None)       (0x00))
+    ((Cpp)        (0x01))
+    ((Go)         (0x02))
+    ((Rust)       (0x04))
+    ((Python)     (0x08))
+    ((JavaScript) (0x10))
+    ((CppGo)      (0x03))
+    ((All)        (0x1f))
+    ((Unknown)    (0x20))
+);
+DEFINE_ENUM_UNKNOWN_VALUE(ELangsWithUnknown, Unknown);
+
+DEFINE_ENUM(EColorWithUnknown,
+    ((Red)     (0))
+    ((Green)   (1))
+    ((Unknown) (2))
+);
+DEFINE_ENUM_UNKNOWN_VALUE(EColorWithUnknown, Unknown);
+
 TEST(TCastTest, TryCheckedEnumCast)
 {
     EXPECT_EQ((TryCheckedEnumCast<ECardinal, char>(2)), ECardinal::East);
@@ -36,6 +56,35 @@ TEST(TCastTest, TryCheckedEnumCast)
     EXPECT_EQ((TryCheckedEnumCast<EFeatures, ui8>(ToUnderlying(EFeatures::First | EFeatures::Second))), EFeatures::First | EFeatures::Second);
 
     EXPECT_FALSE((TryCheckedEnumCast<EFeatures, ui8>(0x10)));
+
+    EXPECT_FALSE(TryCheckedEnumCast<EColorWithUnknown>(3));
+    EXPECT_EQ(TryCheckedEnumCast<EColorWithUnknown>(3, /*enableUnknown*/ true), EColorWithUnknown::Unknown);
+
+    EXPECT_FALSE(TryCheckedEnumCast<ELangsWithUnknown>(0x40));
+    EXPECT_EQ(TryCheckedEnumCast<ELangsWithUnknown>(0x40, /*enableUnknown*/ true), ELangsWithUnknown::Unknown);
+    EXPECT_EQ(TryCheckedEnumCast<ELangsWithUnknown>(0x41, /*enableUnknown*/ true), ELangsWithUnknown::Unknown | ELangsWithUnknown::Cpp);
+}
+
+TEST(TCastTest, CheckedEnumCast)
+{
+    EXPECT_EQ((CheckedEnumCast<ECardinal, char>(2)), ECardinal::East);
+    EXPECT_EQ((CheckedEnumCast<ECardinal, int>(3)), ECardinal::South);
+
+    EXPECT_THROW((CheckedEnumCast<ECardinal, char>(100)), TSimpleException);
+    EXPECT_THROW((CheckedEnumCast<ECardinal, int>(300)), TSimpleException);
+
+    EXPECT_EQ((CheckedEnumCast<EFeatures, ui8>(0)), EFeatures::None);
+    EXPECT_EQ((CheckedEnumCast<EFeatures, ui8>(ToUnderlying(EFeatures::First))), EFeatures::First);
+    EXPECT_EQ((CheckedEnumCast<EFeatures, ui8>(ToUnderlying(EFeatures::Second))), EFeatures::Second);
+    EXPECT_EQ((CheckedEnumCast<EFeatures, int>(ToUnderlying(EFeatures::First))), EFeatures::First);
+    EXPECT_EQ((CheckedEnumCast<EFeatures, ui8>(ToUnderlying(EFeatures::First | EFeatures::Second))), EFeatures::First | EFeatures::Second);
+
+    EXPECT_THROW((CheckedEnumCast<EFeatures, ui8>(0x10)), TSimpleException);
+
+    EXPECT_EQ(CheckedEnumCast<EColorWithUnknown>(3), EColorWithUnknown::Unknown);
+
+    EXPECT_EQ(CheckedEnumCast<ELangsWithUnknown>(0x40), ELangsWithUnknown::Unknown);
+    EXPECT_EQ(CheckedEnumCast<ELangsWithUnknown>(0x41), ELangsWithUnknown::Unknown | ELangsWithUnknown::Cpp);
 }
 
 ////////////////////////////////////////////////////////////////////////////////