From e2b4e0d3e67de4f625c200baba1b08740d75ac22 Mon Sep 17 00:00:00 2001
From: Tarik Gul <47201679+TarikGul@users.noreply.github.com>
Date: Sun, 18 Aug 2024 14:51:19 -0700
Subject: [PATCH] Rework api.derive.balance.{account, all} to work with the new
 frame account data (#5955)

* Rework api.derive.balance.account to work with the new frame account data

* Ensure api.derive.balances.all correctly returns transferable

* Remove null from transferable

* set back to null
---
 packages/api-derive/src/balances/account.ts | 71 +++++++++++++++------
 packages/api-derive/src/balances/all.ts     |  8 +++
 packages/api-derive/src/balances/types.ts   | 28 ++++++++
 3 files changed, 89 insertions(+), 18 deletions(-)

diff --git a/packages/api-derive/src/balances/account.ts b/packages/api-derive/src/balances/account.ts
index 087cfc87c557..b938c82a841c 100644
--- a/packages/api-derive/src/balances/account.ts
+++ b/packages/api-derive/src/balances/account.ts
@@ -4,6 +4,7 @@
 import type { Observable } from 'rxjs';
 import type { QueryableStorageEntry } from '@polkadot/api-base/types';
 import type { AccountData, AccountId, AccountIndex, AccountInfo, Address, Balance, Index } from '@polkadot/types/interfaces';
+import type { FrameSystemAccountInfo, PalletBalancesAccountData } from '@polkadot/types/lookup';
 import type { ITuple } from '@polkadot/types/types';
 import type { DeriveApi, DeriveBalancesAccount, DeriveBalancesAccountData } from '../types.js';
 
@@ -15,7 +16,9 @@ import { memo } from '../util/index.js';
 
 type BalanceResult = [Balance, Balance, Balance, Balance];
 
-type Result = [Index, BalanceResult[]];
+type Result = [Index, BalanceResult[], AccountType];
+
+interface AccountType { isFrameAccountData: boolean }
 
 type DeriveCustomAccount = DeriveApi['derive'] & Record<string, {
   customAccount?: DeriveApi['query']['balances']['account']
@@ -25,24 +28,38 @@ function zeroBalance (api: DeriveApi) {
   return api.registry.createType('Balance');
 }
 
-function getBalance (api: DeriveApi, [freeBalance, reservedBalance, frozenFee, frozenMisc]: BalanceResult): DeriveBalancesAccountData {
+function getBalance (api: DeriveApi, [freeBalance, reservedBalance, frozenFeeOrFrozen, frozenMiscOrFlags]: BalanceResult, accType: AccountType): DeriveBalancesAccountData {
   const votingBalance = api.registry.createType('Balance', freeBalance.toBn());
 
+  if (accType.isFrameAccountData) {
+    return {
+      frameSystemAccountInfo: {
+        flags: frozenMiscOrFlags,
+        frozen: frozenFeeOrFrozen
+      },
+      freeBalance,
+      frozenFee: api.registry.createType('Balance', 0),
+      frozenMisc: api.registry.createType('Balance', 0),
+      reservedBalance,
+      votingBalance
+    };
+  }
+
   return {
     freeBalance,
-    frozenFee,
-    frozenMisc,
+    frozenFee: frozenFeeOrFrozen,
+    frozenMisc: frozenMiscOrFlags,
     reservedBalance,
     votingBalance
   };
 }
 
-function calcBalances (api: DeriveApi, [accountId, [accountNonce, [primary, ...additional]]]: [AccountId, Result]): DeriveBalancesAccount {
+function calcBalances (api: DeriveApi, [accountId, [accountNonce, [primary, ...additional], accType]]: [AccountId, Result]): DeriveBalancesAccount {
   return objectSpread({
     accountId,
     accountNonce,
-    additional: additional.map((b) => getBalance(api, b))
-  }, getBalance(api, primary));
+    additional: additional.map((b) => getBalance(api, b, accType))
+  }, getBalance(api, primary, accType));
 }
 
 // old
@@ -54,7 +71,8 @@ function queryBalancesFree (api: DeriveApi, accountId: AccountId): Observable<Re
   ]).pipe(
     map(([freeBalance, reservedBalance, accountNonce]): Result => [
       accountNonce,
-      [[freeBalance, reservedBalance, zeroBalance(api), zeroBalance(api)]]
+      [[freeBalance, reservedBalance, zeroBalance(api), zeroBalance(api)]],
+      { isFrameAccountData: false }
     ])
   );
 }
