From 43ea7be0fbca64995eca3c902a4c77e6d44c2de2 Mon Sep 17 00:00:00 2001
From: levlam <levlam@telegram.org>
Date: Tue, 19 Apr 2022 02:33:12 +0300
Subject: [PATCH] Add fast-moved implementation for big HashSet nodes.

---
 tdutils/td/utils/FlatHashTable.h | 12 ++---
 tdutils/td/utils/MapNode.h       | 75 +++++++++++++++++++++++++++++++-
 tdutils/td/utils/SetNode.h       | 61 +++++++++++++++++++++++++-
 3 files changed, 140 insertions(+), 8 deletions(-)

diff --git a/tdutils/td/utils/FlatHashTable.h b/tdutils/td/utils/FlatHashTable.h
index 101d12dc0..2048ea05d 100644
--- a/tdutils/td/utils/FlatHashTable.h
+++ b/tdutils/td/utils/FlatHashTable.h
@@ -316,9 +316,6 @@ class FlatHashTable {
     auto bucket = calc_bucket(key);
     while (true) {
       auto &node = nodes_[bucket];
-      if (EqT()(node.key(), key)) {
-        return {NodePointer(&node), false};
-      }
       if (node.empty()) {
         if (unlikely(used_node_count_ * 5 >= bucket_count_mask_ * 3)) {
           resize(2 * bucket_count_);
@@ -331,6 +328,9 @@ class FlatHashTable {
         used_node_count_++;
         return {NodePointer(&node), true};
       }
+      if (EqT()(node.key(), key)) {
+        return {NodePointer(&node), false};
+      }
       next_bucket(bucket);
     }
   }
@@ -477,12 +477,12 @@ class FlatHashTable {
     auto bucket = calc_bucket(key);
     while (true) {
       auto &node = nodes_[bucket];
-      if (EqT()(node.key(), key)) {
-        return &node;
-      }
       if (node.empty()) {
         return nullptr;
       }
+      if (EqT()(node.key(), key)) {
+        return &node;
+      }
       next_bucket(bucket);
     }
   }
diff --git a/tdutils/td/utils/MapNode.h b/tdutils/td/utils/MapNode.h
index 3304fcd56..806abe0a5 100644
--- a/tdutils/td/utils/MapNode.h
+++ b/tdutils/td/utils/MapNode.h
@@ -10,11 +10,12 @@
 #include "td/utils/HashTableUtils.h"
 
 #include <new>
+#include <type_traits>
 #include <utility>
 
 namespace td {
 
-template <class KeyT, class ValueT>
+template <class KeyT, class ValueT, class Enable = void>
 struct MapNode {
   using first_type = KeyT;
   using second_type = ValueT;
@@ -90,4 +91,76 @@ struct MapNode {
   }
 };
 
+template <class KeyT, class ValueT>
+struct MapNode<KeyT, ValueT, typename std::enable_if_t<(sizeof(KeyT) + sizeof(ValueT) > 6 * sizeof(void *))>> {
+  struct Impl {
+    using first_type = KeyT;
+    using second_type = ValueT;
+
+    KeyT first{};
+    union {
+      ValueT second;
+    };
+
+    template <class InputKeyT, class... ArgsT>
+    Impl(InputKeyT &&key, ArgsT &&...args) : first(std::forward<InputKeyT>(key)) {
+      new (&second) ValueT(std::forward<ArgsT>(args)...);
+      DCHECK(!is_hash_table_key_empty(first));
+    }
+    Impl(const Impl &other) = delete;
+    Impl &operator=(const Impl &other) = delete;
+    Impl(Impl &&other) = delete;
+    void operator=(Impl &&other) = delete;
+    ~Impl() {
+      second.~ValueT();
+    }
+  };
+
+  using first_type = KeyT;
+  using second_type = ValueT;
+  using public_key_type = KeyT;
+  using public_type = Impl;
+
+  unique_ptr<Impl> impl_;
+
+  const KeyT &key() const {
+    DCHECK(!empty());
+    return impl_->first;
+  }
+
+  Impl &get_public() {
+    return *impl_;
+  }
+
+  const Impl &get_public() const {
+    return *impl_;
+  }
+
+  MapNode() {
+  }
+  MapNode(KeyT key, ValueT value) : impl_(td::make_unique<Impl>(std::move(key), std::move(value))) {
+  }
+
+  void copy_from(const MapNode &other) {
+    DCHECK(empty());
+    DCHECK(!other.empty());
+    impl_ = td::make_unique<Impl>(other.impl_->first, other.impl_->second);
+  }
+
+  bool empty() const {
+    return impl_ == nullptr;
+  }
+
+  void clear() {
+    DCHECK(!empty());
+    impl_ = nullptr;
+  }
+
+  template <class... ArgsT>
+  void emplace(KeyT key, ArgsT &&...args) {
+    DCHECK(empty());
+    impl_ = td::make_unique<Impl>(std::move(key), std::forward<ArgsT>(args)...);
+  }
+};
+
 }  // namespace td
diff --git a/tdutils/td/utils/SetNode.h b/tdutils/td/utils/SetNode.h
index 3d1ca522a..f30a93bc8 100644
--- a/tdutils/td/utils/SetNode.h
+++ b/tdutils/td/utils/SetNode.h
@@ -11,7 +11,7 @@
 
 namespace td {
 
-template <class KeyT>
+template <class KeyT, class Enable = void>
 struct SetNode {
   using public_key_type = KeyT;
   using public_type = const KeyT;
@@ -63,4 +63,63 @@ struct SetNode {
   }
 };
 
+template <class KeyT>
+struct SetNode<KeyT, typename std::enable_if_t<(sizeof(KeyT) > 6 * sizeof(void *))>> {
+  struct Impl {
+    using second_type = KeyT;
+
+    KeyT first{};
+
+    template <class InputKeyT>
+    Impl(InputKeyT &&key) : first(std::forward<InputKeyT>(key)) {
+      DCHECK(!is_hash_table_key_empty(first));
+    }
+    Impl(const Impl &other) = delete;
+    Impl &operator=(const Impl &other) = delete;
+    Impl(Impl &&other) = delete;
+    void operator=(Impl &&other) = delete;
+  };
+
+  using public_key_type = KeyT;
+  using public_type = const KeyT;
+  using second_type = KeyT;  // TODO: remove second_type?
+
+  unique_ptr<Impl> impl_;
+
+  const KeyT &key() const {
+    DCHECK(!empty());
+    return impl_->first;
+  }
+
+  const KeyT &get_public() {
+    DCHECK(!empty());
+    return impl_->first;
+  }
+
+  SetNode() {
+  }
+  explicit SetNode(KeyT key) : impl_(td::make_unique<Impl>(std::move(key))) {
+  }
+
+  void copy_from(const SetNode &other) {
+    DCHECK(empty());
+    impl_ = td::make_unique<Impl>(other.impl_->first);
+    DCHECK(!empty());
+  }
+
+  bool empty() const {
+    return impl_ == nullptr;
+  }
+
+  void clear() {
+    DCHECK(!empty());
+    impl_ = nullptr;
+  }
+
+  void emplace(KeyT key) {
+    DCHECK(empty());
+    impl_ = td::make_unique<Impl>(std::move(key));
+  }
+};
+
 }  // namespace td