Introduce TEnumIndexedArray as a refurbished version of TEnumIndexedVector

7 files changed, 207 insertions, 0 deletions

diff --git a/library/cpp/yt/misc/enum.h b/library/cpp/yt/misc/enum.h
index 1494ca7907..9da79b9fd3 100644
--- a/library/cpp/yt/misc/enum.h
+++ b/library/cpp/yt/misc/enum.h
@@ -185,6 +185,7 @@ struct TEnumTraits<T, true>
 
 ////////////////////////////////////////////////////////////////////////////////
 
+// TODO(babenko): drop in favor of TEnumIndexedArray
 //! A statically sized vector with elements of type |T| indexed by
 //! the items of enumeration type |E|.
 /*!
diff --git a/library/cpp/yt/misc/enum_indexed_array-inl.h b/library/cpp/yt/misc/enum_indexed_array-inl.h
new file mode 100644
index 0000000000..edda891683
--- /dev/null
+++ b/library/cpp/yt/misc/enum_indexed_array-inl.h
@@ -0,0 +1,73 @@
+#pragma once
+#ifndef ENUM_INDEXED_ARRAY_INL_H_
+#error "Direct inclusion of this file is not allowed, include enum.h"
+// For the sake of sane code completion.
+#include "enum_indexed_array.h"
+#endif
+
+#include <library/cpp/yt/assert/assert.h>
+
+namespace NYT {
+
+////////////////////////////////////////////////////////////////////////////////
+
+template <class E, class T, E Min, E Max>
+TEnumIndexedArray<E, T, Min, Max>::TEnumIndexedArray(std::initializer_list<std::pair<E, T>> elements)
+{
+    for (const auto& [index, value] : elements) {
+        (*this)[index] = value;
+    }
+}
+
+template <class E, class T, E Min, E Max>
+T& TEnumIndexedArray<E, T, Min, Max>::operator[] (E index)
+{
+    YT_ASSERT(IsValidIndex(index));
+    return Items_[ToUnderlying(index) - ToUnderlying(Min)];
+}
+
+template <class E, class T, E Min, E Max>
+const T& TEnumIndexedArray<E, T, Min, Max>::operator[] (E index) const
+{
+    return const_cast<TEnumIndexedArray&>(*this)[index];
+}
+
+template <class E, class T, E Min, E Max>
+T* TEnumIndexedArray<E, T, Min, Max>::begin()
+{
+    return Items_.data();
+}
+
+template <class E, class T, E Min, E Max>
+const T* TEnumIndexedArray<E, T, Min, Max>::begin() const
+{
+    return Items_.data();
+}
+
+template <class E, class T, E Min, E Max>
+T* TEnumIndexedArray<E, T, Min, Max>::end()
+{
+    return begin() + Size;
+}
+
+template <class E, class T, E Min, E Max>
+const T* TEnumIndexedArray<E, T, Min, Max>::end() const
+{
+    return begin() + Size;
+}
+
+template <class E, class T, E Min, E Max>
+constexpr size_t TEnumIndexedArray<E, T, Min, Max>::size() const
+{
+    return Size;
+}
+
+template <class E, class T, E Min, E Max>
+bool TEnumIndexedArray<E, T, Min, Max>::IsValidIndex(E index)
+{
+    return index >= Min && index <= Max;
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+} // namespace NYT
diff --git a/library/cpp/yt/misc/enum_indexed_array.h b/library/cpp/yt/misc/enum_indexed_array.h
new file mode 100644
index 0000000000..8992456277
--- /dev/null
+++ b/library/cpp/yt/misc/enum_indexed_array.h
@@ -0,0 +1,64 @@
+#pragma once
+
+#include "enum.h"
+
+namespace NYT {
+
+////////////////////////////////////////////////////////////////////////////////
+
+//! A statically sized vector with elements of type |T| indexed by
+//! the items of enumeration type |E|.
+/*!
+ *  By default, valid indexes are in range from the minimum declared value of |E|
+ *  to the maximum declared value of |E|.
+ *
+ *  Items are value-initialized on construction.
+ */
+template <
+    class E,
+    class T,
+    E Min = TEnumTraits<E>::GetMinValue(),
+    E Max = TEnumTraits<E>::GetMaxValue()
+>
+class TEnumIndexedArray
+{
+public:
+    static_assert(Min <= Max);
+
+    using TIndex = E;
+    using TValue = T;
+
+    constexpr TEnumIndexedArray() = default;
+    TEnumIndexedArray(std::initializer_list<std::pair<E, T>> elements);
+
+    constexpr TEnumIndexedArray(const TEnumIndexedArray&) = default;
+    constexpr TEnumIndexedArray(TEnumIndexedArray&&) = default;
+
+    constexpr TEnumIndexedArray& operator=(const TEnumIndexedArray&) = default;
+    constexpr TEnumIndexedArray& operator=(TEnumIndexedArray&&) = default;
+
+    T& operator[] (E index);
+    const T& operator[] (E index) const;
+
+    // STL interop.
+    T* begin();
+    const T* begin() const;
+    T* end();
+    const T* end() const;
+
+    constexpr size_t size() const;
+
+    static bool IsValidIndex(E index);
+
+private:
+    static constexpr size_t Size = static_cast<size_t>(Max) - static_cast<size_t>(Min) + 1;
+    std::array<T, Size> Items_{};
+};
+
+////////////////////////////////////////////////////////////////////////////////
+
+} // namespace NYT
+
+#define ENUM_INDEXED_ARRAY_INL_H_
+#include "enum_indexed_array-inl.h"
+#undef ENUM_INDEXED_ARRAY_INL_H_
diff --git a/library/cpp/yt/misc/unittests/enum_indexed_array_ut.cpp b/library/cpp/yt/misc/unittests/enum_indexed_array_ut.cpp
new file mode 100644
index 0000000000..4111c6f4fe
--- /dev/null
+++ b/library/cpp/yt/misc/unittests/enum_indexed_array_ut.cpp
@@ -0,0 +1,45 @@
+#include <library/cpp/testing/gtest/gtest.h>
+
+#include <library/cpp/yt/misc/enum_indexed_array.h>
+
+namespace NYT {
+namespace {
+
+////////////////////////////////////////////////////////////////////////////////
+
+DEFINE_ENUM(EColor,
+    ((Red)  (10))
+    ((Green)(20))
+    ((Blue) (30))
+);
+
+TEST(TEnumIndexedArrayTest, Size)
+{
+    TEnumIndexedArray<EColor, int> arr;
+    EXPECT_EQ(std::ssize(arr), 21);
+}
+
+TEST(TEnumIndexedArrayTest, IsValidIndex)
+{
+    TEnumIndexedArray<EColor, int> arr;
+    EXPECT_TRUE(arr.IsValidIndex(EColor::Red));
+    EXPECT_TRUE(arr.IsValidIndex(EColor::Green));
+    EXPECT_TRUE(arr.IsValidIndex(EColor::Blue));
+    EXPECT_TRUE(arr.IsValidIndex(static_cast<EColor>(11)));
+    EXPECT_FALSE(arr.IsValidIndex(static_cast<EColor>(9)));
+}
+
+TEST(TEnumIndexedArrayTest, Simple)
+{
+    TEnumIndexedArray<EColor, int> arr;
+    EXPECT_EQ(arr[EColor::Red], 0);
+    arr[EColor::Red] = 1;
+    EXPECT_EQ(arr[EColor::Red], 1);
+    EXPECT_EQ(arr[EColor::Blue], 0);
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+} // namespace
+} // namespace NYT
+
diff --git a/library/cpp/yt/misc/unittests/ya.make b/library/cpp/yt/misc/unittests/ya.make
index 611edd7217..43e9f9b9ff 100644
--- a/library/cpp/yt/misc/unittests/ya.make
+++ b/library/cpp/yt/misc/unittests/ya.make
@@ -4,6 +4,7 @@ INCLUDE(${ARCADIA_ROOT}/library/cpp/yt/ya_cpp.make.inc)
 
 SRCS(
     enum_ut.cpp
+    enum_indexed_array_ut.cpp
     guid_ut.cpp
     preprocessor_ut.cpp
     strong_typedef_ut.cpp
diff --git a/library/cpp/yt/string/format-inl.h b/library/cpp/yt/string/format-inl.h
index c379af712a..c198fdcd65 100644
--- a/library/cpp/yt/string/format-inl.h
+++ b/library/cpp/yt/string/format-inl.h
@@ -12,6 +12,7 @@
 #include <library/cpp/yt/small_containers/compact_vector.h>
 
 #include <library/cpp/yt/misc/enum.h>
+#include <library/cpp/yt/misc/enum_indexed_array.h>
 
 #include <util/system/platform.h>
 
@@ -393,6 +394,27 @@ struct TValueFormatter<THashMultiMap<K, V>>
     }
 };
 
