Matrix: fix crash on reply?

1 files changed, 35 insertions, 15 deletions

diff --git a/src/QuickMedia.cpp b/src/QuickMedia.cpp
index e3653e5..5376535 100644
--- a/src/QuickMedia.cpp
+++ b/src/QuickMedia.cpp
@@ -3446,6 +3446,7 @@ namespace QuickMedia {
         std::unordered_map<std::string, std::shared_ptr<Message>> pending_sent_replies;
 
         // This is needed to keep the message shared ptr alive. TODO: Remove this shit, maybe even use raw pointers.
+        std::unordered_map<std::string, ProvisionalMessage> sent_messages; // |event_id| is always empty in this. Use |message->event_id| instead
         std::optional<ProvisionalMessage> provisional_message;
         MessageQueue<ProvisionalMessage> provisional_message_queue;
         MessageQueue<std::function<ProvisionalMessage()>> post_task_queue;
@@ -3459,6 +3460,33 @@ namespace QuickMedia {
         };
         std::thread post_thread(post_thread_handler);
 
+        auto message_set_replaced_by = [&find_body_item_by_event_id, &tabs, &pending_sent_replies](std::shared_ptr<Message> message) {
+            if(message->related_event_type == RelatedEventType::EDIT) {
+                auto body_item = find_body_item_by_event_id(tabs[MESSAGES_TAB_INDEX].body->items.data(), tabs[MESSAGES_TAB_INDEX].body->items.size(), message->related_event_id);
+                if(body_item) {
+                    Message *orig_message = static_cast<Message*>(body_item->userdata);
+                    orig_message->replaced_by = message;
+                }
+            } else if(message->related_event_type == RelatedEventType::REPLY) {
+                auto pending_sent_reply_it = pending_sent_replies.find(message->transaction_id);
+                if(pending_sent_reply_it != pending_sent_replies.end()) {
+                    pending_sent_reply_it->second->replaced_by = message;
+                    pending_sent_replies.erase(pending_sent_reply_it);
+                }
+            }
+        };
+
+        auto filter_sent_messages = [&sent_messages, &message_set_replaced_by](Messages &messages) {
+            for(auto it = messages.begin(); it != messages.end();) {
+                if(sent_messages.find((*it)->event_id) != sent_messages.end()) {
+                    message_set_replaced_by(*it);
+                    it = messages.erase(it);
+                } else {
+                    ++it;
+                }
+            }
+        };
+
         auto upload_file = [this, &tabs, &current_room](const std::string &filepath) {
             TaskResult post_file_result = run_task_with_loading_screen([this, &current_room, filepath]() {
                 std::string event_id_response;
@@ -3918,22 +3946,10 @@ namespace QuickMedia {
             }
         };
 
-        auto filter_provisional_messages = [&find_body_item_by_event_id, &tabs, &pending_sent_replies](Messages &messages) {
+        auto filter_provisional_messages = [&message_set_replaced_by](Messages &messages) {
             for(auto it = messages.begin(); it != messages.end();) {
                 if(!(*it)->transaction_id.empty()) {
-                    if((*it)->related_event_type == RelatedEventType::EDIT) {
-                        auto body_item = find_body_item_by_event_id(tabs[MESSAGES_TAB_INDEX].body->items.data(), tabs[MESSAGES_TAB_INDEX].body->items.size(), (*it)->related_event_id);
-                        if(body_item) {
-                            Message *orig_message = static_cast<Message*>(body_item->userdata);
-                            orig_message->replaced_by = (*it);
-                        }
-                    } else if((*it)->related_event_type == RelatedEventType::REPLY) {
-                        auto pending_sent_reply_it = pending_sent_replies.find((*it)->transaction_id);
-                        if(pending_sent_reply_it != pending_sent_replies.end()) {
-                            pending_sent_reply_it->second->replaced_by = (*it);
-                            pending_sent_replies.erase(pending_sent_reply_it);
-                        }
-                    }
+                    message_set_replaced_by(*it);
                     it = messages.erase(it);
                 } else {
                     ++it;
@@ -4097,7 +4113,7 @@ namespace QuickMedia {
             }
         };
 
-        auto cleanup_tasks = [&set_read_marker_future, &fetch_message_future, &fetch_users_future, &typing_state_queue, &typing_state_thread, &post_task_queue, &provisional_message_queue, &fetched_messages_set, &pending_sent_replies, &post_thread, &tabs]() {
+        auto cleanup_tasks = [&set_read_marker_future, &fetch_message_future, &fetch_users_future, &typing_state_queue, &typing_state_thread, &post_task_queue, &provisional_message_queue, &fetched_messages_set, &sent_messages, &pending_sent_replies, &post_thread, &tabs]() {
             set_read_marker_future.cancel();
             fetch_message_future.cancel();
             fetch_users_future.cancel();
@@ -4113,6 +4129,7 @@ namespace QuickMedia {
             }
             provisional_message_queue.clear();
             fetched_messages_set.clear();
+            sent_messages.clear();
             pending_sent_replies.clear();
     
             //unreferenced_event_by_room.clear();
@@ -4567,6 +4584,7 @@ namespace QuickMedia {
                 if(!provisional_message->event_id.empty()) {
                     provisional_message->message->event_id = std::move(provisional_message->event_id);
                     provisional_message->body_item->set_description_color(sf::Color::White);
+                    sent_messages[provisional_message->message->event_id] = std::move(provisional_message.value());
                 } else if(provisional_message->body_item) {
                     provisional_message->body_item->set_description("Failed to send: " + provisional_message->body_item->get_description());
                     provisional_message->body_item->set_description_color(sf::Color::Red);
@@ -4598,6 +4616,7 @@ namespace QuickMedia {
                 all_messages.insert(all_messages.end(), new_messages.begin(), new_messages.end());
                 if(new_messages.empty())
                     fetched_enough_messages = true;
+                filter_sent_messages(new_messages);
                 filter_existing_messages(new_messages);
                 fprintf(stderr, "Finished fetching older messages, num new messages: %zu\n", new_messages.size());
                 size_t num_new_messages = new_messages.size();
@@ -4825,6 +4844,7 @@ namespace QuickMedia {
                 }
                 all_messages.insert(all_messages.end(), all_messages_new.begin(), all_messages_new.end());
                 //me = matrix->get_me(current_room);
+                filter_sent_messages(all_messages_new);
                 add_new_messages_to_current_room(all_messages_new);
                 modify_related_messages_in_current_room(all_messages_new);
                 unresolved_reactions.clear();