1 files changed, 31 insertions, 9 deletions
diff --git a/Swift/Controllers/Chat/ChatController.cpp b/Swift/Controllers/Chat/ChatController.cpp
index 22ef68d..d8c9e31 100644
--- a/Swift/Controllers/Chat/ChatController.cpp
+++ b/Swift/Controllers/Chat/ChatController.cpp
@@ -11,6 +11,7 @@
 
 #include <Swift/Controllers/Intl.h>
 #include <Swiften/Base/format.h>
+#include <Swiften/Base/Algorithm.h>
 #include "Swiften/Avatars/AvatarManager.h"
 #include "Swiften/Chat/ChatStateNotifier.h"
 #include "Swiften/Chat/ChatStateTracker.h"
@@ -102,8 +103,23 @@ void ChatController::preHandleIncomingMessage(boost::shared_ptr<MessageEvent> me
 			setToJID(from);
 		}
 	}
-	chatStateNotifier_->receivedMessageFromContact(message->getPayload<ChatState>());
+	boost::shared_ptr<Replace> replace = message->getPayload<Replace>();
+	if (replace) {
+		// Determine the timestamp
+		boost::posix_time::ptime timeStamp = boost::posix_time::microsec_clock::universal_time();
+		boost::optional<boost::posix_time::ptime> messageTimeStamp = getMessageTimestamp(message);
+		if (messageTimeStamp) {
+			timeStamp = *messageTimeStamp;
+		}
+		std::string body = message->getBody();
+		chatWindow_->replaceMessage(body, lastMessageUIID_, timeStamp);
+		replacedMessage_ = true;
+	}
+	else {
+		replacedMessage_ = false;
+	}
 	chatStateTracker_->handleMessageReceived(message);
+	chatStateNotifier_->receivedMessageFromContact(message->getPayload<ChatState>());
 }
 
 void ChatController::postHandleIncomingMessage(boost::shared_ptr<MessageEvent> messageEvent) {
@@ -116,21 +132,27 @@ void ChatController::preSendMessageRequest(boost::shared_ptr<Message> message) {
 }
 
 void ChatController::postSendMessage(const std::string& body, boost::shared_ptr<Stanza> sentStanza) {
-	std::string id = addMessage(body, QT_TRANSLATE_NOOP("", "me"), true, labelsEnabled_ ? chatWindow_->getSelectedSecurityLabel().getLabel() : boost::shared_ptr<SecurityLabel>(), std::string(avatarManager_->getAvatarPath(selfJID_).string()), boost::posix_time::microsec_clock::universal_time());
-	if (stanzaChannel_->getStreamManagementEnabled()) {
-		chatWindow_->setAckState(id, ChatWindow::Pending);
-		unackedStanzas_[sentStanza] = id;
+	boost::shared_ptr<Replace> replace = sentStanza->getPayload<Replace>();
+	if (replace) {
+		chatWindow_->replaceMessage(body, myLastMessageUIID_, boost::posix_time::microsec_clock::universal_time());
+		eraseIf(unackedStanzas_, PairSecondEquals<boost::shared_ptr<Stanza>, std::string>(myLastMessageUIID_));
+	} else {
+		myLastMessageUIID_ = addMessage(body, QT_TRANSLATE_NOOP("", "me"), true, labelsEnabled_ ? chatWindow_->getSelectedSecurityLabel().getLabel() : boost::shared_ptr<SecurityLabel>(), std::string(avatarManager_->getAvatarPath(selfJID_).string()), boost::posix_time::microsec_clock::universal_time());
+	}
+	if (stanzaChannel_->getStreamManagementEnabled() && !myLastMessageUIID_.empty() ) {
+		chatWindow_->setAckState(myLastMessageUIID_, ChatWindow::Pending);
+		unackedStanzas_[sentStanza] = myLastMessageUIID_;
 	}
 	lastWasPresence_ = false;
 	chatStateNotifier_->userSentMessage();
 }
 
 void ChatController::handleStanzaAcked(boost::shared_ptr<Stanza> stanza) {
-	std::string id = unackedStanzas_[stanza];
-	if (id != "") {
-		chatWindow_->setAckState(id, ChatWindow::Received);
+	std::map<boost::shared_ptr<Stanza>, std::string>::iterator unackedStanza = unackedStanzas_.find(stanza);
+	if (unackedStanza != unackedStanzas_.end()) {
+		chatWindow_->setAckState(unackedStanza->second, ChatWindow::Received);
+		unackedStanzas_.erase(unackedStanza);
 	}
-	unackedStanzas_.erase(unackedStanzas_.find(stanza));
 }
 
 void ChatController::setOnline(bool online) {