1 files changed, 13 insertions, 9 deletions

diff --git a/libdino/src/service/message_storage.vala b/libdino/src/service/message_storage.vala
index 8812d950..1cde1a95 100644
--- a/libdino/src/service/message_storage.vala
+++ b/libdino/src/service/message_storage.vala
@@ -11,7 +11,7 @@ public class MessageStorage : StreamInteractionModule, Object {
     private StreamInteractor stream_interactor;
     private Database db;
 
-    private HashMap<Conversation, Gee.List<Message>> messages = new HashMap<Conversation, Gee.List<Entities.Message>>(Conversation.hash_func, Conversation.equals_func);
+    private HashMap<Conversation, Gee.TreeSet<Message>> messages = new HashMap<Conversation, Gee.TreeSet<Message>>(Conversation.hash_func, Conversation.equals_func);
 
     public static void start(StreamInteractor stream_interactor, Database db) {
         MessageStorage m = new MessageStorage(stream_interactor, db);
@@ -31,16 +31,18 @@ public class MessageStorage : StreamInteractionModule, Object {
 
     public Gee.List<Message> get_messages(Conversation conversation, int count = 50) {
         init_conversation(conversation);
-        if (messages[conversation].size > 0) {
-            return messages[conversation][int.max(messages[conversation].size - count - 1, 0) : messages[conversation].size];
+        Gee.List<Message> ret = new ArrayList<Message>(Message.equals_func);
+        foreach (Message message in messages[conversation]) {
+            if (ret.size >= count) break;
+            ret.add(message);
         }
-        return new ArrayList<Message>();
+        return ret;
     }
 
     public Message? get_last_message(Conversation conversation) {
         init_conversation(conversation);
         if (messages[conversation].size > 0) {
-            return messages[conversation][messages[conversation].size - 1];
+            return messages[conversation].first();
         }
         return null;
     }
@@ -52,17 +54,19 @@ public class MessageStorage : StreamInteractionModule, Object {
 
     public Message? get_message_by_id(string stanza_id, Conversation conversation) {
         init_conversation(conversation);
-        for (int i = messages[conversation].size - 1; i > 0; i--) {
-            if (messages[conversation][i].stanza_id == stanza_id) return messages[conversation][i];
+        foreach (Message message in messages[conversation]) {
+            if (message.stanza_id == stanza_id) return message;
         }
         return null;
     }
 
     private void init_conversation(Conversation conversation) {
         if (!messages.has_key(conversation)) {
-            messages[conversation] = db.get_messages(conversation.counterpart, conversation.account, Util.get_message_type_for_conversation(conversation), 50, null);
+            messages[conversation] = new Gee.TreeSet<Message>((a, b) => { return -1 * a.local_time.compare(b.local_time); });
+            Gee.List<Message> db_messages = db.get_messages(conversation.counterpart, conversation.account, Util.get_message_type_for_conversation(conversation), 50, null);
+            messages[conversation].add_all(db_messages);
         }
     }
 }
 
-}
\ No newline at end of file
+}