From b8cd0bb597b012602c80226b08858316406c842d Mon Sep 17 00:00:00 2001
From: levlam <levlam@telegram.org>
Date: Wed, 3 May 2023 01:34:58 +0300
Subject: [PATCH] Improve OrderedMessages::erase.

---
 td/telegram/MessagesManager.cpp | 32 +++++---------------------------
 td/telegram/OrderedMessage.cpp  | 21 +++++++++++++++++++--
 td/telegram/OrderedMessage.h    |  2 +-
 3 files changed, 25 insertions(+), 30 deletions(-)

diff --git a/td/telegram/MessagesManager.cpp b/td/telegram/MessagesManager.cpp
index 51179d06a..becb91b1d 100644
--- a/td/telegram/MessagesManager.cpp
+++ b/td/telegram/MessagesManager.cpp
@@ -16330,10 +16330,6 @@ unique_ptr<MessagesManager::Message> MessagesManager::do_delete_message(Dialog *
       << d->being_deleted_message_id << " " << message_id << " " << source;
   d->being_deleted_message_id = message_id;
 
-  const auto message_it = d->ordered_messages.get_iterator(message_id);
-  CHECK(*message_it != nullptr);
-  CHECK((*message_it)->get_message_id() == message_id);
-
   bool need_get_history = false;
   if (!only_from_memory) {
     LOG(INFO) << "Deleting " << full_message_id << " from " << source;
@@ -16343,6 +16339,10 @@ unique_ptr<MessagesManager::Message> MessagesManager::do_delete_message(Dialog *
     delete_active_live_location(d->dialog_id, m);
     remove_message_file_sources(d->dialog_id, m);
 
+    const auto message_it = d->ordered_messages.get_const_iterator(message_id);
+    CHECK(*message_it != nullptr);
+    CHECK((*message_it)->get_message_id() == message_id);
+
     if (message_id == d->last_message_id) {
       if ((*message_it)->have_previous_) {
         auto it = message_it;
@@ -16438,33 +16438,11 @@ unique_ptr<MessagesManager::Message> MessagesManager::do_delete_message(Dialog *
     }
   }
 
-  if ((*message_it)->have_previous_ && (only_from_memory || !(*message_it)->have_next_)) {
-    auto it = message_it;
-    --it;
-    OrderedMessage *prev_m = *it;
-    if (prev_m != nullptr) {
-      prev_m->have_next_ = false;
-    } else {
-      LOG(ERROR) << "Have have_previous is true, but there is no previous for " << full_message_id << " from "
-                 << source;
-    }
-  }
-  if ((*message_it)->have_next_ && (only_from_memory || !(*message_it)->have_previous_)) {
-    auto it = message_it;
-    ++it;
-    OrderedMessage *next_m = *it;
-    if (next_m != nullptr) {
-      next_m->have_previous_ = false;
-    } else {
-      LOG(ERROR) << "Have have_next is true, but there is no next for " << full_message_id << " from " << source;
-    }
-  }
-
   auto result = std::move(d->messages[message_id]);
   CHECK(m == result.get());
   d->messages.erase(message_id);
 
-  d->ordered_messages.erase(message_id);
+  d->ordered_messages.erase(message_id, only_from_memory);
 
   d->being_deleted_message_id = MessageId();
 
diff --git a/td/telegram/OrderedMessage.cpp b/td/telegram/OrderedMessage.cpp
index 91fe00395..8a3d14aba 100644
--- a/td/telegram/OrderedMessage.cpp
+++ b/td/telegram/OrderedMessage.cpp
@@ -59,7 +59,7 @@ void OrderedMessages::insert(MessageId message_id, bool was_auto_attached, bool
   }
 }
 
-void OrderedMessages::erase(MessageId message_id) {
+void OrderedMessages::erase(MessageId message_id, bool only_from_memory) {
   unique_ptr<OrderedMessage> *v = &messages_;
   while (*v != nullptr) {
     if ((*v)->message_id_.get() < message_id.get()) {
@@ -71,8 +71,25 @@ void OrderedMessages::erase(MessageId message_id) {
     }
   }
 
+  CHECK(*v != nullptr);
+  if ((*v)->have_previous_ && (only_from_memory || !(*v)->have_next_)) {
+    Iterator it(messages_.get(), message_id);
+    CHECK(*it == v->get());
+    --it;
+    OrderedMessage *prev_m = *it;
+    CHECK(prev_m != nullptr);
+    prev_m->have_next_ = false;
+  }
+  if ((*v)->have_next_ && (only_from_memory || !(*v)->have_previous_)) {
+    Iterator it(messages_.get(), message_id);
+    CHECK(*it == v->get());
+    ++it;
+    OrderedMessage *next_m = *it;
+    CHECK(next_m != nullptr);
+    next_m->have_previous_ = false;
+  }
+
   unique_ptr<OrderedMessage> result = std::move(*v);
-  CHECK(result != nullptr);
   unique_ptr<OrderedMessage> left = std::move(result->left_);
   unique_ptr<OrderedMessage> right = std::move(result->right_);
 
diff --git a/td/telegram/OrderedMessage.h b/td/telegram/OrderedMessage.h
index f94a1220e..dd8ef0eb7 100644
--- a/td/telegram/OrderedMessage.h
+++ b/td/telegram/OrderedMessage.h
@@ -171,7 +171,7 @@ class OrderedMessages {
 
   void insert(MessageId message_id, bool was_auto_attached, bool have_previous, bool have_next);
 
-  void erase(MessageId message_id);
+  void erase(MessageId message_id, bool only_from_memory);
 
   struct AttachInfo {
     bool have_previous_ = false;