Redesigned the enumerator

td/telegram/files/FileManager.cpp

@@ -1243,8 +1243,8 @@ Result<FileId> FileManager::register_file(FileData &&data, FileLocationSource fi
   int32 remote_key = 0;
   if (file_view.has_remote_location()) {
     RemoteInfo info{file_view.remote_location(), file_location_source, file_id};
-    remote_key = remote_location_info_.add(info);
-    auto &stored_info = remote_location_info_.get(remote_key);
+    RemoteInfo stored_info;
+    std::tie(remote_key, stored_info) = remote_location_info_.add_and_get(info);
     if (stored_info.file_id_ == file_id) {
       get_file_id_info(file_id)->pin_flag_ = true;
       new_remote = true;
@@ -3068,8 +3068,9 @@ Result<FileId> FileManager::check_input_file_id(FileType type, Result<FileId> re
   int32 remote_id = file_id.get_remote();
   if (remote_id == 0) {
     RemoteInfo info{file_view.remote_location(), FileLocationSource::FromUser, file_id};
-    remote_id = remote_location_info_.add(info);
-    if (remote_location_info_.get(remote_id).file_id_ == file_id) {
+    RemoteInfo stored_info;
+    std::tie(remote_id, stored_info) = remote_location_info_.add_and_get(info);
+    if (stored_info.file_id_ == file_id) {
       get_file_id_info(file_id)->pin_flag_ = true;
     }
   }
@@ -4059,20 +4060,27 @@ void FileManager::memory_cleanup() {
 
   /* DESTROY INVALID remote_location_info_ */
   if (true) {
-    std::set<int32> empty_remote_ids = {};
-    for (auto empty_remote_id : remote_location_info_.get_empty_id()) {
-      empty_remote_ids.insert(empty_remote_id);
+    remote_location_info_.lock_access_mutex();
+
+    std::unordered_map<int32, RemoteInfo> old_remote_info = {};
+    {
+      auto map = remote_location_info_.get_map();
+
+      auto mapSize = map.size();
+      auto mapMaxSize = map.max_size();
+
+      auto it = map.begin();
+      while (it != map.end()) {
+        old_remote_info[it->second] = it->first;
+        it++;
+      }
     }
-    auto map = remote_location_info_.get_map();
 
-    auto emptyIdsSize = empty_remote_ids.size();
-    auto mapSize = map.size();
-
-    auto it = map.begin();
-    while (it != map.end() && (empty_remote_ids.find(it->second) == empty_remote_ids.end())) {
+    auto it = old_remote_info.begin();
+    while (it != old_remote_info.end()) {
       auto is_invalid = false;
 
-      auto find_file = file_id_info_.find(it->first.file_id_.fast_get());
+      auto find_file = file_id_info_.find(it->second.file_id_.fast_get());
       if (find_file != file_id_info_.end()) {
         auto &file = find_file->second;
         auto find_file_node = file_nodes_.find(file.node_id_);
@@ -4085,10 +4093,12 @@ void FileManager::memory_cleanup() {
       }
 
       if (is_invalid) {
-        remote_location_info_.erase(it->second);
+        remote_location_info_.erase_unsafe(it->first);
       }
       it++;
     }
+
+    remote_location_info_.unlock_access_mutex();
   }
 
   /* DESTROY NULL file_id_info_ */

tdutils/td/utils/Enumerator.h

@@ -11,6 +11,7 @@
 #include <limits>
 #include <map>
 #include <tuple>
+#include <shared_mutex>
 
 namespace td {
 
@@ -23,44 +24,74 @@ class Enumerator {
     return map_;
   }
 
-  std::pair<Key, bool> next() {
-    if (!empty_id_.empty()) {
-      auto res = empty_id_.back();
-      empty_id_.pop_back();
-      return std::make_pair(res, true);
-    }
-
-    return std::make_pair((Key) (arr_.size() + 1), false);
+  void lock_access_mutex() const {
+    access_mutex.lock();
   }
 
-  void erase(Key key_y) {
-    empty_id_.push_back(key_y);
+  void unlock_access_mutex() const {
+    access_mutex.unlock();
+  }
+
+  /**
+   *
+   * @return true if the key is new
+   */
+  Key next() {
+    auto id = next_id++;
+
+    return id;
+  }
+
+  void erase_unsafe(Key key_y) {
+    auto find_val = arr_.find(key_y);
+    if (find_val != arr_.end()) {
+      // Erase this
+      map_.erase(find_val->second);
+      // TODO: Not sure about erasing this, instead
+      arr_.erase(key_y);
+    }
   }
 
   Key add(ValueT v) {
-    auto next_id = next();
+    std::lock_guard<std::shared_timed_mutex> writerLock(access_mutex);
+
+    return add_internal(v);
+  }
+
+  Key add_internal(ValueT v) {
+    auto id = next();
     bool was_inserted;
     decltype(map_.begin()) it;
-    std::tie(it, was_inserted) = map_.emplace(std::move(v), next_id.first);
-    if (was_inserted && next_id.second) {
-      arr_[next_id.first - 1] = &it->first;
-    } else if (was_inserted) {
-      arr_.push_back(&it->first);
-    } else if (next_id.second) {
-      empty_id_.push_back(next_id.first);
+    std::tie(it, was_inserted) = map_.emplace(std::move(v), id);
+    if (was_inserted) {
+      arr_[id] = it->first;
     }
     return it->second;
   }
 
   const ValueT &get(Key key) const {
-    auto pos = static_cast<Key>(key - 1);
-    return *arr_[pos];
+    std::shared_lock<std::shared_timed_mutex> readerLock(access_mutex);
+
+    return get_internal(key);
+  }
+
+  const ValueT &get_internal(Key key) const {
+    return arr_.at(key);
+  }
+
+  std::pair<Key,ValueT> add_and_get(ValueT v) {
+    std::lock_guard<std::shared_timed_mutex> writerLock(access_mutex);
+
+    auto remote_key = add_internal(v);
+    auto &stored_info = get_internal(remote_key);
+    return std::make_pair(remote_key, stored_info);
   }
 
  private:
-  std::vector<Key> empty_id_;
+  mutable int32 next_id = 1;
   std::map<ValueT, int32> map_;
-  std::vector<const ValueT *> arr_;
+  std::unordered_map<int32, ValueT> arr_;
+  mutable std::shared_timed_mutex access_mutex;
 };
 
 }  // namespace td