Replace MemStoreControllersComposedState usage with `BackgroundStat…

…eProxy`
MetaMask · Dec 4, 2024 · 99ebad8 · 99ebad8
1 parent 32f90c5
commit 99ebad8
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 22 deletions.
diff --git a/app/scripts/controllers/metametrics-controller.ts b/app/scripts/controllers/metametrics-controller.ts
@@ -64,6 +64,7 @@ import Analytics from '../lib/segment/analytics';
 import { ENVIRONMENT } from '../../../development/build/constants';
 ///: END:ONLY_INCLUDE_IF
 
+import { BackgroundStateProxy } from '../../../shared/types/metamask';
 import type {
   PreferencesControllerState,
   PreferencesControllerGetStateAction,
@@ -963,7 +964,7 @@ export default class MetaMetricsController extends BaseController<
     }
   }
 
-  handleMetaMaskStateUpdate(newState: MetaMaskState): void {
+  handleMetaMaskStateUpdate(newState: BackgroundStateProxy): void {
     const userTraits = this._buildUserTraitsObject(newState);
     if (userTraits) {
       this.identify(userTraits);
@@ -1128,7 +1129,7 @@ export default class MetaMetricsController extends BaseController<
    * @returns traits that have changed since last update
    */
   _buildUserTraitsObject(
-    metamaskState: MetaMaskState,
+    metamaskState: BackgroundStateProxy,
   ): Partial<MetaMetricsUserTraits> | null {
     ///: BEGIN:ONLY_INCLUDE_IF(build-mmi)
     const mmiAccountAddress =
@@ -1258,11 +1259,13 @@ export default class MetaMetricsController extends BaseController<
    *
    * @param allNfts
    */
-  #getAllNFTsFlattened = memoize((allNfts: MetaMaskState['allNfts'] = {}) => {
-    return Object.values(allNfts).reduce((result: Nft[], chainNFTs) => {
-      return result.concat(...Object.values(chainNFTs));
-    }, []);
-  });
+  #getAllNFTsFlattened = memoize(
+    (allNfts: BackgroundStateProxy['NftController']['allNfts'] = {}) => {
+      return Object.values(allNfts).reduce((result: Nft[], chainNFTs) => {
+        return result.concat(...Object.values(chainNFTs));
+      }, []);
+    },
+  );
 
   /**
    * Returns the number of unique NFT addresses the user
@@ -1271,7 +1274,7 @@ export default class MetaMetricsController extends BaseController<
    * @param allNfts
    */
   #getAllUniqueNFTAddressesLength(
-    allNfts: MetaMaskState['allNfts'] = {},
+    allNfts: BackgroundStateProxy['NftController']['allNfts'] = {},
   ): number {
     const allNFTAddresses = this.#getAllNFTsFlattened(allNfts).map(
       (nft) => nft.address,
@@ -1284,7 +1287,9 @@ export default class MetaMetricsController extends BaseController<
    * @param allTokens
    * @returns number of unique token addresses
    */
-  #getNumberOfTokens(allTokens: MetaMaskState['allTokens']): number {
+  #getNumberOfTokens(
+    allTokens: BackgroundStateProxy['TokensController']['allTokens'],
+  ): number {
     return Object.values(allTokens).reduce((result, accountsByChain) => {
       return result + sum(Object.values(accountsByChain).map(size));
     }, 0);
@@ -1508,8 +1513,9 @@ export default class MetaMetricsController extends BaseController<
    *
    * @param metamaskState
    */
-  #getPetnameAddressCount(metamaskState: MetaMaskState): number {
-    const addressNames = metamaskState.names?.[NameType.ETHEREUM_ADDRESS] ?? {};
+  #getPetnameAddressCount(metamaskState: BackgroundStateProxy): number {
+    const addressNames =
+      metamaskState.NameController.names?.[NameType.ETHEREUM_ADDRESS] ?? {};
 
     return Object.keys(addressNames).reduce((totalCount, address) => {
       const addressEntry = addressNames[address];

diff --git a/app/scripts/lib/PatchStore.ts b/app/scripts/lib/PatchStore.ts
@@ -1,7 +1,7 @@
 import { createProjectLogger, getKnownPropertyNames } from '@metamask/utils';
 import { Patch } from 'immer';
 import { v4 as uuid } from 'uuid';
-import { MemStoreControllersComposedState as BackgroundState } from '../../../shared/types/metamask';
+import type { BackgroundStateProxy } from '../../../shared/types/metamask';
 import ComposableObservableStore from './ComposableObservableStore';
 import { sanitizeUIState } from './state-utils';
 
@@ -16,8 +16,8 @@ export class PatchStore {
 
   private listener: (request: {
     controllerKey: string;
-    oldState: BackgroundState;
-    newState: BackgroundState;
+    oldState: BackgroundStateProxy;
+    newState: BackgroundStateProxy;
   }) => void;
 
   constructor(observableStore: ComposableObservableStore) {
@@ -52,8 +52,8 @@ export class PatchStore {
     newState,
   }: {
     controllerKey: string;
-    oldState: BackgroundState;
-    newState: BackgroundState;
+    oldState: BackgroundStateProxy;
+    newState: BackgroundStateProxy;
   }) {
     const sanitizedNewState = sanitizeUIState(newState);
     const patches = this._generatePatches(oldState, sanitizedNewState);
@@ -81,10 +81,10 @@ export class PatchStore {
   }
 
   private _generatePatches(
-    oldState: BackgroundState,
-    newState: BackgroundState,
+    oldState: BackgroundStateProxy,
+    newState: BackgroundStateProxy,
   ): Patch[] {
-    return getKnownPropertyNames<keyof BackgroundState>(newState).reduce<
+    return getKnownPropertyNames<keyof BackgroundStateProxy>(newState).reduce<
       Patch[]
     >((patches, controllerName) => {
       Object.keys(oldState[controllerName]).forEach((key) => {

diff --git a/app/scripts/lib/state-utils.ts b/app/scripts/lib/state-utils.ts
@@ -1,10 +1,12 @@
 import { SnapControllerState } from '@metamask/snaps-controllers';
 import { isSnapId, Snap } from '@metamask/snaps-utils';
-import { MemStoreControllersComposedState as BackgroundState } from '../../../shared/types/metamask';
+import { BackgroundStateProxy } from '../../../shared/types/metamask';
 
 const REMOVE_KEYS = ['snapStates', 'unencryptedSnapStates', 'vault'] as const;
 
-export function sanitizeUIState(state: BackgroundState): BackgroundState {
+export function sanitizeUIState(
+  state: BackgroundStateProxy,
+): BackgroundStateProxy {
   const newState = { ...state };
 
   for (const key of REMOVE_KEYS) {
@@ -20,7 +22,7 @@ export function sanitizeUIState(state: BackgroundState): BackgroundState {
   return newState;
 }
 
-function sanitizeSnapData(state: BackgroundState) {
+function sanitizeSnapData(state: BackgroundStateProxy) {
   const snapsData: SnapControllerState['snaps'] | undefined =
     state.SnapController.snaps;