pull out stash changes from entrykit

packages/stash/package.json

@@ -39,13 +39,12 @@
     "test:ci": "pnpm run test"
   },
   "dependencies": {
-    "@ark/util": "catalog:",
+    "@ark/util": "0.2.2",
     "@latticexyz/common": "workspace:*",
     "@latticexyz/config": "workspace:*",
     "@latticexyz/protocol-parser": "workspace:*",
     "@latticexyz/schema-type": "workspace:*",
-    "@latticexyz/store": "workspace:*",
-    "viem": "catalog:"
+    "@latticexyz/store": "workspace:*"
   },
   "devDependencies": {
     "@testing-library/react": "^16.0.0",
@@ -54,12 +53,14 @@
     "eslint-plugin-react": "7.31.11",
     "eslint-plugin-react-hooks": "4.6.0",
     "fast-deep-equal": "^3.1.3",
-    "react": "^18.2.0",
-    "react-dom": "^18.2.0",
-    "tsup": "^6.7.0"
+    "react": "18.2.0",
+    "react-dom": "18.2.0",
+    "tsup": "^6.7.0",
+    "viem": "2.21.19"
   },
   "peerDependencies": {
-    "react": "18.x"
+    "react": "18.x",
+    "viem": "2.x"
   },
   "publishConfig": {
     "access": "public"

packages/stash/src/actions/applyUpdates.ts

@@ -0,0 +1,82 @@
+import { schemaAbiTypeToDefaultValue } from "@latticexyz/schema-type/internal";
+import { Key, Stash, StoreUpdates, TableRecord } from "../common";
+import { encodeKey } from "./encodeKey";
+import { Table } from "@latticexyz/config";
+import { registerTable } from "./registerTable";
+
+export type StashUpdate<table extends Table = Table> = {
+  table: table;
+  key: Key<table>;
+  value: undefined | Partial<TableRecord<table>>;
+};
+
+export type ApplyUpdatesArgs = {
+  stash: Stash;
+  updates: StashUpdate[];
+};
+
+const pendingUpdates = new Map<Stash, StoreUpdates>();
+
+export function applyUpdates({ stash, updates }: ApplyUpdatesArgs): void {
+  const storeUpdates = pendingUpdates.get(stash) ?? { config: {}, records: {} };
+  if (!pendingUpdates.has(stash)) pendingUpdates.set(stash, storeUpdates);
+
+  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 prevUpdate = storeUpdates.records[table.namespaceLabel]?.[table.label]?.[encodedKey];
+    const update = {
+      // preserve the initial prev state if we already have a pending update
+      // TODO: change subscribers to an array of updates instead of an object
+      prev: prevUpdate ? prevUpdate.prev : prevRecord,
+      current: nextRecord,
+    };
+    ((storeUpdates.records[table.namespaceLabel] ??= {})[table.label] ??= {})[encodedKey] ??= update;
+  }
+
+  // queueMicrotask(() => {
+  notifySubscribers(stash);
+  // });
+}
+
+function notifySubscribers(stash: Stash) {
+  const storeUpdates = pendingUpdates.get(stash);
+  if (!storeUpdates) return;
+
+  // Notify table subscribers
+  for (const [namespaceLabel, tableUpdates] of Object.entries(storeUpdates.records)) {
+    for (const [label, updates] of Object.entries(tableUpdates)) {
+      stash._.tableSubscribers[namespaceLabel]?.[label]?.forEach((subscriber) => subscriber(updates));
+    }
+  }
+  // Notify stash subscribers
+  stash._.storeSubscribers.forEach((subscriber) => subscriber(storeUpdates));
+
+  pendingUpdates.delete(stash);
+}

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

@@ -45,17 +45,7 @@ describe("deleteRecord", () => {
       key: { field2: 1, field3: 2 },
     });
 
-    attest(stash.get().records).snap({
-      namespace1: {
-        table1: {
-          "3|1": {
-            field1: "world",
-            field2: 3,
-            field3: 1,
-          },
-        },
-      },
-    });
+    attest(stash.get().records).snap({ namespace1: { table1: { "3|1": { field1: "world", field2: 3, field3: 1 } } } });
   });
 
   it("should throw a type error if an invalid key is provided", () => {

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/runQuery.test.ts

@@ -70,12 +70,7 @@ describe("runQuery", () => {
 
   it("should return all keys that are in the Position and Health table", () => {
     const result = runQuery({ stash, query: [In(Position), In(Health)] });
-    attest(result).snap({
-      keys: {
-        "0x3": { player: "0x3" },
-        "0x4": { player: "0x4" },
-      },
-    });
+    attest(result).snap({ keys: { "0x3": { player: "0x3" }, "0x4": { player: "0x4" } } });
   });
 
   it("should return all keys that have Position.x = 4 and are included in Health", () => {
@@ -85,13 +80,7 @@ describe("runQuery", () => {
 
   it("should return all keys that are in Position but not Health", () => {
     const result = runQuery({ stash, query: [In(Position), Not(In(Health))] });
-    attest(result).snap({
-      keys: {
-        "0x0": { player: "0x0" },
-        "0x1": { player: "0x1" },
-        "0x2": { player: "0x2" },
-      },
-    });
+    attest(result).snap({ keys: { "0x0": { player: "0x0" }, "0x1": { player: "0x1" }, "0x2": { player: "0x2" } } });
   });
 
   it("should return all keys that don't include a gold item in the Inventory table", () => {
@@ -116,23 +105,10 @@ describe("runQuery", () => {
   it("should include all matching records from the tables if includeRecords is set", () => {
     const result = runQuery({ stash, query: [In(Position), In(Health)], options: { includeRecords: true } });
     attest(result).snap({
-      keys: {
-        "0x3": { player: "0x3" },
-        "0x4": { player: "0x4" },
-      },
+      keys: { "0x3": { player: "0x3" }, "0x4": { player: "0x4" } },
       records: {
-        namespace1: {
-          Position: {
-            "0x3": { player: "0x3", x: 3, y: 2 },
-            "0x4": { player: "0x4", x: 4, y: 1 },
-          },
-        },
-        namespace2: {
-          Health: {
-            "0x3": { player: "0x3", health: 3 },
-            "0x4": { player: "0x4", health: 4 },
-          },
-        },
+        namespace1: { Position: { "0x3": { player: "0x3", x: 3, y: 2 }, "0x4": { player: "0x4", x: 4, y: 1 } } },
+        namespace2: { Health: { "0x3": { player: "0x3", health: 3 }, "0x4": { player: "0x4", health: 4 } } },
       },
     });
   });

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

@@ -39,10 +39,7 @@ describe("setRecord", () => {
 
     attest(stash.get().records).snap({
       namespace1: {
-        table1: {
-          "1|2": { field1: "hello", field2: 1, field3: 2 },
-          "2|1": { field1: "world", field2: 2, field3: 1 },
-        },
+        table1: { "1|2": { field1: "hello", field2: 1, field3: 2 }, "2|1": { field1: "world", field2: 2, field3: 1 } },
       },
     });
   });

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),
+    })),
+  });
 }