fix(assets-controller): avoid recomputing getStateForTransactionPay on every transaction state change

packages/assets-controller/CHANGELOG.md

@@ -15,10 +15,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Changed
 
 - Persist `assetsPrice` in `AssetsController` state so last-known prices survive app restart / state rehydration (refreshed on the next price fetch via existing `PriceDataSource` TTL) ([#9566](https://github.com/MetaMask/core/pull/9566))
-- Memoize hot-path asset ID / legacy-format conversions to cut repeated keccak256 (`toChecksumAddress`) and CAIP parsing ([#9555](https://github.com/MetaMask/core/pull/9555))
+- Memoize hot-path asset ID / legacy-format conversions to cut repeated keccak256 (`toChecksumAddress`) and CAIP parsing ([#9546](https://github.com/MetaMask/core/pull/9546), [#9555](https://github.com/MetaMask/core/pull/9555))
   - `normalizeAssetId` uses lodash `memoize` so already-normalized IDs skip re-checksumming across the pipeline
   - `formatExchangeRatesForBridge` caches the last result on input identity (`===` for BaseController slices, lodash `isEqual` for rebuilt native maps); exports `FormatExchangeRatesForBridgeParams`
-  - `formatStateForTransactionPay` caches the last result the same way (`isEqual` for rebuilt `accounts` / `nativeAssetIdentifiers`); exports `FormatStateForTransactionPayParams`
+  - `formatStateForTransactionPay` caches the last result the same way (`isEqual` for rebuilt `accounts` / `nativeAssetIdentifiers`) and freezes the cached result so later mutations cannot poison the cache; exports `FormatStateForTransactionPayParams`
   - `#getNativeAssetMap` returns a stable empty object when uncached so memoized formatters are not busted by `?? {}` identity churn
 - `TokenDataSource` EVM spam filtering now uses per-chain floors from Token API `GET /v1/suggestedOccurrenceFloors` (`queryApiClient.token.fetchV1SuggestedOccurrenceFloors`) instead of a hardcoded minimum of 3 occurrences. Chains missing from the response (or a failed floors fetch) still fall back to 3 ([#9537](https://github.com/MetaMask/core/pull/9537))
 - `TokensApiClient` (used by `RpcDataSource` / `TokenDetector`) now sets the token-list `occurrenceFloor` query param from the same Token API `GET /v1/suggestedOccurrenceFloors` endpoint (cached 1h), replacing the hardcoded Linea/MegaETH/Tempo special cases. Missing chains or failed fetches fall back to 3 ([#9537](https://github.com/MetaMask/core/pull/9537))

packages/assets-controller/src/utils/formatStateForTransactionPay.test.ts

@@ -3,6 +3,7 @@ import {
   clearFormatStateForTransactionPayCacheForTesting,
   formatStateForTransactionPay,
   AccountForLegacyFormat,
+  FormatStateForTransactionPayParams,
 } from './formatStateForTransactionPay';
 
 function price(
@@ -389,25 +390,21 @@ describe('formatStateForTransactionPay', () => {
       clearFormatStateForTransactionPayCacheForTesting();
     });
 
-    it('returns the same object when inputs are unchanged by identity / value', () => {
-      const assetsBalance = {
+    const makeParams = (): FormatStateForTransactionPayParams => ({
+      accounts: [{ ...ACCOUNT_1 }],
+      assetsBalance: {
         [ACCOUNT_1.id]: {
-          [ETH_NATIVE_ID]: { amount: '100' } as AssetBalance,
+          [USDC_ASSET_ID]: { amount: '1000000' } as AssetBalance,
         },
-      };
-      const assetsInfo = {
-        [ETH_NATIVE_ID]: ETH_NATIVE_METADATA,
-      };
-      const assetsPrice = {};
-      const params = {
-        accounts: [ACCOUNT_1],
-        assetsBalance,
-        assetsInfo,
-        assetsPrice,
-        selectedCurrency: 'usd',
-        nativeAssetIdentifiers: EVM_NATIVE_IDS,
-        networkConfigurationsByChainId: EVM_NETWORK_CONFIGS,
-      };
+      },
+      assetsInfo: {},
+      assetsPrice: {},
+      selectedCurrency: 'usd',
+      nativeAssetIdentifiers: { ...EVM_NATIVE_IDS },
+    });
+
+    it('returns the cached result when called again with identical inputs', () => {
+      const params = makeParams();
 
       const first = formatStateForTransactionPay(params);
       const second = formatStateForTransactionPay({
@@ -419,38 +416,70 @@ describe('formatStateForTransactionPay', () => {
       expect(second).toBe(first);
     });
 
-    it('recomputes when assetsBalance identity changes', () => {
-      const assetsInfo = {
-        [ETH_NATIVE_ID]: ETH_NATIVE_METADATA,
-      };
-      const first = formatStateForTransactionPay({
-        accounts: [ACCOUNT_1],
+    it('recomputes when a state slice reference changes', () => {
+      const params = makeParams();
+
+      const first = formatStateForTransactionPay(params);
+      const second = formatStateForTransactionPay({
+        ...params,
         assetsBalance: {
           [ACCOUNT_1.id]: {
-            [ETH_NATIVE_ID]: { amount: '100' } as AssetBalance,
+            [USDC_ASSET_ID]: { amount: '2000000' } as AssetBalance,
           },
         },
-        assetsInfo,
-        assetsPrice: {},
-        selectedCurrency: 'usd',
-        nativeAssetIdentifiers: EVM_NATIVE_IDS,
-        networkConfigurationsByChainId: EVM_NETWORK_CONFIGS,
       });
+
+      expect(second).not.toBe(first);
+      const accountLower = ACCOUNT_1.address.toLowerCase();
+      expect(second.tokenBalances[accountLower]['0x1'][USDC_ADDRESS]).toBe(
+        '0x1e8480',
+      );
+    });
+
+    it('recomputes when the accounts change', () => {
+      const params = makeParams();
+
+      const first = formatStateForTransactionPay(params);
       const second = formatStateForTransactionPay({
-        accounts: [ACCOUNT_1],
-        assetsBalance: {
-          [ACCOUNT_1.id]: {
-            [ETH_NATIVE_ID]: { amount: '200' } as AssetBalance,
+        ...params,
+        accounts: [
+          {
+            id: 'account-2',
+            address: '0x0Ac1dF02185025F65202660F8167210A80dD5086',
           },
-        },
-        assetsInfo,
-        assetsPrice: {},
-        selectedCurrency: 'usd',
-        nativeAssetIdentifiers: EVM_NATIVE_IDS,
-        networkConfigurationsByChainId: EVM_NETWORK_CONFIGS,
+        ],
       });
 
       expect(second).not.toBe(first);
+      expect(second.tokenBalances).toStrictEqual({});
+    });
+
+    it('recomputes when the selected currency changes', () => {
+      const params = makeParams();
+
+      const first = formatStateForTransactionPay(params);
+      const second = formatStateForTransactionPay({
+        ...params,
+        selectedCurrency: 'eur',
+      });
+
+      expect(second).not.toBe(first);
+      expect(second.currentCurrency).toBe('eur');
+    });
+
+    it('freezes the result so consumers cannot poison the cache', () => {
+      const result = formatStateForTransactionPay(makeParams());
+
+      expect(Object.isFrozen(result)).toBe(true);
+    });
+
+    it('produces the same output whether served from cache or recomputed', () => {
+      const params = makeParams();
+
+      const cached = formatStateForTransactionPay(params);
+      const recomputed = formatStateForTransactionPay(makeParams());
+
+      expect(recomputed).toStrictEqual(cached);
     });
   });
 });

packages/assets-controller/src/utils/formatStateForTransactionPay.ts

@@ -61,58 +61,75 @@ export type FormatStateForTransactionPayParams = {
   networkConfigurationsByChainId?: Record<string, { nativeCurrency?: string }>;
 };
 
-let lastCall: {
-  params: FormatStateForTransactionPayParams;
-  result: TransactionPayLegacyFormat;
-} | null = null;
-
 function amountToHex(amount: string): `0x${string}` {
   const hexString = BigInt(amount).toString(16);
   return `0x${hexString}`;
 }
 
+/**
+ * Determines whether two sets of {@link formatStateForTransactionPay}
+ * parameters are identical for memoization purposes.
+ *
+ * State slices (`assetsBalance`, `assetsInfo`, `assetsPrice`,
+ * `networkConfigurationsByChainId`) are compared by reference since
+ * BaseController state updates are immutable. The `accounts` and
+ * `nativeAssetIdentifiers` inputs are rebuilt on every call, so they are
+ * compared by value instead.
+ *
+ * @param a - Previous parameters.
+ * @param b - Next parameters.
+ * @returns True if the parameters are identical.
+ */
+function isTransactionPayParamsIdentical(
+  a: FormatStateForTransactionPayParams,
+  b: FormatStateForTransactionPayParams,
+): boolean {
+  return (
+    a.assetsBalance === b.assetsBalance &&
+    a.assetsInfo === b.assetsInfo &&
+    a.assetsPrice === b.assetsPrice &&
+    a.selectedCurrency === b.selectedCurrency &&
+    a.networkConfigurationsByChainId === b.networkConfigurationsByChainId &&
+    isEqual(a.accounts, b.accounts) &&
+    isEqual(a.nativeAssetIdentifiers, b.nativeAssetIdentifiers)
+  );
+}
+
 function getAmountFromBalance(balance: AssetBalance): string {
   return typeof balance === 'object' && balance !== null && 'amount' in balance
     ? String((balance as { amount: string }).amount)
     : '0';
 }
 
+let lastCall: {
+  params: FormatStateForTransactionPayParams;
+  result: TransactionPayLegacyFormat;
+} | null = null;
+
 /**
  * Converts AssetsController state into the legacy format consumed by
  * transaction-pay-controller (TokenBalancesController, AccountTrackerController,
  * TokensController, TokenRatesController, CurrencyRateController shapes).
  *
- * Memoized on input identity for BaseController state slices (`===`) and
- * lodash `isEqual` for rebuilt arrays/maps (`accounts`, `nativeAssetIdentifiers`).
- * `AssetsController:getStateForTransactionPay` is invoked on every
- * `TransactionController:stateChange` while its inputs only change when the
- * assets pipeline updates; recomputing runs keccak256 (`toChecksumAddress`) and
- * CAIP parsing per asset.
+ * The last result is memoized on input identity: this function is invoked (via
+ * `AssetsController:getStateForTransactionPay`) on every
+ * `TransactionController:stateChange`, but its inputs only change when the
+ * assets pipeline updates. Recomputing runs keccak256 (`toChecksumAddress`)
+ * per asset per call, which dominates CPU profiles during transaction
+ * approval.
  *
  * @param params - Conversion parameters.
  * @returns Legacy-compatible state for transaction-pay-controller.
  */
 export function formatStateForTransactionPay(
   params: FormatStateForTransactionPayParams,
 ): TransactionPayLegacyFormat {
-  if (
-    lastCall?.params.assetsBalance === params.assetsBalance &&
-    lastCall.params.assetsInfo === params.assetsInfo &&
-    lastCall.params.assetsPrice === params.assetsPrice &&
-    lastCall.params.selectedCurrency === params.selectedCurrency &&
-    lastCall.params.networkConfigurationsByChainId ===
-      params.networkConfigurationsByChainId &&
-    isEqual(lastCall.params.accounts, params.accounts) &&
-    isEqual(
-      lastCall.params.nativeAssetIdentifiers,
-      params.nativeAssetIdentifiers,
-    )
-  ) {
+  if (lastCall && isTransactionPayParamsIdentical(lastCall.params, params)) {
     return lastCall.result;
   }
 
   const result = computeStateForTransactionPay(params);
-  lastCall = { params, result };
+  lastCall = { params, result: Object.freeze(result) };
   return result;
 }
 
@@ -124,7 +141,7 @@ export function clearFormatStateForTransactionPayCacheForTesting(): void {
 }
 
 /**
- * Performs the actual legacy-format conversion for
+ * Performs the actual state conversion for
  * {@link formatStateForTransactionPay}.
  *
  * @param params - Conversion parameters.