YT-19441: Support allocator in TCompactSet.

2 files changed, 42 insertions, 34 deletions

diff --git a/library/cpp/yt/small_containers/compact_set-inl.h b/library/cpp/yt/small_containers/compact_set-inl.h
index 24f74e5101..e93d11f126 100644
--- a/library/cpp/yt/small_containers/compact_set-inl.h
+++ b/library/cpp/yt/small_containers/compact_set-inl.h
@@ -10,11 +10,11 @@ namespace NYT {
 
 ////////////////////////////////////////////////////////////////////////////////
 
-template <typename T, size_t N,  typename C>
-class TCompactSet<T, N, C>::const_iterator
+template <typename T, size_t N,  typename C, typename A>
+class TCompactSet<T, N, C, A>::const_iterator
 {
 private:
-    friend class TCompactSet<T, N, C>;
+    friend class TCompactSet<T, N, C, A>;
 
     union
     {
@@ -202,26 +202,31 @@ public:
 
 ////////////////////////////////////////////////////////////////////////////////
 
-template <typename T, size_t N,  typename C>
-bool TCompactSet<T, N, C>::empty() const
+template <typename T, size_t N,  typename C, typename A>
+TCompactSet<T, N, C, A>::TCompactSet(const A& allocator)
+    : Set_(allocator)
+{ }
+
+template <typename T, size_t N,  typename C, typename A>
+bool TCompactSet<T, N, C, A>::empty() const
 {
     return Vector_.empty() && Set_.empty();
 }
 
-template <typename T, size_t N,  typename C>
-typename TCompactSet<T, N, C>::size_type TCompactSet<T, N, C>::size() const
+template <typename T, size_t N,  typename C, typename A>
+typename TCompactSet<T, N, C, A>::size_type TCompactSet<T, N, C, A>::size() const
 {
     return is_small() ? Vector_.size() : Set_.size();
 }
 
-template <typename T, size_t N,  typename C>
-const T& TCompactSet<T, N, C>::front() const
+template <typename T, size_t N,  typename C, typename A>
+const T& TCompactSet<T, N, C, A>::front() const
 {
     return is_small() ? Vector_.front() : *Set_.begin();
 }
 
-template <typename T, size_t N,  typename C>
-typename TCompactSet<T, N, C>::size_type TCompactSet<T, N, C>::count(const T& v) const
+template <typename T, size_t N,  typename C, typename A>
+typename TCompactSet<T, N, C, A>::size_type TCompactSet<T, N, C, A>::count(const T& v) const
 {
     if (is_small()) {
         return std::binary_search(Vector_.begin(), Vector_.end(), v, C()) ? 1 : 0;
@@ -230,14 +235,14 @@ typename TCompactSet<T, N, C>::size_type TCompactSet<T, N, C>::count(const T& v)
     }
 }
 
-template <typename T, size_t N,  typename C>
-bool TCompactSet<T, N, C>::contains(const T& v) const
+template <typename T, size_t N,  typename C, typename A>
+bool TCompactSet<T, N, C, A>::contains(const T& v) const
 {
     return count(v) == 1;
 }
 
-template <typename T, size_t N,  typename C>
-std::pair<typename TCompactSet<T, N, C>::const_iterator, bool> TCompactSet<T, N, C>::insert(const T& v)
+template <typename T, size_t N,  typename C, typename A>
+std::pair<typename TCompactSet<T, N, C, A>::const_iterator, bool> TCompactSet<T, N, C, A>::insert(const T& v)
 {
     if (!is_small()) {
         auto [it, inserted] = Set_.insert(v);
@@ -262,17 +267,17 @@ std::pair<typename TCompactSet<T, N, C>::const_iterator, bool> TCompactSet<T, N,
     return {const_iterator(std::move(newIt)), true};
 }
 
-template <typename T, size_t N,  typename C>
+template <typename T, size_t N,  typename C, typename A>
 template <typename TIter>
-void TCompactSet<T, N, C>::insert(TIter i, TIter e)
+void TCompactSet<T, N, C, A>::insert(TIter i, TIter e)
 {
     for (; i != e; ++i) {
         insert(*i);
     }
 }
 
-template <typename T, size_t N,  typename C>
-bool TCompactSet<T, N, C>::erase(const T& v)
+template <typename T, size_t N,  typename C, typename A>
+bool TCompactSet<T, N, C, A>::erase(const T& v)
 {
     if (!is_small()) {
         return Set_.erase(v);
@@ -287,39 +292,39 @@ bool TCompactSet<T, N, C>::erase(const T& v)
     }
 }
 
-template <typename T, size_t N,  typename C>
-void TCompactSet<T, N, C>::clear()
+template <typename T, size_t N,  typename C, typename A>
+void TCompactSet<T, N, C, A>::clear()
 {
     Vector_.clear();
     Set_.clear();
 }
 
-template <typename T, size_t N,  typename C>
-typename TCompactSet<T, N, C>::const_iterator TCompactSet<T, N, C>::begin() const
+template <typename T, size_t N,  typename C, typename A>
+typename TCompactSet<T, N, C, A>::const_iterator TCompactSet<T, N, C, A>::begin() const
 {
     return is_small() ? const_iterator(Vector_.begin()) : const_iterator(Set_.begin());
 }
 
-template <typename T, size_t N,  typename C>
-typename TCompactSet<T, N, C>::const_iterator TCompactSet<T, N, C>::cbegin() const
+template <typename T, size_t N,  typename C, typename A>
+typename TCompactSet<T, N, C, A>::const_iterator TCompactSet<T, N, C, A>::cbegin() const
 {
     return begin();
 }
 
-template <typename T, size_t N,  typename C>
-typename TCompactSet<T, N, C>::const_iterator TCompactSet<T, N, C>::end() const
+template <typename T, size_t N,  typename C, typename A>
+typename TCompactSet<T, N, C, A>::const_iterator TCompactSet<T, N, C, A>::end() const
 {
     return is_small() ? const_iterator(Vector_.end()) : const_iterator(Set_.end());
 }
 
-template <typename T, size_t N,  typename C>
-typename TCompactSet<T, N, C>::const_iterator TCompactSet<T, N, C>::cend() const
+template <typename T, size_t N,  typename C, typename A>
+typename TCompactSet<T, N, C, A>::const_iterator TCompactSet<T, N, C, A>::cend() const
 {
     return end();
 }
 
-template <typename T, size_t N,  typename C>
-bool TCompactSet<T, N, C>::is_small() const
+template <typename T, size_t N,  typename C, typename A>
+bool TCompactSet<T, N, C, A>::is_small() const
 {
     return Set_.empty();
 }
diff --git a/library/cpp/yt/small_containers/compact_set.h b/library/cpp/yt/small_containers/compact_set.h
index b325e42743..ca5e3c6efd 100644
--- a/library/cpp/yt/small_containers/compact_set.h
+++ b/library/cpp/yt/small_containers/compact_set.h
@@ -34,7 +34,7 @@ namespace NYT {
  *
  *  Note that any modification of the set may invalidate *all* iterators.
  */
-template <typename T, size_t N, typename C = std::less<T>>
+template <typename T, size_t N, typename C = std::less<T>, typename A = std::allocator<T>>
 class TCompactSet
 {
 public:
@@ -42,6 +42,9 @@ public:
     using size_type = std::size_t;
     using key_type = T;
 
+    TCompactSet() = default;
+    TCompactSet(const A& allocator);
+
     [[nodiscard]] bool empty() const;
 
     size_type size() const;
@@ -72,9 +75,9 @@ private:
     // reach its 'large' stage) to avoid calling the default ctors of elements
     // we will never use.
     TCompactVector<T, N> Vector_;
-    std::set<T, C> Set_;
+    std::set<T, C, A> Set_;
 
-    using TSetConstIterator = typename std::set<T, C>::const_iterator;
+    using TSetConstIterator = typename std::set<T, C, A>::const_iterator;
     using TVectorConstIterator = typename TCompactVector<T, N>::const_iterator;
 
     bool is_small() const;