Track items' parents, and only iterate dirty subtrees when creating a snapshot

Author: FTWinston

src/schema/createSnapshot.ts

@@ -44,29 +44,53 @@ export type Snapshot<T> = DeepReadonly<
 export interface SnapshotContext {
     /** Map of refId → Schema object from decoder.root.refs */
     refs: Map<number, any> | undefined;
+    /** Reverse lookup: Schema object → refId (built lazily) */
+    objectToRefId: Map<object, number> | undefined;
     /** Snapshot results from the previous render pass */
     previousResultsByRefId: Map<number, any>;
     /** Snapshot results from the current render pass (for cycle detection) */
     currentResultsByRefId: Map<number, any>;
+    /** Set of refIds that have been modified since the last snapshot */
+    dirtyRefIds: Set<number>;
+    /** Map of childRefId → parentRefId for ancestor tracking */
+    parentRefIdMap: Map<number, number>;
+    /** Current parent refId during traversal (used to build parentRefIdMap) */
+    currentParentRefId: number;
 }
 
 /**
- * Finds the refId for a Schema object by searching the decoder's refs map.
+ * Builds a reverse lookup map from objects to their refIds.
+ */
+function buildObjectToRefIdMap(refs: Map<number, any>): Map<object, number> {
+    const map = new Map<object, number>();
+    for (const [refId, obj] of refs.entries()) {
+        if (obj !== null && typeof obj === "object") {
+            map.set(obj, refId);
+        }
+    }
+    return map;
+}
+
+/**
+ * Finds the refId for a Schema object using the reverse lookup map.
  * 
  * In Colyseus 3.x, each Schema instance is assigned a unique numeric refId
  * that remains stable across encode/decode cycles. This allows us to track
  * object identity even when the JavaScript object references change.
  * 
  * @param node - The Schema object to find the refId for
- * @param refs - The decoder.root.refs map
+ * @param ctx - The snapshot context with the reverse lookup map
  * @returns The refId if found, or -1 if not found
  */
