From f55f450850796b94e5fb18341241659e61f46a04 Mon Sep 17 00:00:00 2001
From: Bruno Windels <bruno@windels.cloud>
Date: Mon, 27 Sep 2021 16:27:52 +0200
Subject: [PATCH] fix findExistingKeys too many (existing but not requested)
 keys

---
 src/matrix/storage/idb/QueryTarget.ts         | 31 ++++++---
 .../storage/idb/stores/TimelineEventStore.ts  | 64 +++++++++++++++++++
 2 files changed, 87 insertions(+), 8 deletions(-)

diff --git a/src/matrix/storage/idb/QueryTarget.ts b/src/matrix/storage/idb/QueryTarget.ts
index f2b6afb5..402140f2 100644
--- a/src/matrix/storage/idb/QueryTarget.ts
+++ b/src/matrix/storage/idb/QueryTarget.ts
@@ -160,20 +160,35 @@ export class QueryTarget<T> {
 
     /**
      * Checks if a given set of keys exist.
-     * Calls `callback(key, found)` for each key in `keys`, in key sorting order (or reversed if backwards=true).
      * If the callback returns true, the search is halted and callback won't be called again.
-     * `callback` is called with the same instances of the key as given in `keys`, so direct comparison can be used.
      */
     async findExistingKeys(keys: IDBValidKey[], backwards: boolean, callback: (key: IDBValidKey, pk: IDBValidKey) => boolean): Promise<void> {
+        const compareKeys = (a, b) => backwards ? -this.idbFactory.cmp(a, b) : this.idbFactory.cmp(a, b);
+        const sortedKeys = keys.slice().sort(compareKeys);
+        const firstKey = sortedKeys[0];
+        const lastKey = sortedKeys[sortedKeys.length - 1];
+        console.log(firstKey, lastKey, sortedKeys);
         const direction = backwards ? "prev" : "next";
-        const sortedKeys = keys.slice().sort((a, b) => backwards ? -this.idbFactory.cmp(a, b) : this.idbFactory.cmp(a, b));
-        const firstKey = backwards ? sortedKeys[sortedKeys.length - 1] : sortedKeys[0];
-        const lastKey = backwards ? sortedKeys[0] : sortedKeys[sortedKeys.length - 1];
         const cursor = this._target.openKeyCursor(this.IDBKeyRange.bound(firstKey, lastKey), direction);
+        let index = 0;
         await iterateCursor(cursor, (value, key, cursor) => {
-            const pk = cursor.primaryKey;
-            const done = callback(key, pk);
-            return done ? DONE : NOT_DONE;
+            while (index < sortedKeys.length && compareKeys(sortedKeys[index], key) < 0) {
+                index += 1;
+            }
+            let done = false;
+            if (sortedKeys[index] === key) {
+                const pk = cursor.primaryKey;
+                done = callback(key, pk);
+                index += 1;
+            }
+            if (done || index >= sortedKeys.length) {
+                return DONE;
+            } else {
+                return {
+                    done: false,
+                    jumpTo: sortedKeys[index],
+                }
+            }
         });
     }
 
diff --git a/src/matrix/storage/idb/stores/TimelineEventStore.ts b/src/matrix/storage/idb/stores/TimelineEventStore.ts
index 7338f96a..d3021a88 100644
--- a/src/matrix/storage/idb/stores/TimelineEventStore.ts
+++ b/src/matrix/storage/idb/stores/TimelineEventStore.ts
@@ -316,3 +316,67 @@ export class TimelineEventStore {
         this._timelineStore.delete(range);
     }
 }
+
+import {createMockStorage} from "../../../../mocks/Storage.js";
+import {createEvent, withTextBody} from "../../../../mocks/event.js";
+import {createEventEntry} from "../../../room/timeline/persistence/common.js";
+
+export function tests() {
+
+    const sortedIds = [
+        "$2wZy1W-QdcwaAwz68nfz1oc-3SsZKVDy8d86ERP1Pm0",
+        "$4RWaZ5142grUgTnQyr_5qiPTOwzAOimt5MsXg6m1diM",
+        "$4izqHE2Wf5US_-e_za942pZ10CDNJjDncUMmhqBUVQw",
+        "$Oil2Afq2cBLqMAeJTAHjA3Is9T5Wmaa2ogVRlFJ_gzE",
+        "$Wyl-7u-YqnPJElkPufIRXRFTYP-eFxQ4iD-SmLQo2Rw",
+        "$b-eWaZtp22vL9mp0h7odbpphOZQ-rnp54qjyTQPARgo",
+        "$sS9rTv8u2m9o4RaMI2jGOnpMtb9t8_0euiQLhNFW380",
+        "$uZLkB9rzTKvJAK2QrQNX-prwQ2Niajdi0fvvRnyCtz8",
+        "$vGecIBZFex9_vlQf1E1LjtQXE3q5GwERIHMiy4mOWv0",
+        "$vdLgAnwjHj0cicU3MA4ynLHUBGOIFhvvksY3loqzjF",
+    ];
+
+    const insertedIds = [
+        sortedIds[5],
+        sortedIds[3],
+        sortedIds[9],
+        sortedIds[7],
+        sortedIds[1],
+    ];
+
+    const checkedIds = [
+        sortedIds[2],
+        sortedIds[4],
+        sortedIds[3],
+        sortedIds[0],
+        sortedIds[8],
+        sortedIds[9],
+        sortedIds[6],
+    ];
+
+    const roomId = "!fjsdf423423jksdfdsf:hs.tld";
+
+    function createEventWithId(id) {
+        return withTextBody("hello", createEvent("m.room.message", id, "@alice:hs.tld"));
+    }
+
+    return {
+        "getEventKeysForIds": async assert => {
+            const storage = await createMockStorage();
+            const txn = await storage.readWriteTxn([storage.storeNames.timelineEvents]);
+            let eventKey = EventKey.defaultLiveKey;
+            for (const insertedId of insertedIds) {
+                assert(await txn.timelineEvents.tryInsert(createEventEntry(eventKey.nextKey(), roomId, createEventWithId(insertedId))));
+                eventKey = eventKey.nextKey();
+            }
+            const eventKeyMap = await txn.timelineEvents.getEventKeysForIds(roomId, checkedIds);
+            assert.equal(eventKeyMap.size, 2);
+            const eventKey1 = eventKeyMap.get("$Oil2Afq2cBLqMAeJTAHjA3Is9T5Wmaa2ogVRlFJ_gzE");
+            assert.equal(eventKey1.fragmentId, 0);
+            assert.equal(eventKey1.eventIndex, 80000001);
+            const eventKey2 = eventKeyMap.get("$vdLgAnwjHj0cicU3MA4ynLHUBGOIFhvvksY3loqzjF");
+            assert.equal(eventKey2.fragmentId, 0);
+            assert.equal(eventKey2.eventIndex, 80000002);
+        }
+    }
+}