feat(stash): various improvements

packages/stash/src/actions/applyUpdates.ts

@@ -0,0 +1,91 @@
+import { schemaAbiTypeToDefaultValue } from "@latticexyz/schema-type/internal";
+import { Key, Stash, TableRecord, TableUpdates } from "../common";
+import { encodeKey } from "./encodeKey";
+import { Table } from "@latticexyz/config";
+import { registerTable } from "./registerTable";
+
+export type PendingStashUpdate<table extends Table = Table> = {
+  table: table;
+  key: Key<table>;
+  value: undefined | Partial<TableRecord<table>>;
+};
+
+export type ApplyUpdatesArgs = {
+  stash: Stash;
+  updates: PendingStashUpdate[];
+};
+
+type PendingUpdates = {
+  [namespaceLabel: string]: {
+    [tableLabel: string]: TableUpdates;
+  };
+};
+
+const pendingStashUpdates = new Map<Stash, PendingUpdates>();
+
+export function applyUpdates({ stash, updates }: ApplyUpdatesArgs): void {
+  const pendingUpdates = pendingStashUpdates.get(stash) ?? {};
+  if (!pendingStashUpdates.has(stash)) pendingStashUpdates.set(stash, pendingUpdates);
+
+  for (const { table, key, value } of updates) {
+    if (stash.get().config[table.namespaceLabel]?.[table.label] == null) {
+      registerTable({ stash, table });
+    }
+    const tableState = ((stash._.state.records[table.namespaceLabel] ??= {})[table.label] ??= {});
+    const encodedKey = encodeKey({ table, key });
+    const prevRecord = tableState[encodedKey];
+
+    // create new record, preserving field order
+    const nextRecord =
+      value == null
+        ? undefined
+        : Object.fromEntries(
+            Object.entries(table.schema).map(([fieldName, { type }]) => [
+              fieldName,
+              key[fieldName] ?? // Use provided key fields
+                value[fieldName] ?? // Or provided value fields
+                prevRecord?.[fieldName] ?? // Keep existing non-overridden fields
+                schemaAbiTypeToDefaultValue[type], // Default values for new fields
+            ]),
+          );
+
+    // apply update to state
+    if (nextRecord != null) {
+      tableState[encodedKey] = nextRecord;
+    } else {
+      delete tableState[encodedKey];
+    }
+
+    // add update to pending updates for notifying subscribers
+    const tableUpdates = ((pendingUpdates[table.namespaceLabel] ??= {})[table.label] ??= []);
+    tableUpdates.push({
+      table,
+      key,
+      previous: prevRecord,
+      current: nextRecord,
+    });
+  }
+
+  queueMicrotask(() => {
+    notifySubscribers(stash);
+  });
+}
+
+function notifySubscribers(stash: Stash) {
+  const pendingUpdates = pendingStashUpdates.get(stash);
+  if (!pendingUpdates) return;
+
+  // Notify table subscribers
+  for (const [namespaceLabel, namespaceUpdates] of Object.entries(pendingUpdates)) {
+    for (const [tableLabel, tableUpdates] of Object.entries(namespaceUpdates)) {
+      stash._.tableSubscribers[namespaceLabel]?.[tableLabel]?.forEach((subscriber) => subscriber(tableUpdates));
+    }
+  }
+  // Notify stash subscribers
+  const updates = Object.values(pendingUpdates)
+    .map((namespaceUpdates) => Object.values(namespaceUpdates))
+    .flat(2);
+  stash._.storeSubscribers.forEach((subscriber) => subscriber({ type: "records", updates }));
+
+  pendingStashUpdates.delete(stash);
+}

packages/stash/src/actions/deleteRecord.ts

@@ -1,7 +1,6 @@
 import { Table } from "@latticexyz/config";
 import { Key, Stash } from "../common";
