From 996869d4ca83aa59b9621b0c46bc45d3d808fad9 Mon Sep 17 00:00:00 2001
From: levlam <levlam@telegram.org>
Date: Thu, 31 Jan 2019 00:37:38 +0300
Subject: [PATCH] Implement store/parse FileSourceId.

GitOrigin-RevId: ff0d80bf2a8d479dae0cdca2c806a9731ff499d9
---
 CMakeLists.txt                       |  3 +-
 td/telegram/AnimationsManager.cpp    | 12 ++--
 td/telegram/AnimationsManager.h      |  2 +
 td/telegram/ContactsManager.cpp      | 78 +++++++++++++++++++++++--
 td/telegram/ContactsManager.h        | 11 ++++
 td/telegram/FileReferenceManager.cpp | 10 ++--
 td/telegram/FileReferenceManager.h   | 13 ++++-
 td/telegram/FileReferenceManager.hpp | 86 ++++++++++++++++++++++++++++
 td/telegram/MessagesManager.cpp      |  7 ++-
 td/telegram/MessagesManager.h        |  4 +-
 td/telegram/SetWithPosition.h        | 46 +++++++--------
 td/telegram/Td.cpp                   | 28 +++++++++
 td/telegram/WallpaperManager.cpp     |  9 +--
 td/telegram/WallpaperManager.h       |  2 +-
 td/telegram/WebPagesManager.cpp      | 41 +++++++++----
 td/telegram/WebPagesManager.h        |  7 +++
 td/telegram/files/FileData.h         | 32 ++++++++---
 td/telegram/files/FileLocation.h     | 12 ++--
 td/telegram/files/FileManager.cpp    | 81 ++++++++++++++------------
 td/telegram/files/FileManager.h      | 17 +++++-
 20 files changed, 385 insertions(+), 116 deletions(-)
 create mode 100644 td/telegram/FileReferenceManager.hpp

