refactor: create State type to replace Map

packages/federation-sdk/src/services/state.service.spec.ts

@@ -1,5 +1,6 @@
 import { beforeEach, describe, expect, it, spyOn, test } from 'bun:test';
 import { type EventStore } from '@rocket.chat/federation-core';
+import type { State } from '@rocket.chat/federation-room';
 import * as room from '@rocket.chat/federation-room';
 import {
 	EventID,
@@ -11,7 +12,6 @@ import {
 	PersistentEventFactory,
 	RejectCodes,
 	RoomVersion,
-	StateMapKey,
 } from '@rocket.chat/federation-room';
 import { type WithId } from 'mongodb';
 import { EventRepository } from '../repositories/event.repository';
@@ -23,8 +23,6 @@ import { type ConfigService } from './config.service';
 import { DatabaseConnectionService } from './database-connection.service';
 import { StateService } from './state.service';
 
-type State = Map<StateMapKey, PersistentEventBase>;
-
 function getDefaultFields() {
 	return {
 		auth_events: [],

packages/federation-sdk/src/services/state.service.ts

@@ -14,6 +14,7 @@ import {
 	RoomID,
 	RoomState,
 	RoomVersion,
+	State,
 	type StateID,
 	type StateMapKey,
 	StateResolverAuthorizationError,
@@ -26,8 +27,6 @@ import { EventRepository } from '../repositories/event.repository';
 import { StateGraphRepository } from '../repositories/state-graph.repository';
 import { ConfigService } from './config.service';
 
-type State = Map<StateMapKey, PersistentEventBase>;
-
 type StrippedEvent = {
 	content: PduContent;
 	sender: string;
@@ -354,7 +353,7 @@ export class StateService {
 	}
 
 	private buildStateFromEvents(events: PersistentEventBase[]): State {
-		const state = new Map<StateMapKey, PersistentEventBase>();
+		const state: State = new Map();
 		for (const event of events) {
 			state.set(event.getUniqueStateIdentifier(), event);
 		}
@@ -372,7 +371,7 @@ export class StateService {
 		const eventIds = stateMap.values().toArray();
 
 		const events = await this.eventRepository.findByIds(eventIds).toArray();
-		const state = new Map<StateMapKey, PersistentEventBase>();
+		const state: State = new Map();
 		for (const event of events) {
 			const e = PersistentEventFactory.createFromRawEvent(
 				event.event,
@@ -836,7 +835,7 @@ export class StateService {
 	private async _resolveState(
 		stateIds: StateID[],
 		roomVersion: RoomVersion,
-	): Promise<Map<StateMapKey, PersistentEventBase>> {
+	): Promise<State> {
 		if (await this._isSameChain(stateIds)) {
 			// pick the latest id from the chain, that's the state to use
 			const latestDelta =

packages/room/src/manager/event-wrapper.ts

@@ -44,6 +44,16 @@ export type PduWithHashesAndSignaturesOptional<T extends Pdu = Pdu> = Prettify<
 
 export const REDACT_ALLOW_ALL_KEYS: unique symbol = Symbol.for('all');
 
+export interface State extends Map<StateMapKey, PersistentEventBase> {
+	get<T extends StateMapKey>(
+		key: T,
+	): T extends `${infer I}:${string}`
+		? I extends PduType
+			? PersistentEventBase<RoomVersion, I> | undefined
+			: never
+		: never;
+}
+
 // convinient wrapper to manage schema differences when working with same algorithms across different versions
 export abstract class PersistentEventBase<
 	Version extends RoomVersion = RoomVersion,

packages/room/src/state_resolution/definitions/algorithm/v2.ts

@@ -58,7 +58,7 @@
 // i am too early in this to remember everything by heart.
 
 import assert from 'node:assert';
-import { PersistentEventBase } from '../../../manager/event-wrapper';
+import { PersistentEventBase, State } from '../../../manager/event-wrapper';
 import { PowerLevelEvent } from '../../../manager/power-level-event-wrapper';
 import type { EventID, StateMapKey } from '../../../types/_common';
 import {
@@ -83,7 +83,7 @@ export const isTruthy = <T>(
 export async function resolveStateV2Plus(
 	states: ReadonlyMap<StateMapKey, PersistentEventBase>[],
 	store: EventStore, // with cache
-): Promise<Map<StateMapKey, PersistentEventBase>> {
+): Promise<State> {
 	// memory o'memory
 	// const eventMap = new Map<string, PduV3>();
 
@@ -138,7 +138,7 @@ export async function resolveStateV2Plus(
 			assert(event, 'event should not be null');
 			accum.set(stateKey, event);
 			return accum;
-		}, new Map<StateMapKey, PersistentEventBase>());
+		}, new Map() as State);
 
 	if (conflicted.size === 0) {
 		// no conflicted state, return the unconflicted state

packages/room/src/state_resolution/definitions/definitions.ts

@@ -1,11 +1,15 @@
 import { PriorityQueue } from '@datastructures-js/priority-queue';
-import type { EventID, State, StateMapKey } from '../../types/_common';
+import type {
+	EventID,
+	StateEventIdMap,
+	StateMapKey,
+} from '../../types/_common';
 import { type PduType } from '../../types/v3-11';
 
 import assert from 'node:assert';
 import { StateResolverAuthorizationError } from '../../authorizartion-rules/errors';
 import { checkEventAuthWithState } from '../../authorizartion-rules/rules';
-import { PersistentEventBase } from '../../manager/event-wrapper';
+import { PersistentEventBase, State } from '../../manager/event-wrapper';
 import { PowerLevelEvent } from '../../manager/power-level-event-wrapper';
 import { RoomVersion } from '../../manager/type';
 
@@ -47,8 +51,8 @@ export function isPowerEvent(event: PersistentEventBase): boolean {
 // iout map takes care of the non-duplication of a set
 export function partitionState(
 	events: Readonly<Iterable<PersistentEventBase>>,
-): [State, Map<StateMapKey, EventID[]>] {
-	const unconflictedState: State = new Map();
+): [StateEventIdMap, Map<StateMapKey, EventID[]>] {
+	const unconflictedState: StateEventIdMap = new Map();
 
 	// Note that the unconflicted state map only has one event for each key K, whereas the conflicted state set may contain multiple events with the same key.
 	const conflictedStateEventsMap: Map<StateMapKey, EventID[]> = new Map();
@@ -148,7 +152,7 @@ export async function getAuthChain(
 // Auth difference
 // NOTE: https://github.com/element-hq/synapse/blob/a25a37002c851ef419d12925a11dd8bf2233470e/docs/auth_chain_difference_algorithm.md
 export async function getAuthChainDifference(
-	states: Readonly<Iterable<State>>,
+	states: Readonly<Iterable<StateEventIdMap>>,
 	store: EventStore,
 ) {
 	const authChainSets = [] as Set<EventID>[];
@@ -574,13 +578,11 @@ export async function iterativeAuthChecks(
 	events: PersistentEventBase[],
 	stateMap: ReadonlyMap<StateMapKey, PersistentEventBase>,
 	store: EventStore,
-): Promise<Map<StateMapKey, PersistentEventBase>> {
-	const newState = new Map<StateMapKey, PersistentEventBase>(
-		stateMap.entries(),
-	);
+): Promise<State> {
+	const newState: State = new Map(stateMap.entries()) as State;
 
 	for (const event of events) {
-		const authEventStateMap = new Map<StateMapKey, PersistentEventBase>();
+		const authEventStateMap: State = new Map();
 		for (const authEvent of await store.getEvents(event.getAuthEventIds())) {
 			authEventStateMap.set(authEvent.getUniqueStateIdentifier(), authEvent);
 		}

packages/room/src/types/_common.ts

@@ -21,7 +21,7 @@ export type UserID = z.infer<typeof userIdSchema>;
 
 export type StateMapKey = `${PduType}:${StateKey}`;
 
-export type State = Map<StateMapKey, EventID>;
+export type StateEventIdMap = Map<StateMapKey, EventID>;
 
 export type PduForType<P extends PduType = PduType> = Extract<Pdu, { type: P }>;