From a735e14b44f7bd0ed72745610d49b55a181f5401 Mon Sep 17 00:00:00 2001
From: dudendy <a.dude.koval@gmail.com>
Date: Thu, 18 Jan 2024 17:14:24 +0700
Subject: [PATCH] fix(store-sync): improve syncToZustand hydration speed
 (#2145)

Co-authored-by: Kevin Ingersoll <kingersoll@gmail.com>
---
 .changeset/olive-bugs-add.md                  |  5 ++
 .../src/zustand/createStorageAdapter.ts       | 66 ++++++++-----------
 2 files changed, 31 insertions(+), 40 deletions(-)
 create mode 100644 .changeset/olive-bugs-add.md

diff --git a/.changeset/olive-bugs-add.md b/.changeset/olive-bugs-add.md
new file mode 100644
index 0000000000..93164051f9
--- /dev/null
+++ b/.changeset/olive-bugs-add.md
@@ -0,0 +1,5 @@
+---
+"@latticexyz/store-sync": patch
+---
+
+Improved `syncToZustand` speed of hydrating from snapshot by only applying block logs once per block instead of once per log.
diff --git a/packages/store-sync/src/zustand/createStorageAdapter.ts b/packages/store-sync/src/zustand/createStorageAdapter.ts
index a44025b94d..66ed324eba 100644
--- a/packages/store-sync/src/zustand/createStorageAdapter.ts
+++ b/packages/store-sync/src/zustand/createStorageAdapter.ts
@@ -23,6 +23,8 @@ export function createStorageAdapter<tables extends Tables>({
     const updatedIds: string[] = [];
     const deletedIds: string[] = [];
 
+    const rawRecords = { ...store.getState().rawRecords };
+
     for (const log of logs) {
       const table = store.getState().tables[log.args.tableId];
       if (!table) {
@@ -41,20 +43,15 @@ export function createStorageAdapter<tables extends Tables>({
           id,
           log,
         });
+        rawRecords[id] = {
+          id,
+          tableId: log.args.tableId,
+          keyTuple: log.args.keyTuple,
+          staticData: log.args.staticData,
+          encodedLengths: log.args.encodedLengths,
+          dynamicData: log.args.dynamicData,
+        };
         updatedIds.push(id);
-        store.setState({
-          rawRecords: {
-            ...store.getState().rawRecords,
-            [id]: {
-              id,
-              tableId: log.args.tableId,
-              keyTuple: log.args.keyTuple,
-              staticData: log.args.staticData,
-              encodedLengths: log.args.encodedLengths,
-              dynamicData: log.args.dynamicData,
-            },
-          },
-        });
       } else if (log.eventName === "Store_SpliceStaticData") {
         debug("splicing static data", {
           namespace: table.namespace,
@@ -62,8 +59,7 @@ export function createStorageAdapter<tables extends Tables>({
           id,
           log,
         });
-        updatedIds.push(id);
-        const previousRecord = (store.getState().rawRecords[id] as RawRecord | undefined) ?? {
+        const previousRecord = (rawRecords[id] as RawRecord | undefined) ?? {
           id,
           tableId: log.args.tableId,
           keyTuple: log.args.keyTuple,
@@ -72,15 +68,11 @@ export function createStorageAdapter<tables extends Tables>({
           dynamicData: "0x",
         };
         const staticData = spliceHex(previousRecord.staticData, log.args.start, size(log.args.data), log.args.data);
-        store.setState({
-          rawRecords: {
-            ...store.getState().rawRecords,
-            [id]: {
-              ...previousRecord,
-              staticData,
-            },
-          },
-        });
+        rawRecords[id] = {
+          ...previousRecord,
+          staticData,
+        };
+        updatedIds.push(id);
       } else if (log.eventName === "Store_SpliceDynamicData") {
         debug("splicing dynamic data", {
           namespace: table.namespace,
@@ -88,8 +80,7 @@ export function createStorageAdapter<tables extends Tables>({
           id,
           log,
         });
-        updatedIds.push(id);
-        const previousRecord = (store.getState().rawRecords[id] as RawRecord | undefined) ?? {
+        const previousRecord = (rawRecords[id] as RawRecord | undefined) ?? {
           id,
           tableId: log.args.tableId,
           keyTuple: log.args.keyTuple,
@@ -99,16 +90,12 @@ export function createStorageAdapter<tables extends Tables>({
         };
         const encodedLengths = log.args.encodedLengths;
         const dynamicData = spliceHex(previousRecord.dynamicData, log.args.start, log.args.deleteCount, log.args.data);
-        store.setState({
-          rawRecords: {
-            ...store.getState().rawRecords,
-            [id]: {
-              ...previousRecord,
-              encodedLengths,
-              dynamicData,
-            },
-          },
-        });
+        rawRecords[id] = {
+          ...previousRecord,
+          encodedLengths,
+          dynamicData,
+        };
+        updatedIds.push(id);
       } else if (log.eventName === "Store_DeleteRecord") {
         debug("deleting record", {
           namespace: table.namespace,
@@ -116,9 +103,8 @@ export function createStorageAdapter<tables extends Tables>({
           id,
           log,
         });
+        delete rawRecords[id];
         deletedIds.push(id);
-        const { [id]: deletedRecord, ...rawRecords } = store.getState().rawRecords;
-        store.setState({ rawRecords });
       }
     }
 
@@ -129,7 +115,7 @@ export function createStorageAdapter<tables extends Tables>({
       ...Object.fromEntries(
         updatedIds
           .map((id) => {
-            const rawRecord = store.getState().rawRecords[id];
+            const rawRecord = rawRecords[id] as RawRecord | undefined;
             if (!rawRecord) {
               console.warn("no raw record found for updated ID", id);
               return;
@@ -155,6 +141,6 @@ export function createStorageAdapter<tables extends Tables>({
       ),
     };
 
-    store.setState({ records });
+    store.setState({ rawRecords, records });
   };
 }