Move use state pacakge to query state

src/github.com/matrix-org/dendrite/roomserver/input/events.go

@@ -16,9 +16,6 @@ type RoomEventDatabase interface {
 	// Returns an error if the there is an error talking to the database
 	// or if the event IDs aren't in the database.
 	StateEntriesForEventIDs(eventIDs []string) ([]types.StateEntry, error)
-	// Lookup the numeric IDs for a list of string event state keys.
-	// Returns a map from string state key to numeric ID for the state key.
-	EventStateKeyNIDs(eventStateKeys []string) (map[string]types.EventStateKeyNID, error)
 	// Lookup the Events for a list of numeric event IDs.
 	// Returns a sorted list of events.
 	Events(eventNIDs []types.EventNID) ([]types.Event, error)

src/github.com/matrix-org/dendrite/roomserver/query/query.go

@@ -3,6 +3,7 @@ package query
 import (
 	"encoding/json"
 	"github.com/matrix-org/dendrite/roomserver/api"
+	"github.com/matrix-org/dendrite/roomserver/state"
 	"github.com/matrix-org/dendrite/roomserver/types"
 	"github.com/matrix-org/gomatrixserverlib"
 	"github.com/matrix-org/util"
@@ -12,6 +13,7 @@ import (
 
 // RoomserverQueryAPIDatabase has the storage APIs needed to implement the query API.
 type RoomserverQueryAPIDatabase interface {
+	state.RoomStateDatabase
 	// Lookup the numeric ID for the room.
 	// Returns 0 if the room doesn't exists.
 	// Returns an error if there was a problem talking to the database.
@@ -19,20 +21,6 @@ type RoomserverQueryAPIDatabase interface {
 	// Lookup event references for the latest events in the room and the current state snapshot.
 	// Returns an error if there was a problem talking to the database.
 	LatestEventIDs(roomNID types.RoomNID) ([]gomatrixserverlib.EventReference, types.StateSnapshotNID, error)
-	// Lookup the numeric IDs for a list of string event types.
-	// Returns a map from string event type to numeric ID for the event type.
-	EventTypeNIDs(eventTypes []string) (map[string]types.EventTypeNID, error)
-	// Lookup the numeric IDs for a list of string event state keys.
-	// Returns a map from string state key to numeric ID for the state key.
-	EventStateKeyNIDs(eventStateKeys []string) (map[string]types.EventStateKeyNID, error)
-	// Lookup the numeric state block IDs for each numeric state snapshot ID
-	// The returned slice is sorted by numeric state snapshot ID.
-	StateBlockNIDs(stateNIDs []types.StateSnapshotNID) ([]types.StateBlockNIDList, error)
-	// Lookup the state data for the state key tuples for each numeric state block ID
-	// The returned slice is sorted by numeric state block ID.
-	StateEntriesForTuples(stateBlockNIDs []types.StateBlockNID, stateKeyTuples []types.StateKeyTuple) (
-		[]types.StateEntryList, error,
-	)
 	// Lookup the Events for a list of numeric event IDs.
 	// Returns a sorted list of events.
 	Events(eventNIDs []types.EventNID) ([]types.Event, error)
@@ -64,12 +52,12 @@ func (r *RoomserverQueryAPI) QueryLatestEventsAndState(
 	}
 
 	// Lookup the currrent state for the requested tuples.
-	stateTuples, err := stringTuplesToNumericTuples(r.DB, request.StateToFetch)
+	stateTuples, err := state.StringTuplesToNumericTuples(r.DB, request.StateToFetch)
 	if err != nil {
 		return err
 	}
 
-	stateEntries, err := loadStateAtSnapshotForTuples(r.DB, currentStateSnapshotNID, stateTuples)
+	stateEntries, err := state.LoadStateAtSnapshotForTuples(r.DB, currentStateSnapshotNID, stateTuples)
 	if err != nil {
 		return err
 	}

src/github.com/matrix-org/dendrite/roomserver/query/state.go

@@ -1,98 +0,0 @@
-package query
-
-import (
-	"github.com/matrix-org/dendrite/roomserver/api"
-	"github.com/matrix-org/dendrite/roomserver/types"
-	"github.com/matrix-org/util"
-	"sort"
-)
-
-func stringTuplesToNumericTuples(db RoomserverQueryAPIDatabase, stringTuples []api.StateKeyTuple) ([]types.StateKeyTuple, error) {
-	eventTypes := make([]string, len(stringTuples))
-	stateKeys := make([]string, len(stringTuples))
-	for i := range stringTuples {
-		eventTypes[i] = stringTuples[i].EventType
-		stateKeys[i] = stringTuples[i].EventStateKey
-	}
-	eventTypes = util.UniqueStrings(eventTypes)
-	eventTypeMap, err := db.EventTypeNIDs(eventTypes)
-	if err != nil {
-		return nil, err
-	}
-	stateKeys = util.UniqueStrings(stateKeys)
-	stateKeyMap, err := db.EventStateKeyNIDs(stateKeys)
-	if err != nil {
-		return nil, err
-	}
-
-	var result []types.StateKeyTuple
-	for _, stringTuple := range stringTuples {
-		var numericTuple types.StateKeyTuple
-		var ok1, ok2 bool
-		numericTuple.EventTypeNID, ok1 = eventTypeMap[stringTuple.EventType]
-		numericTuple.EventStateKeyNID, ok2 = stateKeyMap[stringTuple.EventStateKey]
-		if ok1 && ok2 {
-			result = append(result, numericTuple)
-		}
-	}
-
-	return result, nil
-}
-
-// loadStateAtSnapshot loads the state of a list of event type and state key pairs in a room at a snapshot.
-// This is typically the state before an event or the current state of a room.
-// Returns a sorted list of state entries or an error if there was a problem talking to the database.
-func loadStateAtSnapshotForTuples(db RoomserverQueryAPIDatabase, stateNID types.StateSnapshotNID, stateKeyTuples []types.StateKeyTuple) ([]types.StateEntry, error) {
-	stateBlockNIDLists, err := db.StateBlockNIDs([]types.StateSnapshotNID{stateNID})
-	if err != nil {
-		return nil, err
-	}
-	stateBlockNIDList := stateBlockNIDLists[0]
-
-	stateEntryLists, err := db.StateEntriesForTuples(stateBlockNIDList.StateBlockNIDs, stateKeyTuples)
-	if err != nil {
-		return nil, err
-	}
-	stateEntriesMap := stateEntryListMap(stateEntryLists)
-
-	// Combined all the state entries for this snapshot.
-	// The order of state block NIDs in the list tells us the order to combine them in.
-	var fullState []types.StateEntry
-	for _, stateBlockNID := range stateBlockNIDList.StateBlockNIDs {
-		entries, ok := stateEntriesMap.lookup(stateBlockNID)
-		if !ok {
-			// If the block is missing from the map it means that none of its entries matched a requested tuple.
-			continue
-		}
-		fullState = append(fullState, entries...)
-	}
-
-	// Stable sort so that the most recent entry for each state key stays
-	// remains later in the list than the older entries for the same state key.
-	sort.Stable(stateEntryByStateKeySorter(fullState))
-	// Unique returns the last entry and hence the most recent entry for each state key.
-	fullState = fullState[:util.Unique(stateEntryByStateKeySorter(fullState))]
-	return fullState, nil
-}
-
-type stateEntryListMap []types.StateEntryList
-
-func (m stateEntryListMap) lookup(stateBlockNID types.StateBlockNID) (stateEntries []types.StateEntry, ok bool) {
-	list := []types.StateEntryList(m)
-	i := sort.Search(len(list), func(i int) bool {
-		return list[i].StateBlockNID >= stateBlockNID
-	})
-	if i < len(list) && list[i].StateBlockNID == stateBlockNID {
-		ok = true
-		stateEntries = list[i].StateEntries
-	}
-	return
-}
-
-type stateEntryByStateKeySorter []types.StateEntry
-
-func (s stateEntryByStateKeySorter) Len() int { return len(s) }
-func (s stateEntryByStateKeySorter) Less(i, j int) bool {
-	return s[i].StateKeyTuple.LessThan(s[j].StateKeyTuple)
-}
-func (s stateEntryByStateKeySorter) Swap(i, j int) { s[i], s[j] = s[j], s[i] }

src/github.com/matrix-org/dendrite/roomserver/state/state.go

@@ -4,6 +4,7 @@ package state
 
 import (
 	"fmt"
+	"github.com/matrix-org/dendrite/roomserver/api"
 	"github.com/matrix-org/dendrite/roomserver/types"
 	"github.com/matrix-org/util"
 	"sort"
@@ -11,12 +12,24 @@ import (
 
 // A RoomStateDatabase has the storage APIs needed to load state from the database
 type RoomStateDatabase interface {
+	// Lookup the numeric IDs for a list of string event types.
+	// Returns a map from string event type to numeric ID for the event type.
+	EventTypeNIDs(eventTypes []string) (map[string]types.EventTypeNID, error)
+	// Lookup the numeric IDs for a list of string event state keys.
+	// Returns a map from string state key to numeric ID for the state key.
+	EventStateKeyNIDs(eventStateKeys []string) (map[string]types.EventStateKeyNID, error)
 	// Lookup the numeric state data IDs for each numeric state snapshot ID
 	// The returned slice is sorted by numeric state snapshot ID.
 	StateBlockNIDs(stateNIDs []types.StateSnapshotNID) ([]types.StateBlockNIDList, error)
 	// Lookup the state data for each numeric state data ID
 	// The returned slice is sorted by numeric state data ID.
 	StateEntries(stateBlockNIDs []types.StateBlockNID) ([]types.StateEntryList, error)
+	// Lookup the state data for the state key tuples for each numeric state block ID
+	// This is used to fetch a subset of the room state at a snapshot.
+	// The returned slice is sorted by numeric state block ID.
+	StateEntriesForTuples(stateBlockNIDs []types.StateBlockNID, stateKeyTuples []types.StateKeyTuple) (
+		[]types.StateEntryList, error,
+	)
 }
 
 // LoadStateAtSnapshot loads the full state of a room at a particular snapshot.
@@ -182,6 +195,79 @@ func DifferenceBetweeenStateSnapshots(db RoomStateDatabase, oldStateNID, newStat
 	}
 }
 
+// StringTuplesToNumericTuples converts the string state key tuples into numeric IDs
+// If there isn't a numeric ID for either the event type or the event state key then the tuple is discarded.
+// Returns an error if there was a problem talking to the database.
+func StringTuplesToNumericTuples(db RoomStateDatabase, stringTuples []api.StateKeyTuple) ([]types.StateKeyTuple, error) {
+	eventTypes := make([]string, len(stringTuples))
+	stateKeys := make([]string, len(stringTuples))
+	for i := range stringTuples {
+		eventTypes[i] = stringTuples[i].EventType
+		stateKeys[i] = stringTuples[i].EventStateKey
+	}
+	eventTypes = util.UniqueStrings(eventTypes)
+	eventTypeMap, err := db.EventTypeNIDs(eventTypes)
+	if err != nil {
+		return nil, err
+	}
+	stateKeys = util.UniqueStrings(stateKeys)
+	stateKeyMap, err := db.EventStateKeyNIDs(stateKeys)
+	if err != nil {
+		return nil, err
+	}
+
+	var result []types.StateKeyTuple
+	for _, stringTuple := range stringTuples {
+		var numericTuple types.StateKeyTuple
+		var ok1, ok2 bool
+		numericTuple.EventTypeNID, ok1 = eventTypeMap[stringTuple.EventType]
+		numericTuple.EventStateKeyNID, ok2 = stateKeyMap[stringTuple.EventStateKey]
+		// Discard the tuple if there wasn't a numeric ID for either the event type or the state key.
+		if ok1 && ok2 {
+			result = append(result, numericTuple)
+		}
+	}
+
+	return result, nil
+}
+
+// LoadStateAtSnapshotForTuples loads the state for a list of event type and state key pairs at a snapshot.
+// This is used when we only want to load a subset of the room state at a snapshot.
+// This is typically the state before an event or the current state of a room.
+// Returns a sorted list of state entries or an error if there was a problem talking to the database.
+func LoadStateAtSnapshotForTuples(db RoomStateDatabase, stateNID types.StateSnapshotNID, stateKeyTuples []types.StateKeyTuple) ([]types.StateEntry, error) {
+	stateBlockNIDLists, err := db.StateBlockNIDs([]types.StateSnapshotNID{stateNID})
+	if err != nil {
+		return nil, err
+	}
+	stateBlockNIDList := stateBlockNIDLists[0]
+
+	stateEntryLists, err := db.StateEntriesForTuples(stateBlockNIDList.StateBlockNIDs, stateKeyTuples)
+	if err != nil {
+		return nil, err
+	}
+	stateEntriesMap := stateEntryListMap(stateEntryLists)
+
+	// Combined all the state entries for this snapshot.
+	// The order of state block NIDs in the list tells us the order to combine them in.
+	var fullState []types.StateEntry
+	for _, stateBlockNID := range stateBlockNIDList.StateBlockNIDs {
+		entries, ok := stateEntriesMap.lookup(stateBlockNID)
+		if !ok {
+			// If the block is missing from the map it means that none of its entries matched a requested tuple.
+			continue
+		}
+		fullState = append(fullState, entries...)
+	}
+
+	// Stable sort so that the most recent entry for each state key stays
+	// remains later in the list than the older entries for the same state key.
+	sort.Stable(stateEntryByStateKeySorter(fullState))
+	// Unique returns the last entry and hence the most recent entry for each state key.
+	fullState = fullState[:util.Unique(stateEntryByStateKeySorter(fullState))]
+	return fullState, nil
+}
+
 type stateBlockNIDListMap []types.StateBlockNIDList
 
 func (m stateBlockNIDListMap) lookup(stateNID types.StateSnapshotNID) (stateBlockNIDs []types.StateBlockNID, ok bool) {

src/github.com/matrix-org/dendrite/roomserver/storage/storage.go

@@ -145,12 +145,12 @@ func (d *Database) StateEntriesForEventIDs(eventIDs []string) ([]types.StateEntr
 	return d.statements.bulkSelectStateEventByID(eventIDs)
 }
 
-// EventTypeNIDs implements query.RoomserverQueryAPIDatabase
+// EventTypeNIDs implements state.RoomStateDatabase
 func (d *Database) EventTypeNIDs(eventTypes []string) (map[string]types.EventTypeNID, error) {
 	return d.statements.bulkSelectEventTypeNID(eventTypes)
 }
 
-// EventStateKeyNIDs implements input.EventDatabase and query.RoomserverQueryAPIDatabase
+// EventStateKeyNIDs implements state.RoomStateDatabase
 func (d *Database) EventStateKeyNIDs(eventStateKeys []string) (map[string]types.EventStateKeyNID, error) {
 	return d.statements.bulkSelectEventStateKeyNID(eventStateKeys)
 }
@@ -200,12 +200,12 @@ func (d *Database) StateAtEventIDs(eventIDs []string) ([]types.StateAtEvent, err
 	return d.statements.bulkSelectStateAtEventByID(eventIDs)
 }
 
-// StateBlockNIDs implements input.EventDatabase and query.RoomserverQueryAPIDatabase
+// StateBlockNIDs implements state.RoomStateDatabase
 func (d *Database) StateBlockNIDs(stateNIDs []types.StateSnapshotNID) ([]types.StateBlockNIDList, error) {
 	return d.statements.bulkSelectStateBlockNIDs(stateNIDs)
 }
 
-// StateEntries implements input.EventDatabase
+// StateEntries implements state.RoomStateDatabase
 func (d *Database) StateEntries(stateBlockNIDs []types.StateBlockNID) ([]types.StateEntryList, error) {
 	return d.statements.bulkSelectStateBlockEntries(stateBlockNIDs)
 }
@@ -341,7 +341,7 @@ func (d *Database) LatestEventIDs(roomNID types.RoomNID) ([]gomatrixserverlib.Ev
 	return references, currentStateSnapshotNID, nil
 }
 
-// StateEntriesForTuples implements query.RoomserverQueryAPIDB
+// StateEntriesForTuples implements state.RoomStateDatabase
 func (d *Database) StateEntriesForTuples(
 	stateBlockNIDs []types.StateBlockNID, stateKeyTuples []types.StateKeyTuple,
 ) ([]types.StateEntryList, error) {