From e1e34b899492e6e8284c5c524f519559f38af0c8 Mon Sep 17 00:00:00 2001
From: Neil Alexander <neilalexander@users.noreply.github.com>
Date: Mon, 11 Jan 2021 12:47:25 +0000
Subject: [PATCH] Deep-checking forward extremities (#1698)

---
 .../internal/input/input_latest_events.go     | 49 +++++++++----------
 1 file changed, 23 insertions(+), 26 deletions(-)

diff --git a/roomserver/internal/input/input_latest_events.go b/roomserver/internal/input/input_latest_events.go
index 9554bf33..e5ec8c4d 100644
--- a/roomserver/internal/input/input_latest_events.go
+++ b/roomserver/internal/input/input_latest_events.go
@@ -259,34 +259,8 @@ func (u *latestEventsUpdater) calculateLatest(
 	// First of all, get a list of all of the events in our current
 	// set of forward extremities.
 	existingRefs := make(map[string]*types.StateAtEventAndReference)
-	existingNIDs := make([]types.EventNID, len(oldLatest))
 	for i, old := range oldLatest {
 		existingRefs[old.EventID] = &oldLatest[i]
-		existingNIDs[i] = old.EventNID
-	}
-
-	// Look up the old extremity events. This allows us to find their
-	// prev events.
-	events, err := u.api.DB.Events(u.ctx, existingNIDs)
-	if err != nil {
-		return false, fmt.Errorf("u.api.DB.Events: %w", err)
-	}
-
-	// Make a list of all of the prev events as referenced by all of
-	// the current forward extremities.
-	existingPrevs := make(map[string]struct{})
-	for _, old := range events {
-		for _, prevEventID := range old.PrevEventIDs() {
-			existingPrevs[prevEventID] = struct{}{}
-		}
-	}
-
-	// If the "new" event is already referenced by a forward extremity
-	// then do nothing - it's not a candidate to be a new extremity if
-	// it has been referenced.
-	if _, ok := existingPrevs[newEvent.EventID()]; ok {
-		u.latest = oldLatest
-		return false, nil
 	}
 
 	// If the "new" event is already a forward extremity then stop, as
@@ -296,6 +270,29 @@ func (u *latestEventsUpdater) calculateLatest(
 		return false, nil
 	}
 
+	// If the "new" event is already referenced by an existing event
+	// then do nothing - it's not a candidate to be a new extremity if
+	// it has been referenced.
+	if referenced, err := u.updater.IsReferenced(newEvent.EventReference()); err != nil {
+		return false, fmt.Errorf("u.updater.IsReferenced(new): %w", err)
+	} else if referenced {
+		u.latest = oldLatest
+		return false, nil
+	}
+
+	// Then let's see if any of the existing forward extremities now
+	// have entries in the previous events table. If they do then we
+	// will no longer include them as forward extremities.
+	existingPrevs := make(map[string]struct{})
+	for _, l := range existingRefs {
+		referenced, err := u.updater.IsReferenced(l.EventReference)
+		if err != nil {
+			return false, fmt.Errorf("u.updater.IsReferenced: %w", err)
+		} else if referenced {
+			existingPrevs[l.EventID] = struct{}{}
+		}
+	}
+
 	// Include our new event in the extremities.
 	newLatest := []types.StateAtEventAndReference{newStateAndRef}