1 files changed, 17 insertions, 10 deletions
diff --git a/Swift/Controllers/Chat/ChatController.cpp b/Swift/Controllers/Chat/ChatController.cpp
index 22ef68d..513b446 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,8 @@ void ChatController::preHandleIncomingMessage(boost::shared_ptr<MessageEvent> me
 			setToJID(from);
 		}
 	}
-	chatStateNotifier_->receivedMessageFromContact(message->getPayload<ChatState>());
 	chatStateTracker_->handleMessageReceived(message);
+	chatStateNotifier_->receivedMessageFromContact(message->getPayload<ChatState>());
 }
 
 void ChatController::postHandleIncomingMessage(boost::shared_ptr<MessageEvent> messageEvent) {
@@ -116,27 +117,33 @@ 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) {
+		eraseIf(unackedStanzas_, PairSecondEquals<boost::shared_ptr<Stanza>, std::string>(myLastMessageUIID_));
+		chatWindow_->replaceMessage(body, myLastMessageUIID_, boost::posix_time::microsec_clock::universal_time());
+	} 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) {
 	if (!online) {
 		std::map<boost::shared_ptr<Stanza>, std::string>::iterator it = unackedStanzas_.begin();
-		for ( ; it != unackedStanzas_.end(); it++) {
+		for ( ; it != unackedStanzas_.end(); ++it) {
 			chatWindow_->setAckState(it->second, ChatWindow::Failed);
 		}
 		unackedStanzas_.clear();