From b2ad399bc9e8745c68cf60c318fc2d2190bb03d8 Mon Sep 17 00:00:00 2001
From: Arseny Smirnov <arseny30@gmail.com>
Date: Sat, 19 Jan 2019 17:50:20 +0300
Subject: [PATCH] New SetWithPosition implementation

GitOrigin-RevId: baa20eabced4227b065bd74f04a3621b6467a339
---
 td/telegram/SetWithPosition.h | 148 +++++++++++++++++++++++++++++++++-
 test/set_with_position.cpp    |  35 +++++++-
 2 files changed, 179 insertions(+), 4 deletions(-)

diff --git a/td/telegram/SetWithPosition.h b/td/telegram/SetWithPosition.h
index ab0f75d6..cff7449e 100644
--- a/td/telegram/SetWithPosition.h
+++ b/td/telegram/SetWithPosition.h
@@ -5,10 +5,154 @@
 // file LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)
 //
 #pragma once
+#include "td/utils/common.h"
+
+#include <set>
 
 namespace td {
+template <class T>
+class FastSetWithPosition {
+ public:
+  void add(int x) {
+    if (checked_.count(x) != 0) {
+      return;
+    }
+    not_checked_.insert(x);
+  }
+  void remove(int x) {
+    checked_.erase(x);
+    not_checked_.erase(x);
+  }
+  bool has_next() {
+    return !not_checked_.empty();
+  }
+  void reset_position() {
+    not_checked_.insert(checked_.begin(), checked_.end());
+    checked_ = {};
+  }
+
+  T next() {
+    CHECK(has_next());
+    auto res = *not_checked_.begin();
+    not_checked_.erase(not_checked_.begin());
+    checked_.insert(res);
+    return res;
+  }
+
+  void merge(FastSetWithPosition &&other) {
+    if (size() < other.size()) {
+      std::swap(*this, other);
+    }
+    for (auto x : other.checked_) {
+      not_checked_.erase(x);
+      checked_.insert(x);
+    }
+    for (auto x : other.not_checked_) {
+      if (checked_.count(x) != 0) {
+        continue;
+      }
+      not_checked_.insert(x);
+    }
+  }
+  size_t size() const {
+    return checked_.size() + not_checked_.size();
+  }
+
+ private:
+  std::set<T> checked_;
+  std::set<T> not_checked_;
+};
+
 template <class T>
 class SetWithPosition {
+ public:
+  void add(int x) {
+    if (fast_) {
+      fast_->add(x);
+      return;
+    }
+    if (!has_value_) {
+      value_ = x;
+      has_value_ = true;
+      is_cheched_ = false;
+      return;
+    }
+    if (value_ == x) {
+      return;
+    }
+    make_fast();
+    fast_->add(x);
+  }
+  void remove(int x) {
+    if (fast_) {
+      fast_->remove(x);
+      return;
+    }
+    if (has_value_ && value_ == x) {
+      has_value_ = false;
+      is_cheched_ = false;
+    }
+  }
+  bool has_next() {
+    if (fast_) {
+      return fast_->has_next();
+    }
+    return has_value_ && !is_cheched_;
+  }
+  void reset_position() {
+    if (fast_) {
+      fast_->reset_position();
+      return;
+    }
+    is_cheched_ = false;
+  }
+
+  T next() {
+    CHECK(has_next());
+    if (fast_) {
+      return fast_->next();
+    }
+    is_cheched_ = true;
+    return value_;
+  }
+
+  void merge(SetWithPosition &&other) {
+    if (size() < other.size()) {
+      std::swap(*this, other);
+    }
+    if (other.size() == 0) {
+      return;
+    }
+    make_fast();
+    other.make_fast();
+    fast_->merge(std::move(*other.fast_));
+  }
+  size_t size() const {
+    if (fast_) {
+      return fast_->size();
+    }
+    return has_value_;
+  }
+
+ private:
+  T value_;
+  bool has_value_{false};
+  bool is_cheched_{false};
+  unique_ptr<FastSetWithPosition<T>> fast_;
+  void make_fast() {
+    if (fast_) {
+      return;
+    }
+    fast_ = make_unique<FastSetWithPosition<T>>();
+    CHECK(has_value_);
+    fast_->add(value_);
+    if (is_cheched_) {
+      fast_->next();
+    }
+  }
+};
+template <class T>
+class OldSetWithPosition {
  public:
   void add(T value) {
     auto it = std::find(values_.begin(), values_.end(), value);
@@ -37,8 +181,8 @@ class SetWithPosition {
   bool has_next() {
     return pos_ < values_.size();
   }
-  void merge(SetWithPosition &&other) {
-    SetWithPosition res;
+  void merge(OldSetWithPosition &&other) {
+    OldSetWithPosition res;
     for (size_t i = 0; i < pos_; i++) {
       res.add(values_[i]);
     }
diff --git a/test/set_with_position.cpp b/test/set_with_position.cpp
index 211cc5b2..d56ad3c1 100644
--- a/test/set_with_position.cpp
+++ b/test/set_with_position.cpp
@@ -51,6 +51,10 @@ class CheckedSetWithPosition {
   }
 
   void merge(CheckedSetWithPosition &&other) {
+    if (size() < other.size()) {
+      std::swap(*this, other);
+      std::swap(this->s_, other.s_);
+    }
     for (auto x : other.checked_) {
       not_checked_.erase(x);
       checked_.insert(x);
@@ -63,6 +67,9 @@ class CheckedSetWithPosition {
     }
     s_.merge(std::move(other.s_));
   }
+  size_t size() const {
+    return checked_.size() + not_checked_.size();
+  }
 
  private:
   std::set<T> checked_;
@@ -156,12 +163,36 @@ void test_stress() {
     }
   }
 }
+template <template <class> class RawSet>
+void test_speed() {
+  Random::Xorshift128plus rnd(123);
+  using Set = CheckedSetWithPosition<int, RawSet>;
+  std::vector<unique_ptr<Set>> sets(1 << 18);
+  for (size_t i = 0; i < sets.size(); i++) {
+    sets[i] = make_unique<Set>();
+    sets[i]->add(int(i));
+  }
+  for (size_t d = 1; d < sets.size(); d *= 2) {
+    for (size_t i = 0; i < sets.size(); i += 2 * d) {
+      size_t j = i + d;
+      CHECK(j < sets.size());
+      sets[i]->merge(std::move(*sets[j]));
+    }
+  }
+  LOG(ERROR) << sets[0]->size();
+}
 
 TEST(SetWithPosition, hands) {
-  //test_hands<SetWithPosition>();
+  test_hands<FastSetWithPosition>();
+  test_hands<OldSetWithPosition>();
   test_hands<SetWithPosition>();
 }
 TEST(SetWithPosition, stress) {
-  //test_stress<SetWithPosition>();
+  test_stress<FastSetWithPosition>();
+  test_stress<OldSetWithPosition>();
   test_stress<SetWithPosition>();
 }
+TEST(SetWithPosition, speed) {
+  test_speed<FastSetWithPosition>();
+  test_speed<SetWithPosition>();
+}