From d5e163cd9da64260632926137a9ea55be330b7f7 Mon Sep 17 00:00:00 2001
From: Arseny Smirnov <arseny30@gmail.com>
Date: Tue, 8 Feb 2022 20:47:10 +0100
Subject: [PATCH] FlatHashMap: support initializer_list

---
 tdutils/td/utils/FlatHashMap.h | 22 ++++++++++++++++++++--
 tdutils/test/HashSet.cpp       | 31 +++++++++++++++++++------------
 2 files changed, 39 insertions(+), 14 deletions(-)

diff --git a/tdutils/td/utils/FlatHashMap.h b/tdutils/td/utils/FlatHashMap.h
index f56fc09a0..51cd27eef 100644
--- a/tdutils/td/utils/FlatHashMap.h
+++ b/tdutils/td/utils/FlatHashMap.h
@@ -72,7 +72,10 @@ class fixed_vector {
 
 template <class KeyT, class ValueT, class HashT = std::hash<KeyT>>
 class FlatHashMapImpl {
+ public:
   struct Node {
+    using first_type = KeyT;
+    using second_type = ValueT;
     KeyT first{};
     union {
       ValueT second;
@@ -86,6 +89,10 @@ class FlatHashMapImpl {
 
     Node() {
     }
+    Node(KeyT key, ValueT value) : first(std::move(key)) {
+      new (&second) ValueT(std::move(value));
+      CHECK(!empty());
+    }
     ~Node() {
       if (!empty()) {
         second.~ValueT();
@@ -127,10 +134,9 @@ class FlatHashMapImpl {
   using NodeIterator = typename fixed_vector<Node>::iterator;
   using ConstNodeIterator = typename fixed_vector<Node>::const_iterator;
 
- public:
   // define key_type and value_type for benchmarks
   using key_type = KeyT;
-  using value_type = std::pair<const KeyT, ValueT>;
+  using value_type = Node;
 
   struct Iterator {
     using iterator_category = std::bidirectional_iterator_tag;
@@ -218,6 +224,18 @@ class FlatHashMapImpl {
     assign(other.begin(), other.end());
     return *this;
   }
+
+  FlatHashMapImpl(std::initializer_list<Node> nodes) {
+    reserve(nodes.size());
+    for (auto &node : nodes) {
+      size_t bucket = calc_bucket(node.key());
+      while (!nodes_[bucket].empty()) {
+        next_bucket(bucket);
+      }
+      nodes_[bucket].emplace(node.first, node.second);
+    }
+  }
+
   FlatHashMapImpl(FlatHashMapImpl &&other) noexcept : nodes_(std::move(other.nodes_)), used_nodes_(other.used_nodes_) {
     other.used_nodes_ = 0;
   }
diff --git a/tdutils/test/HashSet.cpp b/tdutils/test/HashSet.cpp
index 59132f99c..87757eabb 100644
--- a/tdutils/test/HashSet.cpp
+++ b/tdutils/test/HashSet.cpp
@@ -11,22 +11,29 @@
 #include <array>
 
 TEST(FlatHashMap, basic) {
-  td::FlatHashMap<int, int> map;
-  map[1] = 2;
-  ASSERT_EQ(2, map[1]);
-  ASSERT_EQ(1, map.find(1)->first);
-  ASSERT_EQ(2, map.find(1)->second);
-  // ASSERT_EQ(1, map.find(1)->key());
-  // ASSERT_EQ(2, map.find(1)->value());
-  for (auto &kv : map) {
-    ASSERT_EQ(1, kv.first);
-    ASSERT_EQ(2, kv.second);
+  {
+    td::FlatHashMap<int, int> map;
+    map[1] = 2;
+    ASSERT_EQ(2, map[1]);
+    ASSERT_EQ(1, map.find(1)->first);
+    ASSERT_EQ(2, map.find(1)->second);
+    // ASSERT_EQ(1, map.find(1)->key());
+    // ASSERT_EQ(2, map.find(1)->value());
+    for (auto &kv : map) {
+      ASSERT_EQ(1, kv.first);
+      ASSERT_EQ(2, kv.second);
+    }
+    map.erase(map.find(1));
+    auto map_copy = map;
   }
-  map.erase(map.find(1));
-  auto map_copy = map;
 
   td::FlatHashMap<int, std::array<td::unique_ptr<td::string>, 20>> x;
   auto y = std::move(x);
   x[12];
   x.erase(x.find(12));
+
+  {
+    td::FlatHashMap<int, std::string> map = {{1, "hello"}, {2, "world"}};
+    ASSERT_EQ("hello", map[1]);
+  }
 }