From 8108d98fa744129af809b1b193066d671e9ae252 Mon Sep 17 00:00:00 2001
From: Nicolas Werner <nicolas.werner@hotmail.de>
Date: Sat, 3 Apr 2021 13:15:35 +0200
Subject: [PATCH] Store state events with state keys

---
 src/Cache.cpp | 13 ++++++++-----
 src/Cache_p.h | 43 +++++++++++++++++++++++++++++++++++++++----
 2 files changed, 47 insertions(+), 9 deletions(-)

diff --git a/src/Cache.cpp b/src/Cache.cpp
index f3407e48..3e1057c0 100644
--- a/src/Cache.cpp
+++ b/src/Cache.cpp
@@ -1181,11 +1181,14 @@ Cache::saveState(const mtx::responses::Sync &res)
 
         // Save joined rooms
         for (const auto &room : res.rooms.join) {
-                auto statesdb  = getStatesDb(txn, room.first);
-                auto membersdb = getMembersDb(txn, room.first);
-
-                saveStateEvents(txn, statesdb, membersdb, room.first, room.second.state.events);
-                saveStateEvents(txn, statesdb, membersdb, room.first, room.second.timeline.events);
+                auto statesdb    = getStatesDb(txn, room.first);
+                auto stateskeydb = getStatesKeyDb(txn, room.first);
+                auto membersdb   = getMembersDb(txn, room.first);
+
+                saveStateEvents(
+                  txn, statesdb, stateskeydb, membersdb, room.first, room.second.state.events);
+                saveStateEvents(
+                  txn, statesdb, stateskeydb, membersdb, room.first, room.second.timeline.events);
 
                 saveTimelineMessages(txn, room.first, room.second.timeline);
 
diff --git a/src/Cache_p.h b/src/Cache_p.h
index 62927923..51dc4cbc 100644
--- a/src/Cache_p.h
+++ b/src/Cache_p.h
@@ -273,6 +273,18 @@ public:
                 return false;
         }
 
+        static int compare_state_key(const MDB_val *a, const MDB_val *b)
+        {
+                auto get_skey = [](const MDB_val *v) {
+                        return nlohmann::json::parse(
+                                 std::string_view(static_cast<const char *>(v->mv_data),
+                                                  v->mv_size))
+                          .value("key", "");
+                };
+
+                return get_skey(a).compare(get_skey(b));
+        }
+
 signals:
         void newReadReceipts(const QString &room_id, const std::vector<QString> &event_ids);
         void roomReadStatus(const std::map<QString, bool> &status);
@@ -323,17 +335,19 @@ private:
         template<class T>
         void saveStateEvents(lmdb::txn &txn,
                              lmdb::dbi &statesdb,
+                             lmdb::dbi &stateskeydb,
                              lmdb::dbi &membersdb,
                              const std::string &room_id,
                              const std::vector<T> &events)
         {
                 for (const auto &e : events)
-                        saveStateEvent(txn, statesdb, membersdb, room_id, e);
+                        saveStateEvent(txn, statesdb, stateskeydb, membersdb, room_id, e);
         }
 
         template<class T>
         void saveStateEvent(lmdb::txn &txn,
                             lmdb::dbi &statesdb,
+                            lmdb::dbi &stateskeydb,
                             lmdb::dbi &membersdb,
                             const std::string &room_id,
                             const T &event)
@@ -371,9 +385,22 @@ private:
                 }
 
                 std::visit(
-                  [&txn, &statesdb](auto e) {
-                          if (isStateEvent(e) && e.type != EventType::Unsupported)
-                                  statesdb.put(txn, to_string(e.type), json(e).dump());
+                  [&txn, &statesdb, &stateskeydb](auto e) {
+                          if constexpr (isStateEvent(e))
+                                  if (e.type != EventType::Unsupported) {
+                                          if (e.state_key.empty())
+                                                  statesdb.put(
+                                                    txn, to_string(e.type), json(e).dump());
+                                          else
+                                                  stateskeydb.put(
+                                                    txn,
+                                                    to_string(e.type),
+                                                    json::object({
+                                                                   {"key", e.state_key},
+                                                                   {"id", e.event_id},
+                                                                 })
+                                                      .dump());
+                                  }
                   },
                   event);
         }
@@ -461,6 +488,14 @@ private:
                 return lmdb::dbi::open(txn, std::string(room_id + "/state").c_str(), MDB_CREATE);
         }
 
+        lmdb::dbi getStatesKeyDb(lmdb::txn &txn, const std::string &room_id)
+        {
+                auto db =
+                  lmdb::dbi::open(txn, std::string(room_id + "/state_by_key").c_str(), MDB_CREATE);
+                lmdb::dbi_set_dupsort(txn, db, compare_state_key);
+                return db;
+        }
+
         lmdb::dbi getAccountDataDb(lmdb::txn &txn, const std::string &room_id)
         {
                 return lmdb::dbi::open(