Improve OrderedMessages::erase.

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();
 

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_);
 

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;