diff --git a/CMakeLists.txt b/CMakeLists.txt
index f4963f1ec..60317188f 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -483,6 +483,7 @@ set(TDLIB_SOURCE
   td/telegram/DocumentsManager.h
   td/telegram/DraftMessage.h
   td/telegram/FileReferenceManager.h
+  td/telegram/FileReferenceManager.hpp
   td/telegram/files/FileBitmask.h
   td/telegram/files/FileData.h
   td/telegram/files/FileDb.h
@@ -501,10 +502,10 @@ set(TDLIB_SOURCE
   td/telegram/files/FileLoadManager.h
   td/telegram/files/FileLocation.h
   td/telegram/files/FileManager.h
-  td/telegram/files/FileType.h
   td/telegram/files/FileSourceId.h
   td/telegram/files/FileStats.h
   td/telegram/files/FileStatsWorker.h
+  td/telegram/files/FileType.h
   td/telegram/files/FileUploader.h
   td/telegram/files/PartsManager.h
   td/telegram/files/ResourceManager.h
diff --git a/td/telegram/AnimationsManager.cpp b/td/telegram/AnimationsManager.cpp
index a8196536e..0075615b3 100644
--- a/td/telegram/AnimationsManager.cpp
+++ b/td/telegram/AnimationsManager.cpp
@@ -753,10 +753,7 @@ void AnimationsManager::send_update_saved_animations(bool from_database) {
     }
     std::sort(new_saved_animation_file_ids.begin(), new_saved_animation_file_ids.end());
     if (new_saved_animation_file_ids != saved_animation_file_ids_) {
-      if (!saved_animations_file_source_id_.is_valid()) {
-        saved_animations_file_source_id_ = td_->file_reference_manager_->create_saved_animations_file_source();
-      }
-      td_->file_manager_->change_files_source(saved_animations_file_source_id_, saved_animation_file_ids_,
+      td_->file_manager_->change_files_source(get_saved_animations_file_source_id(), saved_animation_file_ids_,
                                               new_saved_animation_file_ids);
       saved_animation_file_ids_ = std::move(new_saved_animation_file_ids);
     }
@@ -777,6 +774,13 @@ void AnimationsManager::save_saved_animations_to_database() {
   }
 }
 
+FileSourceId AnimationsManager::get_saved_animations_file_source_id() {
+  if (!saved_animations_file_source_id_.is_valid()) {
+    saved_animations_file_source_id_ = td_->file_reference_manager_->create_saved_animations_file_source();
+  }
+  return saved_animations_file_source_id_;
+}
+
 string AnimationsManager::get_animation_search_text(FileId file_id) const {
   auto animation = get_animation(file_id);
   CHECK(animation != nullptr);
diff --git a/td/telegram/AnimationsManager.h b/td/telegram/AnimationsManager.h
index 4d68d1d6c..7f12ca6df 100644
--- a/td/telegram/AnimationsManager.h
+++ b/td/telegram/AnimationsManager.h
@@ -68,6 +68,8 @@ class AnimationsManager : public Actor {
 
   vector<FileId> get_saved_animations(Promise<Unit> &&promise);
 
+  FileSourceId get_saved_animations_file_source_id();
+
   void send_save_gif_query(FileId animation_id, bool unsave, Promise<Unit> &&promise);
 
   void add_saved_animation(const tl_object_ptr<td_api::InputFile> &input_file, Promise<Unit> &&promise);
diff --git a/td/telegram/ContactsManager.cpp b/td/telegram/ContactsManager.cpp
index d1ea24a13..ffe972659 100644
--- a/td/telegram/ContactsManager.cpp
+++ b/td/telegram/ContactsManager.cpp
@@ -49,6 +49,7 @@
 #include "td/utils/tl_helpers.h"
 
 #include <algorithm>
+#include <functional>
 #include <limits>
 #include <tuple>
 #include <utility>
@@ -6915,7 +6916,15 @@ void ContactsManager::do_update_user_photo(User *u, UserId user_id,
 
 void ContactsManager::add_user_photo_id(User *u, UserId user_id, int64 photo_id, const vector<FileId> &photo_file_ids) {
   if (photo_id > 0 && !photo_file_ids.empty() && u->photo_ids.insert(photo_id).second) {
-    auto file_source_id = td_->file_reference_manager_->create_user_photo_file_source(user_id, photo_id);
+    FileSourceId file_source_id;
+    auto it = user_profile_photo_file_source_ids_.find(std::make_pair(user_id, photo_id));
+    if (it != user_profile_photo_file_source_ids_.end()) {
+      VLOG(file_references) << "Move " << it->second << " inside of " << user_id;
+      file_source_id = it->second;
+      user_profile_photo_file_source_ids_.erase(it);
+    } else {
+      file_source_id = td_->file_reference_manager_->create_user_photo_file_source(user_id, photo_id);
+    }
     for (auto &file_id : photo_file_ids) {
       td_->file_manager_->add_file_source(file_id, file_source_id);
     }
@@ -8548,6 +8557,7 @@ std::pair<int32, vector<const Photo *>> ContactsManager::get_user_profile_photos
 }
 
 void ContactsManager::reload_user_profile_photo(UserId user_id, int64 photo_id, Promise<Unit> &&promise) {
+  get_user_force(user_id);
   auto input_user = get_input_user(user_id);
   if (input_user == nullptr) {
     return promise.set_error(Status::Error(6, "User info not found"));
@@ -8558,6 +8568,22 @@ void ContactsManager::reload_user_profile_photo(UserId user_id, int64 photo_id,
   td_->create_handler<GetUserPhotosQuery>(std::move(promise))->send(user_id, std::move(input_user), -1, 1, photo_id);
 }
 
+FileSourceId ContactsManager::get_user_profile_photo_file_source_id(UserId user_id, int64 photo_id) {
+  auto u = get_user(user_id);
+  if (u != nullptr && u->photo_ids.count(photo_id) != 0) {
+    VLOG(file_references) << "Don't need to create file source for photo " << photo_id << " of " << user_id;
+    // photo was already added, source id was registered and shouldn't be needed
+    return FileSourceId();
+  }
+
+  auto &source_id = user_profile_photo_file_source_ids_[std::make_pair(user_id, photo_id)];
+  if (!source_id.is_valid()) {
+    source_id = td_->file_reference_manager_->create_user_photo_file_source(user_id, photo_id);
+  }
+  VLOG(file_references) << "Return " << source_id << " for photo " << photo_id << " of " << user_id;
+  return source_id;
+}
+
 bool ContactsManager::have_chat(ChatId chat_id) const {
   return chats_.count(chat_id) > 0;
 }
@@ -8581,8 +8607,20 @@ ContactsManager::Chat *ContactsManager::get_chat(ChatId chat_id) {
 }
 
 ContactsManager::Chat *ContactsManager::add_chat(ChatId chat_id) {
+  auto c = get_chat(chat_id);
+  if (c != nullptr) {
+    return c;
+  }
+
   CHECK(chat_id.is_valid());
-  return &chats_[chat_id];
+  c = &chats_[chat_id];
+  auto it = chat_photo_file_source_ids_.find(chat_id);
+  if (it != chat_photo_file_source_ids_.end()) {
+    VLOG(file_references) << "Move " << it->second << " inside of " << chat_id;
+    c->photo_source_id = it->second;
+    chat_photo_file_source_ids_.erase(it);
+  }
+  return c;
 }
 
 bool ContactsManager::get_chat(ChatId chat_id, int left_tries, Promise<Unit> &&promise) {
@@ -8771,6 +8809,15 @@ bool ContactsManager::is_appointed_chat_administrator(ChatId chat_id) const {
   }
 }
 
+FileSourceId ContactsManager::get_chat_photo_file_source_id(ChatId chat_id) {
+  auto c = get_chat(chat_id);
+  auto &source_id = c == nullptr ? chat_photo_file_source_ids_[chat_id] : c->photo_source_id;
+  if (!source_id.is_valid()) {
+    source_id = td_->file_reference_manager_->create_chat_photo_file_source(chat_id);
+  }
+  return source_id;
+}
+
 ChannelType ContactsManager::get_channel_type(ChannelId channel_id) const {
   auto c = get_channel(channel_id);
   if (c == nullptr) {
@@ -8819,6 +8866,15 @@ bool ContactsManager::get_channel_sign_messages(const Channel *c) {
   return c->sign_messages;
 }
 
+FileSourceId ContactsManager::get_channel_photo_file_source_id(ChannelId channel_id) {
+  auto c = get_channel(channel_id);
+  auto &source_id = c == nullptr ? channel_photo_file_source_ids_[channel_id] : c->photo_source_id;
+  if (!source_id.is_valid()) {
+    source_id = td_->file_reference_manager_->create_channel_photo_file_source(channel_id);
+  }
+  return source_id;
+}
+
 bool ContactsManager::have_channel(ChannelId channel_id) const {
   return channels_.count(channel_id) > 0;
 }
@@ -8846,11 +8902,20 @@ ContactsManager::Channel *ContactsManager::get_channel(ChannelId channel_id) {
 }
 
 ContactsManager::Channel *ContactsManager::add_channel(ChannelId channel_id, const char *source) {
-  CHECK(channel_id.is_valid());
-  Channel *c = &channels_[channel_id];
-  if (c->debug_source == nullptr) {
-    c->debug_source = source;
+  auto c = get_channel(channel_id);
+  if (c != nullptr) {
+    return c;
   }
+
+  CHECK(channel_id.is_valid());
+  c = &channels_[channel_id];
+  auto it = channel_photo_file_source_ids_.find(channel_id);
+  if (it != channel_photo_file_source_ids_.end()) {
+    VLOG(file_references) << "Move " << it->second << " inside of " << channel_id;
+    c->photo_source_id = it->second;
+    channel_photo_file_source_ids_.erase(it);
+  }
+  c->debug_source = source;
   return c;
 }
 
@@ -8885,6 +8950,7 @@ void ContactsManager::reload_channel(ChannelId channel_id, Promise<Unit> &&promi
     return promise.set_error(Status::Error(6, "Invalid supergroup id"));
   }
 
+  have_channel_force(channel_id);
   auto input_channel = get_input_channel(channel_id);
   if (input_channel == nullptr) {
     return promise.set_error(Status::Error(6, "Supergroup info not found"));
diff --git a/td/telegram/ContactsManager.h b/td/telegram/ContactsManager.h
index f5e854fb9..92dcf3ee2 100644
--- a/td/telegram/ContactsManager.h
+++ b/td/telegram/ContactsManager.h
@@ -364,6 +364,7 @@ class ContactsManager : public Actor {
   std::pair<int32, vector<const Photo *>> get_user_profile_photos(UserId user_id, int32 offset, int32 limit,
                                                                   Promise<Unit> &&promise);
   void reload_user_profile_photo(UserId user_id, int64 photo_id, Promise<Unit> &&promise);
+  FileSourceId get_user_profile_photo_file_source_id(UserId user_id, int64 photo_id);
 
   bool have_chat(ChatId chat_id) const;
   bool have_chat_force(ChatId chat_id);
@@ -374,6 +375,7 @@ class ContactsManager : public Actor {
   bool get_chat_is_active(ChatId chat_id) const;
   DialogParticipantStatus get_chat_status(ChatId chat_id) const;
   bool is_appointed_chat_administrator(ChatId chat_id) const;
+  FileSourceId get_chat_photo_file_source_id(ChatId chat_id);
 
   bool have_channel(ChannelId channel_id) const;
   bool have_min_channel(ChannelId channel_id) const;
@@ -391,6 +393,7 @@ class ContactsManager : public Actor {
   int32 get_channel_date(ChannelId channel_id) const;
   DialogParticipantStatus get_channel_status(ChannelId channel_id) const;
   bool get_channel_sign_messages(ChannelId channel_id) const;
+  FileSourceId get_channel_photo_file_source_id(ChannelId channel_id);
 
   std::pair<int32, vector<UserId>> search_among_users(const vector<UserId> &user_ids, const string &query, int32 limit);
 
@@ -1061,15 +1064,23 @@ class ContactsManager : public Actor {
   std::unordered_map<UserId, UserFull, UserIdHash> users_full_;
   mutable std::unordered_set<UserId, UserIdHash> unknown_users_;
   std::unordered_map<UserId, tl_object_ptr<telegram_api::UserProfilePhoto>, UserIdHash> pending_user_photos_;
+  struct UserIdPhotoIdHash {
+    std::size_t operator()(const std::pair<UserId, int64> &pair) const {
+      return UserIdHash()(pair.first) * 2023654985u + std::hash<int64>()(pair.second);
+    }
+  };
+  std::unordered_map<std::pair<UserId, int64>, FileSourceId, UserIdPhotoIdHash> user_profile_photo_file_source_ids_;
 
   std::unordered_map<ChatId, Chat, ChatIdHash> chats_;
   std::unordered_map<ChatId, ChatFull, ChatIdHash> chats_full_;
   mutable std::unordered_set<ChatId, ChatIdHash> unknown_chats_;
+  std::unordered_map<ChatId, FileSourceId, ChatIdHash> chat_photo_file_source_ids_;
 
   std::unordered_set<ChannelId, ChannelIdHash> min_channels_;
   std::unordered_map<ChannelId, Channel, ChannelIdHash> channels_;
   std::unordered_map<ChannelId, ChannelFull, ChannelIdHash> channels_full_;
   mutable std::unordered_set<ChannelId, ChannelIdHash> unknown_channels_;
+  std::unordered_map<ChannelId, FileSourceId, ChannelIdHash> channel_photo_file_source_ids_;
 
   std::unordered_map<SecretChatId, SecretChat, SecretChatIdHash> secret_chats_;
   mutable std::unordered_set<SecretChatId, SecretChatIdHash> unknown_secret_chats_;
diff --git a/td/telegram/FileReferenceManager.cpp b/td/telegram/FileReferenceManager.cpp
index 35136fdf5..62e014419 100644
--- a/td/telegram/FileReferenceManager.cpp
+++ b/td/telegram/FileReferenceManager.cpp
@@ -83,14 +83,14 @@ FileSourceId FileReferenceManager::create_saved_animations_file_source() {
   return add_file_source_id(source, "saved animations");
 }
 
-void FileReferenceManager::add_file_source(NodeId node_id, FileSourceId file_source_id) {
+bool FileReferenceManager::add_file_source(NodeId node_id, FileSourceId file_source_id) {
   VLOG(file_references) << "Add " << file_source_id << " for file " << node_id;
-  nodes_[node_id].file_source_ids.add(file_source_id);
+  return nodes_[node_id].file_source_ids.add(file_source_id);
 }
 
-void FileReferenceManager::remove_file_source(NodeId node_id, FileSourceId file_source_id) {
+bool FileReferenceManager::remove_file_source(NodeId node_id, FileSourceId file_source_id) {
   VLOG(file_references) << "Remove " << file_source_id << " from file " << node_id;
-  nodes_[node_id].file_source_ids.remove(file_source_id);
+  return nodes_[node_id].file_source_ids.remove(file_source_id);
 }
 
 std::vector<FileSourceId> FileReferenceManager::get_some_file_sources(NodeId node_id) {
@@ -276,7 +276,7 @@ void FileReferenceManager::repair_file_reference(NodeId node_id, Promise<> promi
     node.query = make_unique<Query>();
     node.query->generation = ++query_generation_;
     node.file_source_ids.reset_position();
-    VLOG(file_references) << "Create new file reference repair query with " << query_generation_;
+    VLOG(file_references) << "Create new file reference repair query with generation " << query_generation_;
   }
   node.query->promises.push_back(std::move(promise));
   run_node(node_id);
diff --git a/td/telegram/FileReferenceManager.h b/td/telegram/FileReferenceManager.h
index 69e6088ff..c4a47d287 100644
--- a/td/telegram/FileReferenceManager.h
+++ b/td/telegram/FileReferenceManager.h
@@ -26,6 +26,8 @@
 
 namespace td {
 
+class Td;
+
 extern int VERBOSITY_NAME(file_references);
 
 class FileReferenceManager : public Actor {
@@ -43,14 +45,20 @@ class FileReferenceManager : public Actor {
   using NodeId = FileId;
   void repair_file_reference(NodeId node_id, Promise<> promise);
 
-  void add_file_source(NodeId node_id, FileSourceId file_source_id);
+  bool add_file_source(NodeId node_id, FileSourceId file_source_id);
 
   std::vector<FileSourceId> get_some_file_sources(NodeId node_id);
 
-  void remove_file_source(NodeId node_id, FileSourceId file_source_id);
+  bool remove_file_source(NodeId node_id, FileSourceId file_source_id);
 
   void merge(NodeId to_node_id, NodeId from_node_id);
 
+  template <class StorerT>
+  void store_file_source(FileSourceId file_source_id, StorerT &storer) const;
+
+  template <class ParserT>
+  FileSourceId parse_file_source(Td *td, ParserT &parser);
+
  private:
   struct Destination {
     bool empty() const {
@@ -94,6 +102,7 @@ class FileReferenceManager : public Actor {
     // empty
   };
 
+  // append only
   using FileSource = Variant<FileSourceMessage, FileSourceUserPhoto, FileSourceChatPhoto, FileSourceChannelPhoto,
                              FileSourceWallpapers, FileSourceWebPage, FileSourceSavedAnimations>;
   vector<FileSource> file_sources_;
diff --git a/td/telegram/FileReferenceManager.hpp b/td/telegram/FileReferenceManager.hpp
new file mode 100644
index 000000000..a8b0f2599
--- /dev/null
+++ b/td/telegram/FileReferenceManager.hpp
@@ -0,0 +1,86 @@
+//
+// Copyright Aliaksei Levin (levlam@telegram.org), Arseny Smirnov (arseny30@gmail.com) 2014-2019
+//
+// Distributed under the Boost Software License, Version 1.0. (See accompanying
+// file LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)
+//
+#pragma once
+
+#include "td/telegram/AnimationsManager.h"
+#include "td/telegram/ChannelId.h"
+#include "td/telegram/ChatId.h"
+#include "td/telegram/ContactsManager.h"
+#include "td/telegram/FileReferenceManager.h"
+#include "td/telegram/files/FileSourceId.h"
+#include "td/telegram/MessageId.h"
+#include "td/telegram/MessagesManager.h"
+#include "td/telegram/Td.h"
+#include "td/telegram/UserId.h"
+#include "td/telegram/WallpaperManager.h"
+#include "td/telegram/WebPagesManager.h"
+
+#include "td/utils/overloaded.h"
+#include "td/utils/tl_helpers.h"
+
+namespace td {
+
+template <class StorerT>
+void FileReferenceManager::store_file_source(FileSourceId file_source_id, StorerT &storer) const {
+  auto index = static_cast<size_t>(file_source_id.get()) - 1;
+  CHECK(index < file_sources_.size());
+  auto &source = file_sources_[index];
+  td::store(source.get_offset(), storer);
+  source.visit(overloaded([&](const FileSourceMessage &source) { td::store(source.full_message_id, storer); },
+                          [&](const FileSourceUserPhoto &source) {
+                            td::store(source.user_id, storer);
+                            td::store(source.photo_id, storer);
+                          },
+                          [&](const FileSourceChatPhoto &source) { td::store(source.chat_id, storer); },
+                          [&](const FileSourceChannelPhoto &source) { td::store(source.channel_id, storer); },
+                          [&](const FileSourceWallpapers &source) {},
+                          [&](const FileSourceWebPage &source) { td::store(source.url, storer); },
+                          [&](const FileSourceSavedAnimations &source) {}));
+}
+
+template <class ParserT>
+FileSourceId FileReferenceManager::parse_file_source(Td *td, ParserT &parser) {
+  auto type = parser.fetch_int();
+  switch (type) {
+    case 0: {
+      FullMessageId full_message_id;
+      td::parse(full_message_id, parser);
+      return td->messages_manager_->get_message_file_source_id(full_message_id);
+    }
+    case 1: {
+      UserId user_id;
+      int64 photo_id;
+      td::parse(user_id, parser);
+      td::parse(photo_id, parser);
+      return td->contacts_manager_->get_user_profile_photo_file_source_id(user_id, photo_id);
+    }
+    case 2: {
+      ChatId chat_id;
+      td::parse(chat_id, parser);
+      return td->contacts_manager_->get_chat_photo_file_source_id(chat_id);
+    }
+    case 3: {
+      ChannelId channel_id;
+      td::parse(channel_id, parser);
+      return td->contacts_manager_->get_channel_photo_file_source_id(channel_id);
+    }
+    case 4:
+      return td->wallpaper_manager_->get_wallpapers_file_source_id();
+    case 5: {
+      string url;
+      td::parse(url, parser);
+      return td->web_pages_manager_->get_url_file_source_id(url);
+    }
+    case 6:
+      return td->animations_manager_->get_saved_animations_file_source_id();
+    default:
+      parser.set_error("Invalid type in FileSource");
+      return FileSourceId();
+  }
+}
+
+}  // namespace td
diff --git a/td/telegram/MessagesManager.cpp b/td/telegram/MessagesManager.cpp
index 6aace4ff2..1a633b8d3 100644
--- a/td/telegram/MessagesManager.cpp
+++ b/td/telegram/MessagesManager.cpp
@@ -11660,6 +11660,7 @@ void MessagesManager::get_messages_from_server(vector<FullMessageId> &&message_i
   }
 
   for (auto &it : channel_message_ids) {
+    td_->contacts_manager_->have_channel_force(it.first);
     auto input_channel = td_->contacts_manager_->get_input_channel(it.first);
     if (input_channel == nullptr) {
       LOG(ERROR) << "Can't find info about " << it.first << " to get a message from it";
@@ -13733,8 +13734,10 @@ void MessagesManager::save_active_live_locations() {
 }
 
 FileSourceId MessagesManager::get_message_file_source_id(FullMessageId full_message_id) {
-  if (full_message_id.get_dialog_id().get_type() == DialogType::SecretChat ||
-      !full_message_id.get_message_id().is_server()) {
+  auto dialog_id = full_message_id.get_dialog_id();
+  auto message_id = full_message_id.get_message_id();
+  if (!dialog_id.is_valid() || !message_id.is_valid() || dialog_id.get_type() == DialogType::SecretChat ||
+      !message_id.is_server()) {
     return FileSourceId();
   }
 
diff --git a/td/telegram/MessagesManager.h b/td/telegram/MessagesManager.h
index 8d4ef82dd..d5cf65174 100644
--- a/td/telegram/MessagesManager.h
+++ b/td/telegram/MessagesManager.h
@@ -666,6 +666,8 @@ class MessagesManager : public Actor {
                                   vector<tl_object_ptr<telegram_api::User>> users,
                                   vector<tl_object_ptr<telegram_api::Chat>> chats);
 
+  FileSourceId get_message_file_source_id(FullMessageId full_message_id);
+
   struct MessageNotificationGroup {
     DialogId dialog_id;
     NotificationGroupType type = NotificationGroupType::Calls;
@@ -1847,8 +1849,6 @@ class MessagesManager : public Actor {
 
   void save_active_live_locations();
 
-  FileSourceId get_message_file_source_id(FullMessageId full_message_id);
-
   void add_message_file_sources(DialogId dialog_id, const Message *m);
 
   void remove_message_file_sources(DialogId dialog_id, const Message *m);
diff --git a/td/telegram/SetWithPosition.h b/td/telegram/SetWithPosition.h
index e843d243a..36543d015 100644
--- a/td/telegram/SetWithPosition.h
+++ b/td/telegram/SetWithPosition.h
@@ -10,6 +10,7 @@
 #include "td/utils/logging.h"
 #include "td/utils/misc.h"
 
+#include <algorithm>
 #include <set>
 #include <utility>
 
@@ -20,34 +21,32 @@ class FastSetWithPosition {
  public:
   std::vector<T> get_some_elements() const {
     std::vector<T> res;
-    using std::prev;
-    using std::next;
-    res.reserve(5);
+    res.reserve(4);
     if (!checked_.empty()) {
-      res.push_back(*begin(checked_));
-      res.push_back(*prev(end(checked_)));
+      res.push_back(*checked_.begin());
+      res.push_back(*checked_.rbegin());
     }
     if (!not_checked_.empty()) {
-      res.push_back(*begin(not_checked_));
-      res.push_back(*prev(end(not_checked_)));
+      res.push_back(*not_checked_.begin());
+      res.push_back(*not_checked_.rbegin());
     }
     std::sort(res.begin(), res.end());
     res.erase(std::unique(res.begin(), res.end()), res.end());
     if (res.size() > 2) {
-      res.erase(next(res.begin()), prev(res.end()));
+      res.erase(res.begin() + 1, res.end() - 1);
     }
     return res;
   }
-  void add(T x) {
+
+  bool add(T x) {
     if (checked_.count(x) != 0) {
-      return;
+      return false;
     }
-    not_checked_.insert(x);
+    return not_checked_.insert(x).second;
   }
 
-  void remove(T x) {
-    checked_.erase(x);
-    not_checked_.erase(x);
+  bool remove(T x) {
+    return checked_.erase(x) != 0 || not_checked_.erase(x) != 0;
   }
 
   bool has_next() const {
@@ -118,33 +117,34 @@ class SetWithPosition {
     }
     return {};
   }
-  void add(T x) {
+
+  bool add(T x) {
     if (fast_) {
-      fast_->add(x);
-      return;
+      return fast_->add(x);
     }
     if (!has_value_) {
       value_ = x;
       has_value_ = true;
       is_checked_ = false;
-      return;
+      return true;
     }
     if (value_ == x) {
-      return;
+      return false;
     }
     make_fast();
-    fast_->add(x);
+    return fast_->add(x);
   }
 
-  void remove(T x) {
+  bool remove(T x) {
     if (fast_) {
-      fast_->remove(x);
-      return;
+      return fast_->remove(x);
     }
     if (has_value_ && value_ == x) {
       has_value_ = false;
       is_checked_ = false;
+      return true;
     }
+    return false;
   }
 
   bool has_next() const {
diff --git a/td/telegram/Td.cpp b/td/telegram/Td.cpp
index 9ab444d15..2f703ac05 100644
--- a/td/telegram/Td.cpp
+++ b/td/telegram/Td.cpp
@@ -4073,13 +4073,41 @@ Status Td::init(DbKey key) {
    public:
     explicit FileManagerContext(Td *td) : td_(td) {
     }
+
     void on_new_file(int64 size, int32 cnt) final {
       send_closure(G()->storage_manager(), &StorageManager::on_new_file, size, cnt);
     }
+
     void on_file_updated(FileId file_id) final {
       send_closure(G()->td(), &Td::send_update,
                    make_tl_object<td_api::updateFile>(td_->file_manager_->get_file_object(file_id)));
     }
+
+    bool add_file_source(FileId file_id, FileSourceId file_source_id) final {
+      return td_->file_reference_manager_->add_file_source(file_id, file_source_id);
+    }
+
+    FileSourceId get_wallpapers_file_source_id() final {
+      return td_->wallpaper_manager_->get_wallpapers_file_source_id();
+    }
+
+    bool remove_file_source(FileId file_id, FileSourceId file_source_id) final {
+      return td_->file_reference_manager_->remove_file_source(file_id, file_source_id);
+    }
+
+    void on_merge_files(FileId to_file_id, FileId from_file_id) final {
+      td_->file_reference_manager_->merge(to_file_id, from_file_id);
+    }
+
+    vector<FileSourceId> get_some_file_sources(FileId file_id) final {
+      return td_->file_reference_manager_->get_some_file_sources(file_id);
+    }
+
+    void repair_file_reference(FileId file_id, Promise<Unit> promise) final {
+      send_closure(G()->file_reference_manager(), &FileReferenceManager::repair_file_reference, file_id,
+                   std::move(promise));
+    }
+
     ActorShared<> create_reference() final {
       return td_->create_reference();
     }
diff --git a/td/telegram/WallpaperManager.cpp b/td/telegram/WallpaperManager.cpp
index ce8ba3eff..3ea3a4578 100644
--- a/td/telegram/WallpaperManager.cpp
+++ b/td/telegram/WallpaperManager.cpp
@@ -117,10 +117,7 @@ void WallpaperManager::on_get_wallpapers(Result<vector<telegram_api::object_ptr<
   for (auto &wallpaper : wallpapers_) {
     append(new_file_ids, transform(wallpaper.sizes, [](auto &size) { return size.file_id; }));
   };
-  if (!wallpaper_source_id_.is_valid()) {
-    wallpaper_source_id_ = td_->file_reference_manager_->create_wallpapers_file_source();
-  }
-  td_->file_manager_->change_files_source(wallpaper_source_id_, wallpaper_file_ids_, new_file_ids);
+  td_->file_manager_->change_files_source(get_wallpapers_file_source_id(), wallpaper_file_ids_, new_file_ids);
   wallpaper_file_ids_ = std::move(new_file_ids);
 
   for (auto &promise : promises) {
@@ -128,11 +125,11 @@ void WallpaperManager::on_get_wallpapers(Result<vector<telegram_api::object_ptr<
   }
 }
 
-void WallpaperManager::add_wallpapers_file_source(FileId file_id) {
+FileSourceId WallpaperManager::get_wallpapers_file_source_id() {
   if (!wallpaper_source_id_.is_valid()) {
     wallpaper_source_id_ = td_->file_reference_manager_->create_wallpapers_file_source();
   }
-  td_->file_manager_->add_file_source(file_id, wallpaper_source_id_);
+  return wallpaper_source_id_;
 }
 
 td_api::object_ptr<td_api::wallpapers> WallpaperManager::get_wallpapers_object() const {
diff --git a/td/telegram/WallpaperManager.h b/td/telegram/WallpaperManager.h
index 8dfbc0e16..4885ee60d 100644
--- a/td/telegram/WallpaperManager.h
+++ b/td/telegram/WallpaperManager.h
@@ -31,7 +31,7 @@ class WallpaperManager : public Actor {
 
   td_api::object_ptr<td_api::wallpapers> get_wallpapers_object() const;
 
-  void add_wallpapers_file_source(FileId file_id);
+  FileSourceId get_wallpapers_file_source_id();
 
  private:
   void tear_down() override;
diff --git a/td/telegram/WebPagesManager.cpp b/td/telegram/WebPagesManager.cpp
index 483d0fb9c..2cc72d6dd 100644
--- a/td/telegram/WebPagesManager.cpp
+++ b/td/telegram/WebPagesManager.cpp
@@ -1625,6 +1625,13 @@ void WebPagesManager::update_web_page(unique_ptr<WebPage> web_page, WebPageId we
   if (page != nullptr) {
     old_instant_view = std::move(page->instant_view);
     web_page->logevent_id = page->logevent_id;
+  } else {
+    auto it = url_to_file_source_id_.find(web_page->url);
+    if (it != url_to_file_source_id_.end()) {
+      VLOG(file_references) << "Move " << it->second << " inside of " << web_page_id;
+      web_page->file_source_id = it->second;
+      url_to_file_source_id_.erase(it);
+    }
   }
   page = std::move(web_page);
 
@@ -1632,11 +1639,7 @@ void WebPagesManager::update_web_page(unique_ptr<WebPage> web_page, WebPageId we
 
   auto new_file_ids = get_web_page_file_ids(page.get());
   if (old_file_ids != new_file_ids) {
-    if (!page->file_source_id.is_valid()) {
-      LOG(ERROR) << page->url;
-      page->file_source_id = td_->file_reference_manager_->create_web_page_file_source(page->url);
-    }
-    td_->file_manager_->change_files_source(page->file_source_id, old_file_ids, new_file_ids);
+    td_->file_manager_->change_files_source(get_web_page_file_source_id(page.get()), old_file_ids, new_file_ids);
   }
 
   on_get_web_page_by_url(page->url, web_page_id, from_database);
@@ -1976,11 +1979,7 @@ void WebPagesManager::on_load_web_page_instant_view_from_database(WebPageId web_
 
   auto new_file_ids = get_web_page_file_ids(web_page);
   if (old_file_ids != new_file_ids) {
-    if (!web_page->file_source_id.is_valid()) {
-      LOG(ERROR) << web_page->url;
-      web_page->file_source_id = td_->file_reference_manager_->create_web_page_file_source(web_page->url);
-    }
-    td_->file_manager_->change_files_source(web_page->file_source_id, old_file_ids, new_file_ids);
+    td_->file_manager_->change_files_source(get_web_page_file_source_id(web_page), old_file_ids, new_file_ids);
   }
 
   update_web_page_instant_view_load_requests(web_page_id, false, Unit());
@@ -2888,6 +2887,28 @@ const WebPagesManager::WebPage *WebPagesManager::get_web_page_force(WebPageId we
   return get_web_page(web_page_id);
 }
 
+FileSourceId WebPagesManager::get_web_page_file_source_id(WebPage *web_page) {
+  if (!web_page->file_source_id.is_valid()) {
+    web_page->file_source_id = td_->file_reference_manager_->create_web_page_file_source(web_page->url);
+  }
+  return web_page->file_source_id;
+}
+
+FileSourceId WebPagesManager::get_url_file_source_id(const string &url) {
+  auto web_page_id = get_web_page_by_url(url);
+  if (web_page_id.is_valid()) {
+    auto web_page = get_web_page(web_page_id);
+    if (web_page != nullptr) {
+      if (!web_page->file_source_id.is_valid()) {
+        web_pages_[web_page_id]->file_source_id =
+            td_->file_reference_manager_->create_web_page_file_source(web_page->url);
+      }
+      return web_page->file_source_id;
+    }
+  }
+  return url_to_file_source_id_[url] = td_->file_reference_manager_->create_web_page_file_source(url);
+}
+
 string WebPagesManager::get_web_page_search_text(WebPageId web_page_id) const {
   auto web_page = get_web_page(web_page_id);
   if (web_page == nullptr) {
diff --git a/td/telegram/WebPagesManager.h b/td/telegram/WebPagesManager.h
index 555ed9360..3b4d01310 100644
--- a/td/telegram/WebPagesManager.h
+++ b/td/telegram/WebPagesManager.h
@@ -11,6 +11,7 @@
 
 #include "td/telegram/DialogId.h"
 #include "td/telegram/files/FileId.h"
+#include "td/telegram/files/FileSourceId.h"
 #include "td/telegram/MessageId.h"
 #include "td/telegram/Photo.h"
 #include "td/telegram/WebPageId.h"
@@ -78,6 +79,8 @@ class WebPagesManager : public Actor {
 
   void on_binlog_web_page_event(BinlogEvent &&event);
 
+  FileSourceId get_url_file_source_id(const string &url);
+
   string get_web_page_search_text(WebPageId web_page_id) const;
 
  private:
@@ -216,6 +219,8 @@ class WebPagesManager : public Actor {
 
   void tear_down() override;
 
+  FileSourceId get_web_page_file_source_id(WebPage *web_page);
+
   static vector<FileId> get_web_page_file_ids(const WebPage *web_page);
 
   Td *td_;
@@ -240,6 +245,8 @@ class WebPagesManager : public Actor {
 
   std::unordered_map<string, WebPageId> url_to_web_page_id_;
 
+  std::unordered_map<string, FileSourceId> url_to_file_source_id_;
+
   MultiTimeout pending_web_pages_timeout_{"PendingWebPagesTimeout"};
 };
 
diff --git a/td/telegram/files/FileData.h b/td/telegram/files/FileData.h
index fb9c22292..96fc4bd9d 100644
--- a/td/telegram/files/FileData.h
+++ b/td/telegram/files/FileData.h
@@ -7,9 +7,13 @@
 #pragma once
 
 #include "td/telegram/DialogId.h"
+#include "td/telegram/Global.h"
+#include "td/telegram/FileReferenceManager.h"
+#include "td/telegram/FileReferenceManager.hpp"
 #include "td/telegram/files/FileEncryptionKey.h"
 #include "td/telegram/files/FileLocation.h"
 #include "td/telegram/files/FileSourceId.h"
+#include "td/telegram/Td.h"
 
 #include "td/utils/common.h"
 #include "td/utils/format.h"
@@ -30,7 +34,7 @@ class FileData {
   string remote_name_;
   string url_;
   FileEncryptionKey encryption_key_;
-  std::vector<FileSourceId> sources_;
+  vector<FileSourceId> file_source_ids_;
 
   template <class StorerT>
   void store(StorerT &storer) const {
@@ -38,13 +42,12 @@ class FileData {
     bool has_owner_dialog_id = owner_dialog_id_.is_valid();
     bool has_expected_size = size_ == 0 && expected_size_ != 0;
     bool encryption_key_is_secure = encryption_key_.is_secure();
-    bool has_sources = !sources_.empty();
+    bool has_sources = !file_source_ids_.empty();
     BEGIN_STORE_FLAGS();
     STORE_FLAG(has_owner_dialog_id);
     STORE_FLAG(has_expected_size);
     STORE_FLAG(encryption_key_is_secure);
-    //TODO: uncomment
-    //STORE_FLAG(has_sources);
+    STORE_FLAG(has_sources);
     END_STORE_FLAGS();
 
     if (has_owner_dialog_id) {
@@ -64,8 +67,11 @@ class FileData {
     store(url_, storer);
     store(encryption_key_, storer);
     if (has_sources) {
-      // TODO: uncomment
-      // store(sources_, storer);
+      auto td = G()->td().get_actor_unsafe();
+      store(narrow_cast<int32>(file_source_ids_.size()), storer);
+      for (auto file_source_id : file_source_ids_) {
+        td->file_reference_manager_->store_file_source(file_source_id, storer);
+      }
     }
   }
   template <class ParserT>
@@ -105,8 +111,16 @@ class FileData {
     encryption_key_.parse(encryption_key_is_secure ? FileEncryptionKey::Type::Secure : FileEncryptionKey::Type::Secret,
                           parser);
     if (has_sources) {
-      //TODO: uncomment
-      // parse(sources_, parser);
+      auto td = G()->td().get_actor_unsafe();
+      int32 size;
+      parse(size, parser);
+      if (0 < size && size < 5) {
+        for (int i = 0; i < size; i++) {
+          file_source_ids_.push_back(td->file_reference_manager_->parse_file_source(td, parser));
+        }
+      } else {
+        parser.set_error("Wrong number of file source ids");
+      }
     }
   }
 };
@@ -127,7 +141,7 @@ inline StringBuilder &operator<<(StringBuilder &sb, const FileData &file_data) {
   if (file_data.remote_.type() == RemoteFileLocation::Type::Full) {
     sb << " remote " << file_data.remote_.full();
   }
-  sb << format::as_array(file_data.sources_);
+  sb << format::as_array(file_data.file_source_ids_);
   return sb << "]";
 }
 
diff --git a/td/telegram/files/FileLocation.h b/td/telegram/files/FileLocation.h
index 6c2561b41..bad7e69da 100644
--- a/td/telegram/files/FileLocation.h
+++ b/td/telegram/files/FileLocation.h
@@ -25,16 +25,17 @@
 #include <tuple>
 
 namespace td {
+
 class FileReferenceView {
  public:
   static Slice invalid_file_reference() {
     return "#";
   }
-  static std::string create_one(td::Slice first) {
-    char second_length = char(255);
-    return PSTRING() << second_length << first;
+  static std::string create_one(Slice first) {
+    unsigned char second_length = 255;
+    return PSTRING() << static_cast<char>(second_length) << first;
   }
-  static std::string create_two(td::Slice first, td::Slice second = {}) {
+  static std::string create_two(Slice first, Slice second = {}) {
     if (second.size() >= 255) {
       LOG(ERROR) << "File reference is too big " << base64_encode(second);
       second = invalid_file_reference();
@@ -42,7 +43,7 @@ class FileReferenceView {
     char second_length = narrow_cast<unsigned char>(second.size());
     return PSTRING() << second_length << first << second;
   }
-  std::string create(td::Slice first, td::Slice second) const {
+  std::string create(Slice first, Slice second) const {
     if (size() == 1) {
       return create_one(first);
     }
@@ -55,7 +56,6 @@ class FileReferenceView {
     }
 
     unsigned char second_size = data.ubegin()[0];
-
     if (second_size == 255) {
       first_ = data.substr(1);
       second_ = data.substr(1);
diff --git a/td/telegram/files/FileManager.cpp b/td/telegram/files/FileManager.cpp
index f1a8aec08..18ccf2642 100644
--- a/td/telegram/files/FileManager.cpp
+++ b/td/telegram/files/FileManager.cpp
@@ -9,7 +9,6 @@
 #include "td/telegram/telegram_api.h"
 
 #include "td/telegram/ConfigShared.h"
-#include "td/telegram/FileReferenceManager.h"
 #include "td/telegram/files/FileData.h"
 #include "td/telegram/files/FileDb.h"
 #include "td/telegram/files/FileLoaderUtils.h"
@@ -18,7 +17,6 @@
 #include "td/telegram/misc.h"
 #include "td/telegram/SecureStorage.h"
 #include "td/telegram/TdDb.h"
-#include "td/telegram/WallpaperManager.h"
 
 #include "td/utils/base64.h"
 #include "td/utils/format.h"
@@ -299,8 +297,9 @@ bool FileNode::need_pmc_flush() const {
     has_generate_location = false;
   }
 
-  if (remote_.type() == RemoteFileLocation::Type::Full &&
-      (has_generate_location || local_.type() != LocalFileLocation::Type::Empty)) {
+  if (remote_.type() == RemoteFileLocation::Type::Full/* &&
+      (has_generate_location || local_.type() != LocalFileLocation::Type::Empty)*/) {
+    // we need to always save file sources
     return true;
   }
   if (local_.type() == LocalFileLocation::Type::Full &&
@@ -969,6 +968,12 @@ Result<FileId> FileManager::register_file(FileData &&data, FileLocationSource fi
   try_flush_node(get_file_node(file_id), "register_file");
   auto main_file_id = get_file_node(file_id)->main_file_id_;
   try_forget_file_id(file_id);
+  for (auto file_source_id : data.file_source_ids_) {
+    VLOG(file_references) << "Loaded " << data.file_source_ids_ << " for file " << main_file_id << " from " << source;
+    if (file_source_id.is_valid()) {
+      context_->add_file_source(main_file_id, file_source_id);
+    }
+  }
   return FileId(main_file_id.get(), remote_key);
 }
 
@@ -1300,13 +1305,11 @@ Result<FileId> FileManager::merge(FileId x_file_id, FileId y_file_id, bool no_sy
   node->can_search_locally_ &= other_node->can_search_locally_;
 
   if (main_file_id_i == other_node_i) {
-    send_closure(G()->file_reference_manager(), &FileReferenceManager::merge, other_node->main_file_id_,
-                 node->main_file_id_);
+    context_->on_merge_files(other_node->main_file_id_, node->main_file_id_);
     node->main_file_id_ = other_node->main_file_id_;
     node->main_file_id_priority_ = other_node->main_file_id_priority_;
   } else {
-    send_closure(G()->file_reference_manager(), &FileReferenceManager::merge, node->main_file_id_,
-                 other_node->main_file_id_);
+    context_->on_merge_files(node->main_file_id_, other_node->main_file_id_);
   }
 
   bool send_updates_flag = false;
@@ -1369,10 +1372,10 @@ void FileManager::add_file_source(FileId file_id, FileSourceId file_source_id) {
   }
 
   CHECK(file_source_id.is_valid());
-  send_closure(G()->file_reference_manager(), &FileReferenceManager::add_file_source, node->main_file_id_,
-               file_source_id);
-  node->on_pmc_changed();
-  try_flush_node_pmc(node, "add_file_source");
+  if (context_->add_file_source(node->main_file_id_, file_source_id)) {
+    node->on_pmc_changed();
+    try_flush_node_pmc(node, "add_file_source");
+  }
 }
 
 void FileManager::remove_file_source(FileId file_id, FileSourceId file_source_id) {
@@ -1382,10 +1385,10 @@ void FileManager::remove_file_source(FileId file_id, FileSourceId file_source_id
   }
 
   CHECK(file_source_id.is_valid());
-  send_closure(G()->file_reference_manager(), &FileReferenceManager::remove_file_source, node->main_file_id_,
-               file_source_id);
-  node->on_pmc_changed();
-  try_flush_node_pmc(node, "remove_file_source");
+  if (context_->remove_file_source(node->main_file_id_, file_source_id)) {
+    node->on_pmc_changed();
+    try_flush_node_pmc(node, "remove_file_source");
+  }
 }
 
 void FileManager::change_files_source(FileSourceId file_source_id, const vector<FileId> &old_file_ids,
@@ -1426,7 +1429,7 @@ void FileManager::try_flush_node_full(FileNodePtr node, bool new_remote, bool ne
   if (node->need_pmc_flush()) {
     if (file_db_) {
       load_from_pmc(node, true, true, true);
-      flush_to_pmc(node, new_remote, new_local, new_generate);
+      flush_to_pmc(node, new_remote, new_local, new_generate, "try_flush_node_full");
       if (other_pmc_id.is_valid() && node->pmc_id_ != other_pmc_id) {
         file_db_->set_file_data_ref(other_pmc_id, node->pmc_id_);
       }
@@ -1446,7 +1449,7 @@ void FileManager::try_flush_node_pmc(FileNodePtr node, const char *source) {
   if (node->need_pmc_flush()) {
     if (file_db_) {
       load_from_pmc(node, true, true, true);
-      flush_to_pmc(node, false, false, false);
+      flush_to_pmc(node, false, false, false, source);
     }
     node->on_pmc_flushed();
   }
@@ -1489,7 +1492,8 @@ void FileManager::clear_from_pmc(FileNodePtr node) {
   node->pmc_id_ = FileDbId();
 }
 
-void FileManager::flush_to_pmc(FileNodePtr node, bool new_remote, bool new_local, bool new_generate) {
+void FileManager::flush_to_pmc(FileNodePtr node, bool new_remote, bool new_local, bool new_generate,
+                               const char *source) {
   if (!file_db_) {
     return;
   }
@@ -1526,7 +1530,9 @@ void FileManager::flush_to_pmc(FileNodePtr node, bool new_remote, bool new_local
   data.encryption_key_ = node->encryption_key_;
   data.url_ = node->url_;
   data.owner_dialog_id_ = node->owner_dialog_id_;
-  data.sources_ = G()->file_reference_manager().get_actor_unsafe()->get_some_file_sources(view.file_id());
+  data.file_source_ids_ = context_->get_some_file_sources(view.file_id());
+  VLOG(file_references) << "Save " << data.file_source_ids_ << " to database for file " << view.file_id() << " from "
+                        << source;
 
   file_db_->set_file_data(node->pmc_id_, data, (create_flag || new_remote), (create_flag || new_local),
                           (create_flag || new_generate));
@@ -1822,18 +1828,18 @@ void FileManager::run_download(FileNodePtr node) {
     }
     node->download_was_update_file_reference_ = true;
 
-    send_closure(G()->file_reference_manager(), &FileReferenceManager::repair_file_reference, file_id,
-                 PromiseCreator::lambda([id, actor_id = actor_id(this), file_id](Result<Unit> res) {
-                   Status error;
-                   if (res.is_ok()) {
-                     error = Status::Error("FILE_DOWNLOAD_RESTART_WITH_FILE_REFERENCE");
-                   } else {
-                     error = res.move_as_error();
-                   }
-                   VLOG(file_references) << "run_download: Got result from FileSourceManager for file " << file_id
-                                         << ": " << error;
-                   send_closure(actor_id, &FileManager::on_error, id, std::move(error));
-                 }));
+    context_->repair_file_reference(
+        file_id, PromiseCreator::lambda([id, actor_id = actor_id(this), file_id](Result<Unit> res) {
+          Status error;
+          if (res.is_ok()) {
+            error = Status::Error("FILE_DOWNLOAD_RESTART_WITH_FILE_REFERENCE");
+          } else {
+            error = res.move_as_error();
+          }
+          VLOG(file_references) << "run_download: Got result from FileSourceManager for file " << file_id << ": "
+                                << error;
+          send_closure(actor_id, &FileManager::on_error, id, std::move(error));
+        }));
     return;
   }
 
@@ -2105,11 +2111,10 @@ void FileManager::run_upload(FileNodePtr node, std::vector<int> bad_parts) {
     node->upload_id_ = id;
     node->upload_was_update_file_reference_ = true;
 
-    send_closure(G()->file_reference_manager(), &FileReferenceManager::repair_file_reference, file_id,
-                 PromiseCreator::lambda([id, actor_id = actor_id(this)](Result<Unit> res) {
-                   send_closure(actor_id, &FileManager::on_error, id,
-                                Status::Error("FILE_UPLOAD_RESTART_WITH_FILE_REFERENCE"));
-                 }));
+    context_->repair_file_reference(file_id, PromiseCreator::lambda([id, actor_id = actor_id(this)](Result<Unit> res) {
+                                      send_closure(actor_id, &FileManager::on_error, id,
+                                                   Status::Error("FILE_UPLOAD_RESTART_WITH_FILE_REFERENCE"));
+                                    }));
     return;
   }
 
@@ -2232,7 +2237,7 @@ Result<FileId> FileManager::from_persistent_id_v2(Slice binary, FileType file_ty
   auto file_id =
       register_file(std::move(data), FileLocationSource::FromUser, "from_persistent_id_v2", false).move_as_ok();
   if (real_file_type == FileType::Wallpaper && file_id.is_valid()) {
-    send_closure(G()->wallpaper_manager(), &WallpaperManager::add_wallpapers_file_source, file_id);
+    add_file_source(file_id, context_->get_wallpapers_file_source_id());
   }
   return file_id;
 }
diff --git a/td/telegram/files/FileManager.h b/td/telegram/files/FileManager.h
index c5bd2d021..e40b952ad 100644
--- a/td/telegram/files/FileManager.h
+++ b/td/telegram/files/FileManager.h
@@ -320,8 +320,23 @@ class FileManager : public FileLoadManager::Callback {
   class Context {
    public:
     virtual void on_new_file(int64 size, int32 cnt) = 0;
+
     virtual void on_file_updated(FileId size) = 0;
+
+    virtual bool add_file_source(FileId file_id, FileSourceId file_source_id) = 0;
+
+    virtual FileSourceId get_wallpapers_file_source_id() = 0;
+
+    virtual bool remove_file_source(FileId file_id, FileSourceId file_source_id) = 0;
+
+    virtual void on_merge_files(FileId to_file_id, FileId from_file_id) = 0;
+
+    virtual vector<FileSourceId> get_some_file_sources(FileId file_id) = 0;
+
+    virtual void repair_file_reference(FileId file_id, Promise<Unit> promise) = 0;
+
     virtual ActorShared<> create_reference() = 0;
+
     Context() = default;
     Context(const Context &) = delete;
     Context &operator=(const Context &) = delete;
@@ -509,7 +524,7 @@ class FileManager : public FileLoadManager::Callback {
   void try_flush_node_info(FileNodePtr node, const char *source);
   void try_flush_node_pmc(FileNodePtr node, const char *source);
   void clear_from_pmc(FileNodePtr node);
-  void flush_to_pmc(FileNodePtr node, bool new_remote, bool new_local, bool new_generate);
+  void flush_to_pmc(FileNodePtr node, bool new_remote, bool new_local, bool new_generate, const char *source);
   void load_from_pmc(FileNodePtr node, bool new_remote, bool new_local, bool new_generate);
 
   string get_persistent_id(const FullGenerateFileLocation &location);