From 99ebad88eec829f5209159d3321347b67acc550a Mon Sep 17 00:00:00 2001 From: Jongsun Suh Date: Wed, 4 Dec 2024 10:21:57 -0500 Subject: [PATCH] Replace `MemStoreControllersComposedState` usage with `BackgroundStateProxy` --- .../controllers/metametrics-controller.ts | 28 +++++++++++-------- app/scripts/lib/PatchStore.ts | 16 +++++------ app/scripts/lib/state-utils.ts | 8 ++++-- 3 files changed, 30 insertions(+), 22 deletions(-) diff --git a/app/scripts/controllers/metametrics-controller.ts b/app/scripts/controllers/metametrics-controller.ts index 947ad4369be4..63c520e5d299 100644 --- 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 | 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 index ee69efc954b4..43b75537ac9e 100644 --- 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(newState).reduce< + return getKnownPropertyNames(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 index cbcb81779e99..2a829b038a73 100644 --- 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;