perf: use redux state patches (#26738)

Send state patches to the UI.
Create `PatchStore` for each UI connection.

app/scripts/constants/sentry-state.ts

@@ -279,8 +279,6 @@ export const SENTRY_BACKGROUND_STATE = {
     },
   },
   SnapController: {
-    unencryptedSnapStates: false,
-    snapStates: false,
     snaps: false,
   },
   SnapInterfaceController: {

app/scripts/controllers/preferences-controller.ts

@@ -521,8 +521,11 @@ export default class PreferencesController {
    */
   addKnownMethodData(fourBytePrefix: string, methodData: string): void {
     const { knownMethodData } = this.store.getState();
-    knownMethodData[fourBytePrefix] = methodData;
-    this.store.updateState({ knownMethodData });
+
+    const updatedKnownMethodData = { ...knownMethodData };
+    updatedKnownMethodData[fourBytePrefix] = methodData;
+
+    this.store.updateState({ knownMethodData: updatedKnownMethodData });
   }
 
   /**
@@ -765,11 +768,16 @@ export default class PreferencesController {
     const addresses = Object.values(accounts).map((account) =>
       account.address.toLowerCase(),
     );
-    Object.keys(identities).forEach((identity) => {
-      if (addresses.includes(identity.toLowerCase())) {
-        lostIdentities[identity] = identities[identity];
-      }
-    });
+
+    const updatedLostIdentities = Object.keys(identities).reduce(
+      (acc, identity) => {
+        if (addresses.includes(identity.toLowerCase())) {
+          acc[identity] = identities[identity];
+        }
+        return acc;
+      },
+      { ...(lostIdentities ?? {}) },
+    );
 
     const updatedIdentities = Object.values(accounts).reduce(
       (identitiesMap: Record<string, AccountIdentityEntry>, account) => {
@@ -786,7 +794,7 @@ export default class PreferencesController {
 
     this.store.updateState({
       identities: updatedIdentities,
-      lostIdentities,
+      lostIdentities: updatedLostIdentities,
       selectedAddress: selectedAccount?.address || '', // it will be an empty string during onboarding
     });
   }

app/scripts/lib/ComposableObservableStore.js

@@ -59,7 +59,7 @@ export default class ComposableObservableStore extends ObservableStore {
       const store = config[key];
       if (store.subscribe) {
         config[key].subscribe((state) => {
-          this.updateState({ [key]: state });
+          this.#onStateChange(key, state);
         });
       } else {
         this.controllerMessenger.subscribe(
@@ -69,7 +69,7 @@ export default class ComposableObservableStore extends ObservableStore {
             if (this.persist) {
               updatedState = getPersistentState(state, config[key].metadata);
             }
-            this.updateState({ [key]: updatedState });
+            this.#onStateChange(key, updatedState);
           },
         );
       }
@@ -104,4 +104,12 @@ export default class ComposableObservableStore extends ObservableStore {
     }
     return flatState;
   }
+
+  #onStateChange(controllerKey, newState) {
+    const oldState = this.getState()[controllerKey];
+
+    this.updateState({ [controllerKey]: newState });
+
+    this.emit('stateChange', { oldState, newState, controllerKey });
+  }
 }

app/scripts/lib/PatchStore.test.ts

@@ -0,0 +1,183 @@
+import ComposableObservableStore from './ComposableObservableStore';
+import { PatchStore } from './PatchStore';
+import { sanitizeUIState } from './state-utils';
+
+jest.mock('./state-utils');
+
+function createComposableStoreMock() {
+  return {
+    on: jest.fn(),
+    removeListener: jest.fn(),
+  } as unknown as jest.Mocked<ComposableObservableStore>;
+}
+
+function triggerStateChange(
+  composableStoreMock: jest.Mocked<ComposableObservableStore>,
+  oldState: Record<string, unknown>,
+  newState: Record<string, unknown>,
+) {
+  composableStoreMock.on.mock.calls[0][1]({
+    controllerKey: 'test-controller',
+    oldState,
+    newState,
+  });
+}
+
+describe('PatchStore', () => {
+  const sanitizeUIStateMock = jest.mocked(sanitizeUIState);
+
+  beforeEach(() => {
+    jest.resetAllMocks();
+    sanitizeUIStateMock.mockImplementation((state) => state);
+  });
+
+  describe('flushPendingPatches', () => {
+    it('returns pending patches created by composable store events', () => {
+      const composableStoreMock = createComposableStoreMock();
+      const patchStore = new PatchStore(composableStoreMock);
+
+      triggerStateChange(
+        composableStoreMock,
+        { test1: 'value1' },
+        { test1: 'value2' },
+      );
+
+      triggerStateChange(
+        composableStoreMock,
+        { test2: true },
+        { test2: false },
+      );
+
+      const patches = patchStore.flushPendingPatches();
+
+      expect(patches).toEqual([
+        {
+          op: 'replace',
+          path: ['test1'],
+          value: 'value2',
+        },
+        {
+          op: 'replace',
+          path: ['test2'],
+          value: false,
+        },
+      ]);
+    });
+
+    it('ignores state properties if old and new state is shallow equal', () => {
+      const objectMock = {};
+      const composableStoreMock = createComposableStoreMock();
+      const patchStore = new PatchStore(composableStoreMock);
+
+      triggerStateChange(
+        composableStoreMock,
+        { test1: 'value1' },
+        { test1: 'value1' },
+      );
+
+      triggerStateChange(
+        composableStoreMock,
+        { test2: objectMock },
+        { test2: objectMock },
+      );
+
+      triggerStateChange(
+        composableStoreMock,
+        { test3: { test: 'value' } },
+        { test3: { test: 'value' } },
+      );
+
+      const patches = patchStore.flushPendingPatches();
+
+      expect(patches).toEqual([
+        {
+          op: 'replace',
+          path: ['test3'],
+          value: { test: 'value' },
+        },
+      ]);
+    });
+
+    it('returns empty array if no composable store events', () => {
+      const composableStoreMock = createComposableStoreMock();
+      const patchStore = new PatchStore(composableStoreMock);
+
+      const patches = patchStore.flushPendingPatches();
+
+      expect(patches).toEqual([]);
+    });
+
+    it('clears pending patches', () => {
+      const composableStoreMock = createComposableStoreMock();
+      const patchStore = new PatchStore(composableStoreMock);
+
+      triggerStateChange(
+        composableStoreMock,
+        { test1: 'value1' },
+        { test1: 'value2' },
+      );
+
+      const patches1 = patchStore.flushPendingPatches();
+      const patches2 = patchStore.flushPendingPatches();
+
+      expect(patches1).toHaveLength(1);
+      expect(patches2).toHaveLength(0);
+    });
+
+    it('sanitizes state in patches', () => {
+      const composableStoreMock = createComposableStoreMock();
+      const patchStore = new PatchStore(composableStoreMock);
+
+      sanitizeUIStateMock.mockReturnValueOnce({ test2: 'value' });
+
+      triggerStateChange(
+        composableStoreMock,
+        { test1: false },
+        { test1: true },
+      );
+
+      const patches = patchStore.flushPendingPatches();
+
+      expect(patches).toEqual([
+        {
+          op: 'replace',
+          path: ['test2'],
+          value: 'value',
+        },
+      ]);
+    });
+
+    it('adds isInitialized patch if vault in new state', () => {
+      const composableStoreMock = createComposableStoreMock();
+      const patchStore = new PatchStore(composableStoreMock);
+
+      triggerStateChange(composableStoreMock, { vault: 0 }, { vault: 123 });
+
+      const patches = patchStore.flushPendingPatches();
+
+      expect(patches).toEqual([
+        {
+          op: 'replace',
+          path: ['vault'],
+          value: 123,
+        },
+        {
+          op: 'replace',
+          path: ['isInitialized'],
+          value: true,
+        },
+      ]);
+    });
+  });
+
+  describe('destroy', () => {
+    it('removes listener from composable store', () => {
+      const composableStoreMock = createComposableStoreMock();
+      const patchStore = new PatchStore(composableStoreMock);
+
+      patchStore.destroy();
+
+      expect(composableStoreMock.removeListener).toHaveBeenCalledTimes(1);
+    });
+  });
+});

app/scripts/lib/PatchStore.ts

@@ -0,0 +1,103 @@
+import { createProjectLogger } from '@metamask/utils';
+import { Patch } from 'immer';
+import { v4 as uuid } from 'uuid';
+import ComposableObservableStore from './ComposableObservableStore';
+import { sanitizeUIState } from './state-utils';
+
+const log = createProjectLogger('patch-store');
+
+export class PatchStore {
+  private id: string;
+
+  private observableStore: ComposableObservableStore;
+
+  private pendingPatches: Map<string, Patch> = new Map();
+
+  private listener: (request: {
+    controllerKey: string;
+    oldState: Record<string, unknown>;
+    newState: Record<string, unknown>;
+  }) => void;
+
+  constructor(observableStore: ComposableObservableStore) {
+    this.id = uuid();
+    this.observableStore = observableStore;
+    this.listener = this._onStateChange.bind(this);
+
+    this.observableStore.on('stateChange', this.listener);
+
+    log('Created', this.id);
+  }
+
+  flushPendingPatches(): Patch[] {
+    const patches = [...this.pendingPatches.values()];
+
+    this.pendingPatches.clear();
+
+    for (const patch of patches) {
+      log('Flushed', patch.path.join('.'), this.id, patch);
+    }
+
+    return patches;
+  }
+
+  destroy() {
+    this.observableStore.removeListener('stateChange', this.listener);
+    log('Destroyed', this.id);
+  }
+
+  private _onStateChange({
+    oldState,
+    newState,
+  }: {
+    controllerKey: string;
+    oldState: Record<string, unknown>;
+    newState: Record<string, unknown>;
+  }) {
+    const sanitizedNewState = sanitizeUIState(newState);
+    const patches = this._generatePatches(oldState, sanitizedNewState);
+    const isInitialized = Boolean(newState.vault);
+
+    if (isInitialized) {
+      patches.push({
+        op: 'replace',
+        path: ['isInitialized'],
+        value: isInitialized,
+      });
+    }
+
+    if (!patches.length) {
+      return;
+    }
+
+    for (const patch of patches) {
+      const path = patch.path.join('.');
+
+      this.pendingPatches.set(path, patch);
+
+      log('Updated', path, this.id, patch);
+    }
+  }
+
+  private _generatePatches(
+    oldState: Record<string, unknown>,
+    newState: Record<string, unknown>,
+  ): Patch[] {
+    return Object.keys(newState)
+      .map((key) => {
+        const oldData = oldState[key];
+        const newData = newState[key];
+
+        if (oldData === newData) {
+          return null;
+        }
+
+        return {
+          op: 'replace',
+          path: [key],
+          value: newData,
+        };
+      })
+      .filter(Boolean) as Patch[];
+  }
+}