Always write local drafts the same way.

Telegram/SourceFiles/history/history_widget.cpp

@@ -1440,14 +1440,7 @@ void HistoryWidget::saveCloudDraft() {
 void HistoryWidget::writeDraftTexts() {
 	Expects(_history != nullptr);
 
-	session().local().writeDrafts(
+	session().local().writeDrafts(_history);
-		_history,
-		_editMsgId ? Data::DraftKey::LocalEdit() : Data::DraftKey::Local(),
-		Storage::MessageDraft{
-			_editMsgId ? _editMsgId : _replyToId,
-			_field->getTextWithTags(),
-			_previewState,
-		});
 	if (_migrated) {
 		_migrated->clearDrafts();
 		session().local().writeDrafts(_migrated);
@@ -1457,10 +1450,7 @@ void HistoryWidget::writeDraftTexts() {
 void HistoryWidget::writeDraftCursors() {
 	Expects(_history != nullptr);
 
-	session().local().writeDraftCursors(
+	session().local().writeDraftCursors(_history);
-		_history,
-		_editMsgId ? Data::DraftKey::LocalEdit() : Data::DraftKey::Local(),
-		MessageCursor(_field));
 	if (_migrated) {
 		_migrated->clearDrafts();
 		session().local().writeDraftCursors(_migrated);
@@ -1691,7 +1681,8 @@ void HistoryWidget::applyDraft(FieldHistoryAction fieldHistoryAction) {
 		clearFieldText(0, fieldHistoryAction);
 		_field->setFocus();
 		_replyEditMsg = nullptr;
-		_editMsgId = _replyToId = 0;
+		_replyToId = 0;
+		setEditMsgId(0);
 		if (fieldWillBeHiddenAfterEdit) {
 			updateControlsVisibility();
 			updateControlsGeometry();
@@ -1715,11 +1706,11 @@ void HistoryWidget::applyDraft(FieldHistoryAction fieldHistoryAction) {
 
 	_replyEditMsg = nullptr;
 	if (const auto editDraft = _history->localEditDraft()) {
-		_editMsgId = editDraft->msgId;
+		setEditMsgId(editDraft->msgId);
 		_replyToId = 0;
 	} else {
-		_editMsgId = 0;
 		_replyToId = readyToForward() ? 0 : _history->localDraft()->msgId;
+		setEditMsgId(0);
 	}
 	updateCmdStartShown();
 	updateControlsVisibility();
@@ -1875,7 +1866,7 @@ void HistoryWidget::showHistory(
 		_scrollToAnimation.stop();
 
 		clearAllLoadRequests();
-		_history = _migrated = nullptr;
+		setHistory(nullptr);
 		_list = nullptr;
 		_peer = nullptr;
 		_channel = NoChannel;
@@ -1944,8 +1935,7 @@ void HistoryWidget::showHistory(
 	_itemsRevealHeight = 0;
 
 	if (_peer) {
-		_history = _peer->owner().history(_peer);
+		setHistory(_peer->owner().history(_peer));
-		_migrated = _history->migrateFrom();
 		if (_migrated
 			&& !_migrated->isEmpty()
 			&& (!_history->loadedAtTop() || !_migrated->loadedAtBottom())) {
@@ -2072,6 +2062,56 @@ void HistoryWidget::showHistory(
 	crl::on_main(this, [=] { controller()->widget()->setInnerFocus(); });
 }
 
+void HistoryWidget::setHistory(History *history) {
+	if (_history == history) {
+		return;
+	}
+	unregisterDraftSources();
+	_history = history;
+	_migrated = _history ? _history->migrateFrom() : nullptr;
+	registerDraftSource();
+}
+
+void HistoryWidget::unregisterDraftSources() {
+	if (!_history) {
+		return;
+	}
+	session().local().unregisterDraftSource(
+		_history,
+		Data::DraftKey::Local());
+	session().local().unregisterDraftSource(
+		_history,
+		Data::DraftKey::LocalEdit());
+}
+
+void HistoryWidget::registerDraftSource() {
+	if (!_history) {
+		return;
+	}
+	const auto editMsgId = _editMsgId;
+	const auto draft = [=] {
+		return Storage::MessageDraft{
+			editMsgId ? editMsgId : _replyToId,
+			_field->getTextWithTags(),
+			_previewState,
+		};
+	};
+	auto draftSource = Storage::MessageDraftSource{
+		.draft = draft,
+		.cursor = [=] { return MessageCursor(_field); },
+	};
+	session().local().registerDraftSource(
+		_history,
+		editMsgId ? Data::DraftKey::LocalEdit() : Data::DraftKey::Local(),
+		std::move(draftSource));
+}
+
+void HistoryWidget::setEditMsgId(MsgId msgId) {
+	unregisterDraftSources();
+	_editMsgId = msgId;
+	registerDraftSource();
+}
+
 void HistoryWidget::clearDelayedShowAt() {
 	_delayedShowAtMsgId = -1;
 	clearDelayedShowAtRequest();
@@ -6122,7 +6162,7 @@ void HistoryWidget::cancelEdit() {
 	}
 
 	_replyEditMsg = nullptr;
-	_editMsgId = 0;
+	setEditMsgId(0);
 	_history->clearLocalEditDraft();
 	applyDraft();
 
@@ -6998,6 +7038,7 @@ HistoryWidget::~HistoryWidget() {
 		session().api().saveDraftToCloudDelayed(_history);
 
 		clearAllLoadRequests();
+		unregisterDraftSources();
 	}
 	setTabbedPanel(nullptr);
 }

Telegram/SourceFiles/history/history_widget.h

@@ -558,6 +558,11 @@ private:
 		TextUpdateEvents events = 0,
 		FieldHistoryAction fieldHistoryAction = FieldHistoryAction::Clear);
 
+	void unregisterDraftSources();
+	void registerDraftSource();
+	void setHistory(History *history);
+	void setEditMsgId(MsgId msgId);
+
 	HistoryItem *getItemFromHistoryOrMigrated(MsgId genericMsgId) const;
 	void animatedScrollToItem(MsgId msgId);
 	void animatedScrollToY(int scrollTo, HistoryItem *attachTo = nullptr);

Telegram/SourceFiles/history/view/controls/history_view_compose_controls.cpp

@@ -626,6 +626,7 @@ ComposeControls::ComposeControls(
 
 ComposeControls::~ComposeControls() {
 	saveFieldToHistoryLocalDraft();
+	unregisterDraftSources();
 	setTabbedPanel(nullptr);
 	session().api().request(_inlineBotResolveRequestId).cancel();
 }
@@ -650,7 +651,9 @@ void ComposeControls::setHistory(SetHistoryArgs &&args) {
 	//if (_history == history) {
 	//	return;
 	//}
+	unregisterDraftSources();
 	_history = history;
+	registerDraftSource();
 	_window->tabbedSelector()->setCurrentPeer(
 		history ? history->peer.get() : nullptr);
 	initWebpageProcess();
@@ -877,7 +880,11 @@ TextWithTags ComposeControls::getTextWithAppliedMarkdown() const {
 }
 
 void ComposeControls::clear() {
-	setText({});
+	// Otherwise cancelReplyMessage() will save the draft.
+	const auto saveTextDraft = !replyingToMessage();
+	setFieldText(
+		{},
+		saveTextDraft ? TextUpdateEvent::SaveDraft : TextUpdateEvent());
 	cancelReplyMessage();
 }
 
@@ -993,7 +1000,9 @@ void ComposeControls::init() {
 
 	_header->editMsgId(
 	) | rpl::start_with_next([=](const auto &id) {
+		unregisterDraftSources();
 		updateSendButtonType();
+		registerDraftSource();
 	}, _wrap->lifetime());
 
 	_header->previewCancelled(
@@ -1396,23 +1405,51 @@ void ComposeControls::saveDraft(bool delayed) {
 void ComposeControls::writeDraftTexts() {
 	Expects(_history != nullptr);
 
-	session().local().writeDrafts(
+	session().local().writeDrafts(_history);
-		_history,
-		draftKeyCurrent(),
-		Storage::MessageDraft{
-			_header->getDraftMessageId(),
-			_field->getTextWithTags(),
-			_previewState,
-		});
 }
 
 void ComposeControls::writeDraftCursors() {
 	Expects(_history != nullptr);
 
-	session().local().writeDraftCursors(
+	session().local().writeDraftCursors(_history);
-		_history,
+}
-		draftKeyCurrent(),
+
-		MessageCursor(_field));
+void ComposeControls::unregisterDraftSources() {
+	if (!_history) {
+		return;
+	}
+	const auto normal = draftKey(DraftType::Normal);
+	const auto edit = draftKey(DraftType::Edit);
+	if (normal != Data::DraftKey::None()) {
+		session().local().unregisterDraftSource(_history, normal);
+	}
+	if (edit != Data::DraftKey::None()) {
+		session().local().unregisterDraftSource(_history, edit);
+	}
+}
+
+void ComposeControls::registerDraftSource() {
+	if (!_history) {
+		return;
+	}
+	const auto key = draftKeyCurrent();
+	if (key != Data::DraftKey::None()) {
+		const auto draft = [=] {
+			return Storage::MessageDraft{
+				_header->getDraftMessageId(),
+				_field->getTextWithTags(),
+				_previewState,
+			};
+		};
+		auto draftSource = Storage::MessageDraftSource{
+			.draft = draft,
+			.cursor = [=] { return MessageCursor(_field); },
+		};
+		session().local().registerDraftSource(
+			_history,
+			key,
+			std::move(draftSource));
+	}
 }
 
 void ComposeControls::writeDrafts() {

Telegram/SourceFiles/history/view/controls/history_view_compose_controls.h

@@ -270,6 +270,9 @@ private:
 		FieldHistoryAction fieldHistoryAction = FieldHistoryAction::Clear);
 	void saveFieldToHistoryLocalDraft();
 
+	void unregisterDraftSources();
+	void registerDraftSource();
+
 	const not_null<QWidget*> _parent;
 	const not_null<Window::SessionController*> _window;
 	History *_history = nullptr;

Telegram/SourceFiles/storage/storage_account.cpp

@@ -52,6 +52,7 @@ constexpr auto kSinglePeerTypeChannel = qint32(3);
 constexpr auto kSinglePeerTypeSelf = qint32(4);
 constexpr auto kSinglePeerTypeEmpty = qint32(0);
 constexpr auto kMultiDraftTag = quint64(0xFFFFFFFFFFFFFF01ULL);
+constexpr auto kMultiDraftCursorsTag = quint64(0xFFFFFFFFFFFFFF02ULL);
 
 enum { // Local Storage Keys
 	lskUserMap = 0x00,
@@ -79,6 +80,14 @@ enum { // Local Storage Keys
 	lskMasksKeys = 0x16, // no data
 };
 
+auto EmptyMessageDraftSources()
+-> const base::flat_map<Data::DraftKey, MessageDraftSource> & {
+	static const auto result = base::flat_map<
+		Data::DraftKey,
+		MessageDraftSource>();
+	return result;
+}
+
 [[nodiscard]] FileKey ComputeDataNameKey(const QString &dataName) {
 	// We dropped old test authorizations when migrated to multi auth.
 	//const auto testAddition = (cTestMode() ? qsl(":/test/") : QString());
@@ -957,12 +966,10 @@ void EnumerateDrafts(
 		const Data::HistoryDrafts &map,
 		Data::Draft *cloudDraft,
 		bool supportMode,
-		Data::DraftKey replaceKey,
+		const base::flat_map<Data::DraftKey, MessageDraftSource> &sources,
-		const MessageDraft &replaceDraft,
-		const MessageCursor &replaceCursor,
 		Callback &&callback) {
 	for (const auto &[key, draft] : map) {
-		if (key == Data::DraftKey::Cloud() || key == replaceKey) {
+		if (key == Data::DraftKey::Cloud() || sources.contains(key)) {
 			continue;
 		} else if (key == Data::DraftKey::Local()
 			&& !supportMode
@@ -976,23 +983,45 @@ void EnumerateDrafts(
 			draft->previewState,
 			draft->cursor);
 	}
-	if (replaceKey
+	for (const auto &[key, source] : sources) {
-		&& (replaceDraft.msgId
+		const auto draft = source.draft();
-			|| !replaceDraft.textWithTags.text.isEmpty()
+		const auto cursor = source.cursor();
-			|| replaceCursor != MessageCursor())) {
+		if (draft.msgId
-		callback(
+			|| !draft.textWithTags.text.isEmpty()
-			replaceKey,
+			|| cursor != MessageCursor()) {
-			replaceDraft.msgId,
+			callback(
-			replaceDraft.textWithTags,
+				key,
-			replaceDraft.previewState,
+				draft.msgId,
-			replaceCursor);
+				draft.textWithTags,
+				draft.previewState,
+				cursor);
+		}
 	}
 }
 
-void Account::writeDrafts(
+void Account::registerDraftSource(
 		not_null<History*> history,
-		Data::DraftKey replaceKey,
+		Data::DraftKey key,
-		MessageDraft replaceDraft) {
+		MessageDraftSource source) {
+	Expects(source.draft != nullptr);
+	Expects(source.cursor != nullptr);
+
+	_draftSources[history][key] = std::move(source);
+}
+
+void Account::unregisterDraftSource(
+		not_null<History*> history,
+		Data::DraftKey key) {
+	const auto i = _draftSources.find(history);
+	if (i != _draftSources.end()) {
+		i->second.remove(key);
+		if (i->second.empty()) {
+			_draftSources.erase(i);
+		}
+	}
+}
+
+void Account::writeDrafts(not_null<History*> history) {
 	const auto peerId = history->peer->id;
 	const auto &map = history->draftsMap();
 	const auto cloudIt = map.find(Data::DraftKey::Cloud());
@@ -1000,14 +1029,16 @@ void Account::writeDrafts(
 		? cloudIt->second.get()
 		: nullptr;
 	const auto supportMode = _owner->session().supportMode();
+	const auto sourcesIt = _draftSources.find(history);
+	const auto &sources = (sourcesIt != _draftSources.end())
+		? sourcesIt->second
+		: EmptyMessageDraftSources();
 	auto count = 0;
 	EnumerateDrafts(
 		map,
 		cloudDraft,
 		supportMode,
-		replaceKey,
+		sources,
-		replaceDraft,
-		MessageCursor(),
 		[&](auto&&...) { ++count; });
 	if (!count) {
 		auto i = _draftsMap.find(peerId);
@@ -1043,9 +1074,7 @@ void Account::writeDrafts(
 		map,
 		cloudDraft,
 		supportMode,
-		replaceKey,
+		sources,
-		replaceDraft,
-		MessageCursor(),
 		sizeCallback);
 
 	EncryptedDescriptor data(size);
@@ -1071,9 +1100,7 @@ void Account::writeDrafts(
 		map,
 		cloudDraft,
 		supportMode,
-		replaceKey,
+		sources,
-		replaceDraft,
-		MessageCursor(),
 		writeCallback);
 
 	FileWriteDescriptor file(i->second, _basePath);
@@ -1082,10 +1109,7 @@ void Account::writeDrafts(
 	_draftsNotReadMap.remove(peerId);
 }
 
-void Account::writeDraftCursors(
+void Account::writeDraftCursors(not_null<History*> history) {
-		not_null<History*> history,
-		Data::DraftKey replaceKey,
-		MessageCursor replaceCursor) {
 	const auto peerId = history->peer->id;
 	const auto &map = history->draftsMap();
 	const auto cloudIt = map.find(Data::DraftKey::Cloud());
@@ -1093,14 +1117,16 @@ void Account::writeDraftCursors(
 		? cloudIt->second.get()
 		: nullptr;
 	const auto supportMode = _owner->session().supportMode();
+	const auto sourcesIt = _draftSources.find(history);
+	const auto &sources = (sourcesIt != _draftSources.end())
+		? sourcesIt->second
+		: EmptyMessageDraftSources();
 	auto count = 0;
 	EnumerateDrafts(
 		map,
 		cloudDraft,
 		supportMode,
-		replaceKey,
+		sources,
-		MessageDraft(),
-		replaceCursor,
 		[&](auto&&...) { ++count; });
 	if (!count) {
 		clearDraftCursors(peerId);
@@ -1112,21 +1138,24 @@ void Account::writeDraftCursors(
 		writeMapQueued();
 	}
 
-	auto size = int(sizeof(quint64) * 2 + sizeof(quint32) * 4);
+	auto size = int(sizeof(quint64) * 2
+		+ sizeof(quint32)
+		+ sizeof(qint32) * 4 * count);
 
 	EncryptedDescriptor data(size);
 	data.stream
-		<< quint64(kMultiDraftTag)
+		<< quint64(kMultiDraftCursorsTag)
 		<< SerializePeerId(peerId)
 		<< quint32(count);
 
 	const auto writeCallback = [&](
-			auto&&, // key
+			const Data::DraftKey &key,
 			MsgId, // msgId
 			auto&&, // text
 			Data::PreviewState,
 			const MessageCursor &cursor) { // cursor
 		data.stream
+			<< key.serialize()
 			<< qint32(cursor.position)
 			<< qint32(cursor.anchor)
 			<< qint32(cursor.scroll);
@@ -1135,9 +1164,7 @@ void Account::writeDraftCursors(
 		map,
 		cloudDraft,
 		supportMode,
-		replaceKey,
+		sources,
-		MessageDraft(),
-		replaceCursor,
 		writeCallback);
 
 	FileWriteDescriptor file(i->second, _basePath);
@@ -1166,7 +1193,7 @@ void Account::readDraftCursors(PeerId peerId, Data::HistoryDrafts &map) {
 	}
 	quint64 tag = 0;
 	draft.stream >> tag;
-	if (tag != kMultiDraftTag) {
+	if (tag != kMultiDraftTag && tag != kMultiDraftCursorsTag) {
 		readDraftCursorsLegacy(peerId, draft, tag, map);
 		return;
 	}
@@ -1178,10 +1205,18 @@ void Account::readDraftCursors(PeerId peerId, Data::HistoryDrafts &map) {
 		clearDraftCursors(peerId);
 		return;
 	}
+	const auto keysWritten = (tag == kMultiDraftCursorsTag);
 	for (auto i = 0; i != count; ++i) {
+		qint32 keyValue = 0;
+		if (keysWritten) {
+			draft.stream >> keyValue;
+		}
+		const auto key = keysWritten
+			? Data::DraftKey::FromSerialized(keyValue)
+			: Data::DraftKey::Local();
 		qint32 position = 0, anchor = 0, scroll = QFIXED_MAX;
 		draft.stream >> position >> anchor >> scroll;
-		if (const auto i = map.find(Data::DraftKey::Local()); i != end(map)) {
+		if (const auto i = map.find(key); i != end(map)) {
 			i->second->cursor = MessageCursor(position, anchor, scroll);
 		}
 	}

Telegram/SourceFiles/storage/storage_account.h

@@ -56,6 +56,11 @@ struct MessageDraft {
 	Data::PreviewState previewState = Data::PreviewState::Allowed;
 };
 
+struct MessageDraftSource {
+	Fn<MessageDraft()> draft;
+	Fn<MessageCursor()> cursor;
+};
+
 class Account final {
 public:
 	Account(not_null<Main::Account*> owner, const QString &dataName);
@@ -79,15 +84,16 @@ public:
 	void writeMtpData();
 	void writeMtpConfig();
 
-	void writeDrafts(
+	void registerDraftSource(
 		not_null<History*> history,
-		Data::DraftKey replaceKey = Data::DraftKey::None(),
+		Data::DraftKey key,
-		MessageDraft replaceDraft = MessageDraft());
+		MessageDraftSource source);
+	void unregisterDraftSource(
+		not_null<History*> history,
+		Data::DraftKey key);
+	void writeDrafts(not_null<History*> history);
 	void readDraftsWithCursors(not_null<History*> history);
-	void writeDraftCursors(
+	void writeDraftCursors(not_null<History*> history);
-		not_null<History*> history,
-		Data::DraftKey replaceKey = Data::DraftKey::None(),
-		MessageCursor replaceCursor = MessageCursor());
 	[[nodiscard]] bool hasDraftCursors(PeerId peerId);
 	[[nodiscard]] bool hasDraft(PeerId peerId);
 
@@ -242,6 +248,9 @@ private:
 	base::flat_map<PeerId, FileKey> _draftsMap;
 	base::flat_map<PeerId, FileKey> _draftCursorsMap;
 	base::flat_map<PeerId, bool> _draftsNotReadMap;
+	base::flat_map<
+		not_null<History*>,
+		base::flat_map<Data::DraftKey, MessageDraftSource>> _draftSources;
 
 	QMultiMap<MediaKey, Core::FileLocation> _fileLocations;
 	QMap<QString, QPair<MediaKey, Core::FileLocation>> _fileLocationPairs;