-function findRefId(node: object, refs: Map<number, any> | undefined): number {
-    if (!refs) return -1;
-    for (const [refId, obj] of refs.entries()) {
-        if (obj === node) return refId;
+function findRefId(node: object, ctx: SnapshotContext): number {
+    if (!ctx.refs) return -1;
+    
+    // Build the reverse lookup map lazily on first use.
+    if (!ctx.objectToRefId) {
+        ctx.objectToRefId = buildObjectToRefIdMap(ctx.refs);
     }
-    return -1;
+    
+    return ctx.objectToRefId.get(node) ?? -1;
 }
 
 /**
@@ -109,13 +133,14 @@ function createSnapshotForArraySchema(
     previousResult: any[] | undefined,
     ctx: SnapshotContext
 ): any[] {
-    const items = Array.from(node);
-    const snapshotted: any[] = [];
-    let hasChanged = !previousResult || !Array.isArray(previousResult) || items.length !== previousResult.length;
+    const length = node.length;
+    let hasChanged = !previousResult || !Array.isArray(previousResult) || length !== previousResult.length;
+
+    const snapshotted: any[] = new Array(length);
 
-    for (let i = 0; i < items.length; i++) {
-        const snapshottedValue = createSnapshot(items[i], ctx);
-        snapshotted.push(snapshottedValue);
+    for (let i = 0; i < length; i++) {
+        const snapshottedValue = createSnapshot(node.at(i), ctx);
+        snapshotted[i] = snapshottedValue;
 
         if (!hasChanged && previousResult && previousResult[i] !== snapshottedValue) {
             hasChanged = true;
@@ -187,17 +212,17 @@ function createSnapshotForSchema(
  */
 export function createSnapshot<T>(node: T, ctx: SnapshotContext): Snapshot<T> {
     // Pass through primitives and null/undefined.
-    if (
-        node === null ||
-        node === undefined ||
-        typeof node !== "object" ||
-        typeof node === "function"
-    ) {
+    if (node === null || node === undefined || typeof node !== "object") {
         return node as Snapshot<T>;
     }
 
     // Find the stable refId for this object.
-    const refId = findRefId(node, ctx.refs);
+    const refId = findRefId(node, ctx);
+
+    // Record the parent relationship for ancestor tracking.
+    if (refId !== -1 && ctx.currentParentRefId !== -1) {
+        ctx.parentRefIdMap.set(refId, ctx.currentParentRefId);
+    }
 
     // Check if we've already snapshotted this object in the current pass (cycle detection).
     if (refId !== -1 && ctx.currentResultsByRefId.has(refId)) {
@@ -207,6 +232,18 @@ export function createSnapshot<T>(node: T, ctx: SnapshotContext): Snapshot<T> {
     // Get the previous result for structural sharing comparison.
     const previousResult = refId !== -1 ? ctx.previousResultsByRefId.get(refId) : undefined;
 
+    // If this node is not dirty and we have a previous result,
+    // we can skip the entire subtree. With ancestor tracking, if any descendant
+    // changed, this node would have been marked dirty too.
+    if (refId !== -1 && previousResult !== undefined && !ctx.dirtyRefIds.has(refId)) {
+        ctx.currentResultsByRefId.set(refId, previousResult);
+        return previousResult as Snapshot<T>;
+    }
+
+    // Set this node as the parent for any children we process.
+    const savedParentRefId = ctx.currentParentRefId;
+    ctx.currentParentRefId = refId;
+
     let result: any;
 
     if (node instanceof MapSchema) {
@@ -220,7 +257,8 @@ export function createSnapshot<T>(node: T, ctx: SnapshotContext): Snapshot<T> {
         result = node;
     }
 
-    // Cache the result for this snapshot pass.
+    // Restore parent and cache result.
+    ctx.currentParentRefId = savedParentRefId;
     if (refId !== -1) {
         ctx.currentResultsByRefId.set(refId, result);
     }

src/schema/getOrCreateSubscription.ts

@@ -5,11 +5,19 @@ import { Schema, Decoder, DataChange } from "@colyseus/schema";
  * Subscription state for a single room, shared across all hook instances
  * consuming the same room state.
  */
-interface StateSubscription {
+export interface StateSubscription {
     /** Set of callbacks to invoke when state changes */
     listeners: Set<() => void>;
     /** Cached snapshot results from the previous render, keyed by refId */
     previousResultsByRefId: Map<number, any>;
+    /** Set of refIds that have been modified since the last snapshot */
+    dirtyRefIds: Set<number>;
+    /** Map of childRefId → parentRefId for ancestor tracking */
+    parentRefIdMap: Map<number, number>;
+    /** Reverse lookup: Schema object → refId (rebuilt on each change) */
+    objectToRefId: Map<object, number> | undefined;
+    /** Counter for periodic pruning of stale cache entries */
+    cleanupCounter: number;
     /** Original triggerChanges function from the decoder */
     originalTrigger?: (changes: DataChange[]) => void;
 }
@@ -37,6 +45,10 @@ export function getOrCreateSubscription(roomState: Schema, decoder: Decoder): St
     subscription = {
         listeners: new Set(),
         previousResultsByRefId: new Map(),
+        dirtyRefIds: new Set(),
+        parentRefIdMap: new Map(),
+        objectToRefId: undefined,
+        cleanupCounter: 0,
     };
 
     // Wrap the decoder's triggerChanges to notify React subscribers.
@@ -48,6 +60,30 @@ export function getOrCreateSubscription(roomState: Schema, decoder: Decoder): St
             subscription.originalTrigger(changes);
         }
 
+        // Rebuild reverse lookup since refs may have changed.
+        const refs = decoder.root?.refs;
+        if (refs) {
+            subscription.objectToRefId = new Map();
+            for (const [refId, obj] of refs.entries()) {
+                if (obj !== null && typeof obj === "object") {
+                    subscription.objectToRefId.set(obj, refId);
+                }
+            }
+        }
+
+        // Mark all changed refIds as dirty, walking up the parent chain.
+        for (const change of changes) {
+            const refId = subscription.objectToRefId?.get(change.ref) ?? -1;
+            if (refId !== -1) {
+                // Mark this ref and all its ancestors as dirty.
+                let currentRefId: number | undefined = refId;
+                while (currentRefId !== undefined) {
+                    subscription.dirtyRefIds.add(currentRefId);
+                    currentRefId = subscription.parentRefIdMap.get(currentRefId);
+                }
+            }
+        }
+
         // Notify all React subscribers that state has changed.
         subscription.listeners.forEach((callback) => callback());
     };

src/schema/useColyseusState.ts

@@ -53,8 +53,12 @@ export function useColyseusState<T extends Schema = Schema, U = T>(
         // Create context for this snapshot pass.
         const ctx: SnapshotContext = {
             refs: decoder.root?.refs,
+            objectToRefId: subscription.objectToRefId,
             previousResultsByRefId: subscription.previousResultsByRefId,
             currentResultsByRefId: new Map(),
+            dirtyRefIds: subscription.dirtyRefIds,
+            parentRefIdMap: subscription.parentRefIdMap,
+            currentParentRefId: -1, // No parent for root
         };
 
         const result = createSnapshot(selectedState, ctx);
@@ -64,6 +68,23 @@ export function useColyseusState<T extends Schema = Schema, U = T>(
             subscription.previousResultsByRefId.set(refId, value);
         }
 
+        // Save the objectToRefId map for reuse.
+        subscription.objectToRefId = ctx.objectToRefId;
+
+        // Clear dirty refs after snapshot is complete.
+        subscription.dirtyRefIds.clear();
+
+        // Periodically prune stale cache entries (every 100 snapshots).
+        if (++subscription.cleanupCounter >= 100 && ctx.refs) {
+            subscription.cleanupCounter = 0;
+            for (const refId of subscription.previousResultsByRefId.keys()) {
+                if (!ctx.refs.has(refId)) {
+                    subscription.previousResultsByRefId.delete(refId);
+                    subscription.parentRefIdMap.delete(refId);
+                }
+            }
+        }
+
         return result;
     };