Intermediate changes

3 files changed, 91 insertions, 3 deletions

diff --git a/library/cpp/yt/small_containers/compact_flat_map-inl.h b/library/cpp/yt/small_containers/compact_flat_map-inl.h
index 74557f2467..740a7c2df9 100644
--- a/library/cpp/yt/small_containers/compact_flat_map-inl.h
+++ b/library/cpp/yt/small_containers/compact_flat_map-inl.h
@@ -210,6 +210,30 @@ TCompactFlatMap<K, V, N>::equal_range(const K& k) const
     return result;
 }
 
+template <class K, class V, size_t N>
+typename TCompactFlatMap<K, V, N>::const_iterator TCompactFlatMap<K, V, N>::lower_bound(const K& k) const
+{
+    return std::lower_bound(Storage_.begin(), Storage_.end(), k, TKeyComparer());
+}
+
+template <class K, class V, size_t N>
+typename TCompactFlatMap<K, V, N>::iterator TCompactFlatMap<K, V, N>::lower_bound(const K& k)
+{
+    return std::lower_bound(Storage_.begin(), Storage_.end(), k, TKeyComparer());
+}
+
+template <class K, class V, size_t N>
+typename TCompactFlatMap<K, V, N>::const_iterator TCompactFlatMap<K, V, N>::upper_bound(const K& k) const
+{
+    return std::upper_bound(Storage_.begin(), Storage_.end(), k, TKeyComparer());
+}
+
+template <class K, class V, size_t N>
+typename TCompactFlatMap<K, V, N>::iterator TCompactFlatMap<K, V, N>::upper_bound(const K& k)
+{
+    return std::upper_bound(Storage_.begin(), Storage_.end(), k, TKeyComparer());
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 
 } // namespace NYT
diff --git a/library/cpp/yt/small_containers/compact_flat_map.h b/library/cpp/yt/small_containers/compact_flat_map.h
index 6263a3fada..afb229c06d 100644
--- a/library/cpp/yt/small_containers/compact_flat_map.h
+++ b/library/cpp/yt/small_containers/compact_flat_map.h
@@ -81,6 +81,13 @@ public:
     iterator find(const K& k);
     const_iterator find(const K& k) const;
 
+    iterator lower_bound(const K& k);
+    const_iterator lower_bound(const K& k) const;
+    iterator upper_bound(const K& k);
+    const_iterator upper_bound(const K& k) const;
+    std::pair<iterator, iterator> equal_range(const K& k);
+    std::pair<const_iterator, const_iterator> equal_range(const K& k) const;
+
     bool contains(const K& k) const;
 
     std::pair<iterator, bool> insert(const value_type& value);
@@ -101,9 +108,6 @@ public:
 private:
     TStorage Storage_;
 
-    std::pair<iterator, iterator> equal_range(const K& k);
-    std::pair<const_iterator, const_iterator> equal_range(const K& k) const;
-
     template <class TArg>
     std::pair<iterator, bool> do_insert(TArg&& value);
 };
diff --git a/library/cpp/yt/small_containers/unittests/compact_flat_map_ut.cpp b/library/cpp/yt/small_containers/unittests/compact_flat_map_ut.cpp
index 3f738e78b4..a946245736 100644
--- a/library/cpp/yt/small_containers/unittests/compact_flat_map_ut.cpp
+++ b/library/cpp/yt/small_containers/unittests/compact_flat_map_ut.cpp
@@ -219,6 +219,66 @@ TEST(TCompactFlatMapTest, GrowShrinkGrow)
     // Must not crash or trigger asan.
 }
 
+TEST(TCompactFlatMapTest, LowerBound)
+{
+    TMap m;
+    EXPECT_EQ(m.lower_bound("a"), m.end());
+
+    m.emplace("b", "value");
+    EXPECT_EQ(m.lower_bound("a")->second, "value");
+    EXPECT_EQ(m.lower_bound("b")->second, "value");
+
+    m.emplace("c", "value2");
+
+    // Grows here.
+    m.emplace("d", "value3");
+    EXPECT_EQ(m.lower_bound("a")->second, "value");
+    EXPECT_EQ(m.lower_bound("e"), m.end());
+}
+
+TEST(TCompactFlatMapTest, UpperBound)
+{
+    using TKeyValue = std::pair<std::string, std::string>;
+    TMap m;
+    EXPECT_EQ(m.upper_bound("a"), m.end());
+
+    m.emplace("b", "value");
+    EXPECT_EQ(m.upper_bound("a")->second, "value");
+    EXPECT_EQ(m.upper_bound("b"), m.end());
+
+    m.emplace("c", "value2");
+
+    // Grows here.
+    m.emplace("d", "value3");
+
+    EXPECT_EQ(*m.upper_bound("a"), TKeyValue("b", "value"));
+    EXPECT_EQ(*m.upper_bound("b"), TKeyValue("c", "value2"));
+    EXPECT_EQ(m.upper_bound("d"), m.end());
+}
+
+TEST(TCompactFlatMapTest, EqualRange)
+{
+    TMap m;
+    EXPECT_EQ(m.equal_range("a"), std::make_pair(m.end(), m.end()));
+
+    m.emplace("b", "value-b");
+    EXPECT_EQ(m.equal_range("a"), std::make_pair(m.begin(), m.begin()));
+    EXPECT_EQ(m.equal_range("b"), std::make_pair(m.begin(), m.end()));
+    EXPECT_EQ(m.equal_range("c"), std::make_pair(m.end(), m.end()));
+
+    m.emplace("c", "value-c");
+    m.emplace("d", "value-d");
+
+    auto it = m.begin();
+    EXPECT_EQ(m.equal_range("a"), std::make_pair(it, it));
+    EXPECT_EQ(m.equal_range("b"), std::make_pair(it, std::next(it)));
+    ++it;
+    EXPECT_EQ(m.equal_range("c"), std::make_pair(it, std::next(it)));
+    ++it;
+    EXPECT_EQ(m.equal_range("d"), std::make_pair(it, std::next(it)));
+    EXPECT_EQ(m.equal_range("e"), std::make_pair(m.end(), m.end()));
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 
 } // namespace