-import { encodeKey } from "./encodeKey";
-import { registerTable } from "./registerTable";
+import { applyUpdates } from "./applyUpdates";
 
 export type DeleteRecordArgs<table extends Table = Table> = {
   stash: Stash;
@@ -12,26 +11,5 @@ export type DeleteRecordArgs<table extends Table = Table> = {
 export type DeleteRecordResult = void;
 
 export function deleteRecord<table extends Table>({ stash, table, key }: DeleteRecordArgs<table>): DeleteRecordResult {
-  const { namespaceLabel, label } = table;
-
-  if (stash.get().config[namespaceLabel] == null) {
-    registerTable({ stash, table });
-  }
-
-  const encodedKey = encodeKey({ table, key });
-  const prevRecord = stash.get().records[namespaceLabel]?.[label]?.[encodedKey];
-
-  // Early return if this record doesn't exist
-  if (prevRecord == null) return;
-
-  // Delete record
-  delete stash._.state.records[namespaceLabel]?.[label]?.[encodedKey];
-
-  // Notify table subscribers
-  const updates = { [encodedKey]: { prev: prevRecord && { ...prevRecord }, current: undefined } };
-  stash._.tableSubscribers[namespaceLabel]?.[label]?.forEach((subscriber) => subscriber(updates));
-
-  // Notify stash subscribers
-  const storeUpdate = { config: {}, records: { [namespaceLabel]: { [label]: updates } } };
-  stash._.storeSubscribers.forEach((subscriber) => subscriber(storeUpdate));
+  applyUpdates({ stash, updates: [{ table, key, value: undefined }] });
 }

packages/stash/src/actions/getTable.test.ts

@@ -297,6 +297,8 @@ describe("getTable", () => {
 
   describe("subscribe", () => {
     it("should notify subscriber of table change", () => {
+      vi.useFakeTimers({ toFake: ["queueMicrotask"] });
+
       const config1 = defineTable({
         label: "table1",
         schema: { a: "address", b: "uint256", c: "uint32" },
@@ -315,6 +317,7 @@ describe("getTable", () => {
       table1.subscribe({ subscriber });
 
       table1.setRecord({ key: { a: "0x00" }, value: { b: 1n, c: 2 } });
+      vi.advanceTimersToNextTimer();
 
       expect(subscriber).toHaveBeenCalledTimes(1);
       expect(subscriber).toHaveBeenNthCalledWith(1, {
@@ -326,9 +329,12 @@ describe("getTable", () => {
 
       // Expect unrelated updates to not notify subscribers
       table2.setRecord({ key: { a: "0x01" }, value: { b: 1n, c: 2 } });
+      vi.advanceTimersToNextTimer();
+
       expect(subscriber).toHaveBeenCalledTimes(1);
 
       table1.setRecord({ key: { a: "0x00" }, value: { b: 1n, c: 3 } });
+      vi.advanceTimersToNextTimer();
 
       expect(subscriber).toHaveBeenCalledTimes(2);
       expect(subscriber).toHaveBeenNthCalledWith(2, {

packages/stash/src/actions/getTable.ts

@@ -67,7 +67,7 @@ export function getTable<table extends Table>({ stash, table }: GetTableArgs<tab
     getRecords: (args?: TableBoundGetRecordsArgs<table>) => getRecords({ stash, table, ...args }),
     setRecord: (args: TableBoundSetRecordArgs<table>) => setRecord({ stash, table, ...args }),
     setRecords: (args: TableBoundSetRecordsArgs<table>) => setRecords({ stash, table, ...args }),
-    subscribe: (args: TableBoundSubscribeTableArgs) => subscribeTable({ stash, table, ...args }),
+    subscribe: (args: TableBoundSubscribeTableArgs<table>) => subscribeTable({ stash, table, ...args }),
 
     // TODO: dynamically add setters and getters for individual fields of the table
   };

packages/stash/src/actions/registerTable.ts

@@ -27,11 +27,9 @@ export function registerTable<table extends Table>({
   (stash._.tableSubscribers[namespaceLabel] ??= {})[label] ??= new Set();
 
   // Notify stash subscribers
-  const storeUpdate = {
-    config: { [namespaceLabel]: { [label]: { prev: undefined, current: tableConfig } } },
-    records: {},
-  };
-  stash._.storeSubscribers.forEach((subscriber) => subscriber(storeUpdate));
+  stash._.storeSubscribers.forEach((subscriber) =>
+    subscriber({ type: "config", updates: [{ previous: undefined, current: tableConfig }] }),
+  );
 
   return getTable({ stash, table });
 }

packages/stash/src/actions/setRecord.ts

@@ -1,6 +1,6 @@
-import { Key, TableRecord, Stash } from "../common";
-import { setRecords } from "./setRecords";
 import { Table } from "@latticexyz/config";
+import { Key, TableRecord, Stash } from "../common";
+import { applyUpdates } from "./applyUpdates";
 
 export type SetRecordArgs<table extends Table = Table> = {
   stash: Stash;
@@ -12,12 +12,5 @@ export type SetRecordArgs<table extends Table = Table> = {
 export type SetRecordResult = void;
 
 export function setRecord<table extends Table>({ stash, table, key, value }: SetRecordArgs<table>): SetRecordResult {
-  setRecords({
-    stash,
-    table,
-    records: [
-      // Stored record should include key
-      { ...value, ...key },
-    ],
-  });
+  applyUpdates({ stash, updates: [{ table, key, value }] });
 }

packages/stash/src/actions/setRecords.ts

@@ -1,8 +1,7 @@
-import { dynamicAbiTypeToDefaultValue, staticAbiTypeToDefaultValue } from "@latticexyz/schema-type/internal";
-import { Stash, TableRecord, TableUpdates } from "../common";
-import { encodeKey } from "./encodeKey";
 import { Table } from "@latticexyz/config";
-import { registerTable } from "./registerTable";
+import { getKey, getValue } from "@latticexyz/protocol-parser/internal";
+import { Stash, TableRecord } from "../common";
+import { applyUpdates } from "./applyUpdates";
 
 export type SetRecordsArgs<table extends Table = Table> = {
   stash: Stash;
@@ -13,38 +12,12 @@ export type SetRecordsArgs<table extends Table = Table> = {
 export type SetRecordsResult = void;
 
 export function setRecords<table extends Table>({ stash, table, records }: SetRecordsArgs<table>): SetRecordsResult {
-  const { namespaceLabel, label, schema } = table;
-
-  if (stash.get().config[namespaceLabel]?.[label] == null) {
-    registerTable({ stash, table });
-  }
-
-  // Construct table updates
-  const updates: TableUpdates = {};
-  for (const record of records) {
-    const encodedKey = encodeKey({ table, key: record as never });
-    const prevRecord = stash.get().records[namespaceLabel]?.[label]?.[encodedKey];
-    const newRecord = Object.fromEntries(
-      Object.keys(schema).map((fieldName) => [
-        fieldName,
-        record[fieldName] ?? // Override provided record fields
-          prevRecord?.[fieldName] ?? // Keep existing non-overridden fields
-          staticAbiTypeToDefaultValue[schema[fieldName] as never] ?? // Default values for new fields
-          dynamicAbiTypeToDefaultValue[schema[fieldName] as never],
-      ]),
-    );
-    updates[encodedKey] = { prev: prevRecord, current: newRecord };
-  }
-
-  // Update records
-  for (const [encodedKey, { current }] of Object.entries(updates)) {
-    ((stash._.state.records[namespaceLabel] ??= {})[label] ??= {})[encodedKey] = current as never;
-  }
-
-  // Notify table subscribers
-  stash._.tableSubscribers[namespaceLabel]?.[label]?.forEach((subscriber) => subscriber(updates));
-
-  // Notify stash subscribers
-  const storeUpdate = { config: {}, records: { [namespaceLabel]: { [label]: updates } } };
-  stash._.storeSubscribers.forEach((subscriber) => subscriber(storeUpdate));
+  applyUpdates({
+    stash,
+    updates: Object.values(records).map((record) => ({
+      table,
+      key: getKey(table, record),
+      value: getValue(table, record),
+    })),
+  });
 }

packages/stash/src/actions/subscribeQuery.test.ts

@@ -32,6 +32,7 @@ describe("defineQuery", () => {
 
   beforeEach(() => {
     stash = createStash(config);
+    vi.useFakeTimers({ toFake: ["queueMicrotask"] });
 
     // Add some mock data
     const items = ["0xgold", "0xsilver"] as const;
@@ -45,6 +46,8 @@ describe("defineQuery", () => {
         setRecord({ stash, table: Inventory, key: { player: `0x${String(i)}`, item }, value: { amount: i } });
       }
     }
+
+    vi.advanceTimersToNextTimer();
   });
 
   it("should return the matching keys and keep it updated", () => {
@@ -55,21 +58,26 @@ describe("defineQuery", () => {
     });
 
     setRecord({ stash, table: Health, key: { player: `0x2` }, value: { health: 2 } });
+    vi.advanceTimersToNextTimer();
 
-    attest(result.keys).snap({
-      "0x2": { player: "0x2" },
-      "0x3": { player: "0x3" },
-      "0x4": { player: "0x4" },
-    });
+    attest(result.keys).snap({ "0x3": { player: "0x3" }, "0x4": { player: "0x4" }, "0x2": { player: "0x2" } });
   });
 
   it("should notify subscribers when a matching key is updated", () => {
     let lastUpdate: unknown;
     const subscriber = vi.fn((update: QueryUpdate) => (lastUpdate = update));
-    const result = subscribeQuery({ stash, query: [Matches(Position, { x: 4 }), In(Health)] });
+
+    const result = subscribeQuery({
+      stash,
+      query: [Matches(Position, { x: 4 }), In(Health)],
+    });
     result.subscribe(subscriber);
+    vi.advanceTimersToNextTimer();
+
+    expect(subscriber).toBeCalledTimes(0);
 
     setRecord({ stash, table: Position, key: { player: "0x4" }, value: { y: 2 } });
+    vi.advanceTimersToNextTimer();
 
     expect(subscriber).toBeCalledTimes(1);
     attest(lastUpdate).snap({
@@ -94,7 +102,11 @@ describe("defineQuery", () => {
     const result = subscribeQuery({ stash, query: [In(Position), In(Health)] });
     result.subscribe(subscriber);
 
+    vi.advanceTimersToNextTimer();
+    expect(subscriber).toBeCalledTimes(0);
+
     setRecord({ stash, table: Health, key: { player: `0x2` }, value: { health: 2 } });
+    vi.advanceTimersToNextTimer();
 
     expect(subscriber).toBeCalledTimes(1);
     attest(lastUpdate).snap({
@@ -103,7 +115,7 @@ describe("defineQuery", () => {
           Health: {
             "0x2": {
               prev: undefined,
-              current: { player: `0x2`, health: 2 },
+              current: { player: "0x2", health: 2 },
             },
           },
         },
@@ -119,7 +131,11 @@ describe("defineQuery", () => {
     const result = subscribeQuery({ stash, query: [In(Position), In(Health)] });
     result.subscribe(subscriber);
 
+    vi.advanceTimersToNextTimer();
+    expect(subscriber).toBeCalledTimes(0);
+
     deleteRecord({ stash, table: Position, key: { player: `0x3` } });
+    vi.advanceTimersToNextTimer();
 
     expect(subscriber).toBeCalledTimes(1);
     attest(lastUpdate).snap({