From e0bee7757baf68d47758b709a597ac25cb5c2992 Mon Sep 17 00:00:00 2001
From: levlam <levlam@telegram.org>
Date: Wed, 12 Apr 2023 01:33:48 +0300
Subject: [PATCH] Improve get_message_user_ids and get_message_channel_ids.

---
 td/telegram/MessageContent.cpp  |  9 +++++++++
 td/telegram/MessageContent.h    |  2 ++
 td/telegram/MessagesManager.cpp | 25 +++++++++++++++++++++++++
 3 files changed, 36 insertions(+)

diff --git a/td/telegram/MessageContent.cpp b/td/telegram/MessageContent.cpp
index 0de32ba6d..95d038989 100644
--- a/td/telegram/MessageContent.cpp
+++ b/td/telegram/MessageContent.cpp
@@ -3291,6 +3291,15 @@ std::pair<InputGroupCallId, bool> get_message_content_group_call_info(const Mess
   return {m->input_group_call_id, m->duration >= 0};
 }
 
+UserId get_message_content_contact_user_id(const MessageContent *content) {
+  switch (content->get_type()) {
+    case MessageContentType::Contact:
+      return static_cast<const MessageContact *>(content)->contact.get_user_id();
+    default:
+      return UserId();
+  }
+}
+
 vector<UserId> get_message_content_added_user_ids(const MessageContent *content) {
   CHECK(content->get_type() == MessageContentType::ChatAddUsers);
   return static_cast<const MessageChatAddUsers *>(content)->user_ids;
diff --git a/td/telegram/MessageContent.h b/td/telegram/MessageContent.h
index 176aef9ba..cb743ff9e 100644
--- a/td/telegram/MessageContent.h
+++ b/td/telegram/MessageContent.h
@@ -144,6 +144,8 @@ FullMessageId get_message_content_replied_message_id(DialogId dialog_id, const M
 
 std::pair<InputGroupCallId, bool> get_message_content_group_call_info(const MessageContent *content);
 
+UserId get_message_content_contact_user_id(const MessageContent *content);
+
 vector<UserId> get_message_content_added_user_ids(const MessageContent *content);
 
 UserId get_message_content_deleted_user_id(const MessageContent *content);
diff --git a/td/telegram/MessagesManager.cpp b/td/telegram/MessagesManager.cpp
index 424debdae..9f72f4dae 100644
--- a/td/telegram/MessagesManager.cpp
+++ b/td/telegram/MessagesManager.cpp
@@ -13231,6 +13231,23 @@ vector<UserId> MessagesManager::get_message_user_ids(const Message *m) {
   if (m->sender_user_id.is_valid()) {
     user_ids.push_back(m->sender_user_id);
   }
+  if (m->via_bot_user_id.is_valid()) {
+    user_ids.push_back(m->via_bot_user_id);
+  }
+  if (m->forward_info != nullptr && m->forward_info->sender_user_id.is_valid()) {
+    user_ids.push_back(m->forward_info->sender_user_id);
+  }
+  if (m->content->get_type() == MessageContentType::ChatAddUsers) {
+    append(user_ids, get_message_content_added_user_ids(m->content.get()));
+  }
+  auto deleted_user_id = get_message_content_deleted_user_id(m->content.get());
+  if (deleted_user_id.is_valid()) {
+    user_ids.push_back(deleted_user_id);
+  }
+  auto contact_user_id = get_message_content_contact_user_id(m->content.get());
+  if (contact_user_id.is_valid()) {
+    user_ids.push_back(contact_user_id);
+  }
   return user_ids;
 }
 
@@ -13239,6 +13256,14 @@ vector<ChannelId> MessagesManager::get_message_channel_ids(const Message *m) {
   if (m->sender_dialog_id.is_valid() && m->sender_dialog_id.get_type() == DialogType::Channel) {
     channel_ids.push_back(m->sender_dialog_id.get_channel_id());
   }
+  if (m->forward_info != nullptr && m->forward_info->sender_dialog_id.is_valid() &&
+      m->forward_info->sender_dialog_id.get_type() == DialogType::Channel) {
+    channel_ids.push_back(m->forward_info->sender_dialog_id.get_channel_id());
+  }
+  if (m->forward_info != nullptr && m->forward_info->from_dialog_id.is_valid() &&
+      m->forward_info->from_dialog_id.get_type() == DialogType::Channel) {
+    channel_ids.push_back(m->forward_info->from_dialog_id.get_channel_id());
+  }
   return channel_ids;
 }