+// TEnumIndexedArray
+template <class E, class T>
+struct TValueFormatter<TEnumIndexedArray<E, T>>
+{
+    static void Do(TStringBuilderBase* builder, const TEnumIndexedArray<E, T>& collection, TStringBuf format)
+    {
+        builder->AppendChar('{');
+        bool firstItem = true;
+        for (const auto& index : TEnumTraits<E>::GetDomainValues()) {
+            if (!firstItem) {
+                builder->AppendString(DefaultJoinToStringDelimiter);
+            }
+            FormatValue(builder, index, format);
+            builder->AppendString(": ");
+            FormatValue(builder, collection[index], format);
+            firstItem = false;
+        }
+        builder->AppendChar('}');
+    }
+};
+
 // TEnumIndexedVector
 template <class E, class T>
 struct TValueFormatter<TEnumIndexedVector<E, T>>
diff --git a/library/cpp/yt/string/unittests/enum_ut.cpp b/library/cpp/yt/string/unittests/enum_ut.cpp
index 4059b3d2e6..91a8b04500 100644
--- a/library/cpp/yt/string/unittests/enum_ut.cpp
+++ b/library/cpp/yt/string/unittests/enum_ut.cpp
@@ -13,6 +13,7 @@ namespace {
 // Some compile-time sanity checks.
 DEFINE_ENUM(ESample, (One)(Two));
 static_assert(TFormatTraits<ESample>::HasCustomFormatValue);
+static_assert(TFormatTraits<TEnumIndexedArray<ESample, int>>::HasCustomFormatValue);
 static_assert(TFormatTraits<TEnumIndexedVector<ESample, int>>::HasCustomFormatValue);
 
 DEFINE_ENUM(EColor,