From 17091ee8e607875bcebf00db761067aedf3e9781 Mon Sep 17 00:00:00 2001
From: Neil Alexander <neilalexander@users.noreply.github.com>
Date: Wed, 15 Jun 2022 10:45:44 +0100
Subject: [PATCH] One more time

---
 roomserver/storage/shared/storage.go | 48 +++++++++++++++-------------
 1 file changed, 25 insertions(+), 23 deletions(-)

diff --git a/roomserver/storage/shared/storage.go b/roomserver/storage/shared/storage.go
index 6c5a21659..fa97a50fc 100644
--- a/roomserver/storage/shared/storage.go
+++ b/roomserver/storage/shared/storage.go
@@ -435,32 +435,25 @@ func (d *Database) Events(
 func (d *Database) events(
 	ctx context.Context, txn *sql.Tx, inputEventNIDs []types.EventNID,
 ) ([]types.Event, error) {
-	results := make([]types.Event, len(inputEventNIDs))
-	positions := make(map[types.EventNID]int, len(inputEventNIDs))
-	inCache := make([]bool, len(inputEventNIDs))
-	inDatabase := make([]types.EventNID, 0, len(results))
-	for i, nid := range inputEventNIDs {
+	events := make(map[types.EventNID]*gomatrixserverlib.Event, len(inputEventNIDs))
+	eventNIDs := make([]types.EventNID, 0, len(inputEventNIDs))
+	for _, nid := range inputEventNIDs {
 		if event, ok := d.Cache.GetRoomServerEvent(nid); ok && event != nil {
-			results[i] = types.Event{
-				EventNID: nid,
-				Event:    event,
-			}
-			inCache[i] = true
+			events[nid] = event
 		} else {
-			inDatabase = append(inDatabase, nid)
-			positions[nid] = i
+			eventNIDs = append(eventNIDs, nid)
 		}
 	}
-	eventJSONs, err := d.EventJSONTable.BulkSelectEventJSON(ctx, txn, inDatabase)
+	eventJSONs, err := d.EventJSONTable.BulkSelectEventJSON(ctx, txn, eventNIDs)
 	if err != nil {
 		return nil, err
 	}
-	eventIDs, _ := d.EventsTable.BulkSelectEventID(ctx, txn, inDatabase)
+	eventIDs, _ := d.EventsTable.BulkSelectEventID(ctx, txn, eventNIDs)
 	if err != nil {
 		eventIDs = map[types.EventNID]string{}
 	}
 	var roomNIDs map[types.EventNID]types.RoomNID
-	roomNIDs, err = d.EventsTable.SelectRoomNIDsForEventNIDs(ctx, txn, inDatabase)
+	roomNIDs, err = d.EventsTable.SelectRoomNIDsForEventNIDs(ctx, txn, eventNIDs)
 	if err != nil {
 		return nil, err
 	}
@@ -487,20 +480,29 @@ func (d *Database) events(
 		roomVersions[n] = v
 	}
 	for _, eventJSON := range eventJSONs {
-		eventNID := eventJSON.EventNID
-		result := &results[positions[eventNID]]
-		result.EventNID = eventNID
-		roomNID := roomNIDs[eventNID]
-		result.Event, err = gomatrixserverlib.NewEventFromTrustedJSONWithEventID(
-			eventIDs[eventNID], eventJSON.EventJSON, false, roomVersions[roomNID],
+		roomNID := roomNIDs[eventJSON.EventNID]
+		roomVersion := roomVersions[roomNID]
+		events[eventJSON.EventNID], err = gomatrixserverlib.NewEventFromTrustedJSONWithEventID(
+			eventIDs[eventJSON.EventNID], eventJSON.EventJSON, false, roomVersion,
 		)
 		if err != nil {
 			return nil, err
 		}
-		if result.Event != nil {
-			d.Cache.StoreRoomServerEvent(eventNID, result.Event)
+		if event := events[eventJSON.EventNID]; event != nil {
+			d.Cache.StoreRoomServerEvent(eventJSON.EventNID, event)
 		}
 	}
+	results := make([]types.Event, 0, len(inputEventNIDs))
+	for _, nid := range inputEventNIDs {
+		event, ok := events[nid]
+		if !ok {
+			panic("missing event")
+		}
+		results = append(results, types.Event{
+			EventNID: nid,
+			Event:    event,
+		})
+	}
 	if !redactionsArePermanent {
 		d.applyRedactions(results)
 	}