@@ -62,7 +80,8 @@ function queryBalancesFree (api: DeriveApi, accountId: AccountId): Observable<Re
 function queryNonceOnly (api: DeriveApi, accountId: AccountId): Observable<Result> {
   const fill = (nonce: Index): Result => [
     nonce,
-    [[zeroBalance(api), zeroBalance(api), zeroBalance(api), zeroBalance(api)]]
+    [[zeroBalance(api), zeroBalance(api), zeroBalance(api), zeroBalance(api)]],
+    { isFrameAccountData: false }
   ];
 
   return isFunction(api.query.system.account)
@@ -83,7 +102,8 @@ function queryBalancesAccount (api: DeriveApi, accountId: AccountId, modules: st
 
   const extract = (nonce: Index, data: AccountData[]): Result => [
     nonce,
-    data.map(({ feeFrozen, free, miscFrozen, reserved }): BalanceResult => [free, reserved, feeFrozen, miscFrozen])
+    data.map(({ feeFrozen, free, miscFrozen, reserved }): BalanceResult => [free, reserved, feeFrozen, miscFrozen]),
+    { isFrameAccountData: false }
   ];
 
   // NOTE this is for the first case where we do have instances specified
@@ -106,7 +126,7 @@ function queryBalancesAccount (api: DeriveApi, accountId: AccountId, modules: st
 
 function querySystemAccount (api: DeriveApi, accountId: AccountId): Observable<Result> {
   // AccountInfo is current, support old, eg. Edgeware
-  return api.query.system.account<AccountInfo | ITuple<[Index, AccountData]>>(accountId).pipe(
+  return api.query.system.account<AccountInfo | FrameSystemAccountInfo | ITuple<[Index, AccountData]>>(accountId).pipe(
     map((infoOrTuple): Result => {
       const data = (infoOrTuple as AccountInfo).nonce
         ? (infoOrTuple as AccountInfo).data
@@ -117,16 +137,30 @@ function querySystemAccount (api: DeriveApi, accountId: AccountId): Observable<R
       if (!data || data.isEmpty) {
         return [
           nonce,
-          [[zeroBalance(api), zeroBalance(api), zeroBalance(api), zeroBalance(api)]]
+          [[zeroBalance(api), zeroBalance(api), zeroBalance(api), zeroBalance(api)]],
+          { isFrameAccountData: false }
         ];
       }
 
-      const { feeFrozen, free, miscFrozen, reserved } = data;
+      const isFrameType = !!(infoOrTuple as FrameSystemAccountInfo).data.frozen;
+
+      if (isFrameType) {
+        const { flags, free, frozen, reserved } = (data as unknown as PalletBalancesAccountData);
 
-      return [
-        nonce,
-        [[free, reserved, feeFrozen, miscFrozen]]
-      ];
+        return [
+          nonce,
+          [[free, reserved, frozen, flags]],
+          { isFrameAccountData: true }
+        ];
+      } else {
+        const { feeFrozen, free, miscFrozen, reserved } = data;
+
+        return [
+          nonce,
+          [[free, reserved, feeFrozen, miscFrozen]],
+          { isFrameAccountData: false }
+        ];
+      }
     })
   );
 }
@@ -168,7 +202,8 @@ export function account (instanceId: string, api: DeriveApi): (address: AccountI
           ])
           : of([api.registry.createType('AccountId'), [
             api.registry.createType('Index'),
-            [[zeroBalance(api), zeroBalance(api), zeroBalance(api), zeroBalance(api)]]
+            [[zeroBalance(api), zeroBalance(api), zeroBalance(api), zeroBalance(api)]],
+            { isFrameAccountData: false }
           ]])
         )
       ),
diff --git a/packages/api-derive/src/balances/all.ts b/packages/api-derive/src/balances/all.ts
index 60aaf49cb378..fa69a5cc9885 100644
--- a/packages/api-derive/src/balances/all.ts
+++ b/packages/api-derive/src/balances/all.ts
@@ -54,11 +54,19 @@ function calcLocked (api: DeriveApi, bestNumber: BlockNumber, locks: (PalletBala
 
 function calcShared (api: DeriveApi, bestNumber: BlockNumber, data: DeriveBalancesAccountData, locks: (PalletBalancesBalanceLock | BalanceLockTo212)[]): DeriveBalancesAllAccountData {
   const { allLocked, lockedBalance, lockedBreakdown, vestingLocked } = calcLocked(api, bestNumber, locks);
+  let transferable = null;
+
+  if (data.frameSystemAccountInfo?.frozen) {
+    const frozenReserveDiff = data.frameSystemAccountInfo.frozen.sub(data.reservedBalance);
+
+    transferable = api.registry.createType('Balance', allLocked ? 0 : data?.freeBalance.sub(bnMax(frozenReserveDiff, api.consts.balances.existentialDeposit)));
+  }
 
   return objectSpread({}, data, {
     availableBalance: api.registry.createType('Balance', allLocked ? 0 : bnMax(new BN(0), data?.freeBalance ? data.freeBalance.sub(lockedBalance) : new BN(0))),
     lockedBalance,
     lockedBreakdown,
+    transferable,
     vestingLocked
   });
 }
diff --git a/packages/api-derive/src/balances/types.ts b/packages/api-derive/src/balances/types.ts
index 384776c34fc7..aa28324d27b0 100644
--- a/packages/api-derive/src/balances/types.ts
+++ b/packages/api-derive/src/balances/types.ts
@@ -6,6 +6,10 @@ import type { PalletBalancesBalanceLock, PalletBalancesReserveData } from '@polk
 import type { BN } from '@polkadot/util';
 
 export interface DeriveBalancesAccountData {
+  frameSystemAccountInfo?: {
+    frozen: Balance;
+    flags: Balance;
+  }
   freeBalance: Balance;
   frozenFee: Balance;
   frozenMisc: Balance;
@@ -20,9 +24,33 @@ export interface DeriveBalancesAccount extends DeriveBalancesAccountData {
 }
 
 export interface DeriveBalancesAllAccountData extends DeriveBalancesAccountData {
+  /**
+   * Calculated available balance. This uses the formula: max(0, free - locked)
+   * This is only correct when the return type of `api.query.system.account` is `AccountInfo` which was replaced by `FrameSystemAccountInfo`.
+   * See `transferable` for the correct balance calculation.
+   *
+   * ref: https://github.com/paritytech/substrate/pull/12951
+   */
   availableBalance: Balance;
+  /**
+   * The amount of balance locked away.
+   */
   lockedBalance: Balance;
+  /**
+   * The breakdown of locked balances.
+   */
   lockedBreakdown: (PalletBalancesBalanceLock | BalanceLockTo212)[];
+  /**
+   * Calculated transferable balance. This uses the formula: free - max(frozen - reserve, ed)
+   * This is only correct when the return type of `api.query.system.account` is `FrameSystemAccountInfo`.
+   * Which is the most up to date calulcation for transferrable balances.
+   *
+   * ref: https://github.com/paritytech/polkadot-sdk/issues/1833
+   */
+  transferable: Balance | null;
+  /**
+   * Amount locked in vesting.
+   */
   vestingLocked: Balance;
 }