Introduce instances, apply to balance derives (multi-result) (#3094)

* Introducce instances, apply to balance derives (multi-result)

* Nonce extraction

* Optimize

packages/api-derive/src/balances/account.ts

@@ -1,24 +1,24 @@
 // Copyright 2017-2021 @polkadot/api-derive authors & contributors
 // SPDX-License-Identifier: Apache-2.0
 
-import type { ApiInterfaceRx } from '@polkadot/api/types';
+import type { ApiInterfaceRx, QueryableStorageMultiArg } from '@polkadot/api/types';
 import type { AccountData, AccountId, AccountIndex, AccountInfo, Address, Balance, Index } from '@polkadot/types/interfaces';
 import type { ITuple } from '@polkadot/types/types';
 import type { Observable } from '@polkadot/x-rxjs';
-import type { DeriveBalancesAccount } from '../types';
+import type { DeriveBalancesAccount, DeriveBalancesAccountData } from '../types';
 
-import { isFunction } from '@polkadot/util';
+import { assert, isFunction } from '@polkadot/util';
 import { combineLatest, of } from '@polkadot/x-rxjs';
 import { map, switchMap } from '@polkadot/x-rxjs/operators';
 
 import { memo } from '../util';
 
-type Result = [Balance, Balance, Balance, Balance, Index];
+type BalanceResult = [Balance, Balance, Balance, Balance];
 
-function calcBalances (api: ApiInterfaceRx, [accountId, [freeBalance, reservedBalance, frozenFee, frozenMisc, accountNonce]]: [AccountId, Result]): DeriveBalancesAccount {
+type Result = [Index, BalanceResult[]];
+
+function getBalance (api: ApiInterfaceRx, [freeBalance, reservedBalance, frozenFee, frozenMisc]: BalanceResult): DeriveBalancesAccountData {
   return {
-    accountId,
-    accountNonce,
     freeBalance,
     frozenFee,
     frozenMisc,
@@ -27,6 +27,19 @@ function calcBalances (api: ApiInterfaceRx, [accountId, [freeBalance, reservedBa
   };
 }
 
+function calcBalances (api: ApiInterfaceRx, [accountId, [accountNonce, balances]]: [AccountId, Result]): DeriveBalancesAccount {
+  const primary = balances[0];
+
+  assert(primary, 'No balances retrieved for account');
+
+  return {
+    accountId,
+    accountNonce,
+    additional: balances.filter((_, index) => index !== 0).map((b) => getBalance(api, b)),
+    ...getBalance(api, primary)
+  };
+}
+
 // old
 function queryBalancesFree (api: ApiInterfaceRx, accountId: AccountId): Observable<Result> {
   return api.queryMulti<[Balance, Balance, Index]>([
@@ -34,21 +47,25 @@ function queryBalancesFree (api: ApiInterfaceRx, accountId: AccountId): Observab
     [api.query.balances.reservedBalance, accountId],
     [api.query.system.accountNonce, accountId]
   ]).pipe(
-    map(([freeBalance, reservedBalance, accountNonce]): Result =>
-      [freeBalance, reservedBalance, api.registry.createType('Balance'), api.registry.createType('Balance'), accountNonce]
-    )
+    map(([freeBalance, reservedBalance, accountNonce]): Result => [
+      accountNonce,
+      [[freeBalance, reservedBalance, api.registry.createType('Balance'), api.registry.createType('Balance')]]
+    ])
   );
 }
 
-function queryBalancesAccount (api: ApiInterfaceRx, accountId: AccountId): Observable<Result> {
-  return api.queryMulti<[AccountData, Index]>([
-    [api.query.balances.account, accountId],
-    [api.query.system.accountNonce, accountId]
-  ]).pipe(
-    map(([{ feeFrozen, free, miscFrozen, reserved }, accountNonce]): Result =>
-      [free, reserved, feeFrozen, miscFrozen, accountNonce]
+function queryBalancesAccount (api: ApiInterfaceRx, accountId: AccountId, modules: string[] = ['balances']): Observable<Result> {
+  const balances = modules.map((m): QueryableStorageMultiArg<'rxjs'> => [api.query[m].account, accountId]);
+  const extract = (data: AccountData[]) =>
+    data.map(({ feeFrozen, free, miscFrozen, reserved }): BalanceResult => [free, reserved, feeFrozen, miscFrozen]);
+
+  return isFunction(api.query.system.account)
+    ? api.queryMulti<[AccountInfo, ...(AccountData[])]>([[api.query.system.account, accountId], ...balances]).pipe(
+      map(([{ nonce }, ...balances]): Result => [nonce, extract(balances)])
     )
-  );
+    : api.queryMulti<[Index, ...(AccountData[])]>([[api.query.system.accountNonce, accountId], ...balances]).pipe(
+      map(([nonce, ...balances]): Result => [nonce, extract(balances)])
+    );
 }
 
 function queryCurrent (api: ApiInterfaceRx, accountId: AccountId): Observable<Result> {
@@ -60,7 +77,7 @@ function queryCurrent (api: ApiInterfaceRx, accountId: AccountId): Observable<Re
         : (infoOrTuple as [Index, AccountData])[1];
       const accountNonce = (infoOrTuple as AccountInfo).nonce || (infoOrTuple as [Index, AccountData])[0];
 
-      return [free, reserved, feeFrozen, miscFrozen, accountNonce];
+      return [accountNonce, [[free, reserved, feeFrozen, miscFrozen]]];
     })
   );
 }
@@ -81,19 +98,29 @@ function queryCurrent (api: ApiInterfaceRx, accountId: AccountId): Observable<Re
  * ```
  */
 export function account (instanceId: string, api: ApiInterfaceRx): (address: AccountIndex | AccountId | Address | string) => Observable<DeriveBalancesAccount> {
+  const balanceInstances = api.registry.getModuleInstances(api.runtimeVersion.specName.toString(), 'balances');
+
   return memo(instanceId, (address: AccountIndex | AccountId | Address | string): Observable<DeriveBalancesAccount> =>
     api.derive.accounts.accountId(address).pipe(
       switchMap((accountId): Observable<[AccountId, Result]> =>
         (accountId
           ? combineLatest([
             of(accountId),
-            isFunction(api.query.system.account)
-              ? queryCurrent(api, accountId)
-              : isFunction(api.query.balances.account)
-                ? queryBalancesAccount(api, accountId)
-                : queryBalancesFree(api, accountId)
+            balanceInstances
+              ? queryBalancesAccount(api, accountId, balanceInstances)
+              : isFunction(api.query.system.account)
+                ? queryCurrent(api, accountId)
+                : isFunction(api.query.balances.account)
+                  ? queryBalancesAccount(api, accountId)
+                  : queryBalancesFree(api, accountId)
+          ])
+          : of([
+            api.registry.createType('AccountId'),
+            [
+              api.registry.createType('Index'),
+              [[api.registry.createType('Balance'), api.registry.createType('Balance'), api.registry.createType('Balance'), api.registry.createType('Balance')]]
+            ]
           ])
-          : of([api.registry.createType('AccountId'), [api.registry.createType('Balance'), api.registry.createType('Balance'), api.registry.createType('Balance'), api.registry.createType('Balance'), api.registry.createType('Index')]])
         )
       ),
       map((result): DeriveBalancesAccount => calcBalances(api, result))

packages/api-derive/src/balances/all.ts

@@ -9,7 +9,7 @@ import type { DeriveBalancesAccount, DeriveBalancesAll } from '../types';
 
 import BN from 'bn.js';
 
-import { bnMax, isFunction } from '@polkadot/util';
+import { bnMax } from '@polkadot/util';
 import { combineLatest, of } from '@polkadot/x-rxjs';
 import { map, switchMap } from '@polkadot/x-rxjs/operators';
 
@@ -67,6 +67,7 @@ function calcBalances (api: ApiInterfaceRx, [{ accountId, accountNonce, freeBala
   return {
     accountId,
     accountNonce,
+    additional: [],
     availableBalance,
     freeBalance,
     frozenFee,
@@ -113,11 +114,14 @@ function queryCurrent (api: ApiInterfaceRx, accountId: AccountId): Observable<Re
         [api.query.balances.locks, accountId],
         [api.query.vesting.vesting, accountId]
       ])
-      : api.query.balances.locks(accountId).pipe(
-        map((locks): [Vec<BalanceLock>, Option<VestingInfo>] =>
-          [locks, api.registry.createType('Option<VestingInfo>')]
+      // TODO We need to check module instances here as well, not only the balances module
+      : api.query.balances?.locks
+        ? api.query.balances.locks(accountId).pipe(
+          map((locks): [Vec<BalanceLock>, Option<VestingInfo>] =>
+            [locks, api.registry.createType('Option<VestingInfo>')]
+          )
         )
-      )
+        : of([api.registry.createType('Vec<BalanceLock>'), api.registry.createType('Option<VestingInfo>')] as [Vec<BalanceLock>, Option<VestingInfo>])
   ).pipe(
     map(([locks, optVesting]): ResultBalance =>
       [optVesting.unwrapOr(null), locks]
@@ -148,7 +152,7 @@ export function all (instanceId: string, api: ApiInterfaceRx): (address: Account
           ? combineLatest([
             of(account),
             api.derive.chain.bestNumber(),
-            isFunction(api.query.system.account) || isFunction(api.query.balances.account)
+            api.query.system?.account || api.query.balances?.account
               ? queryCurrent(api, account.accountId)
               : queryOld(api, account.accountId)
           ])

packages/api-derive/src/types.ts

@@ -13,9 +13,7 @@ export * from './parachains/types';
 export * from './session/types';
 export * from './staking/types';
 
-export interface DeriveBalancesAccount {
-  accountId: AccountId;
-  accountNonce: Index;
+export interface DeriveBalancesAccountData {
   freeBalance: Balance;
   frozenFee: Balance;
   frozenMisc: Balance;
@@ -23,6 +21,12 @@ export interface DeriveBalancesAccount {
   votingBalance: Balance;
 }
 
+export interface DeriveBalancesAccount extends DeriveBalancesAccountData {
+  accountId: AccountId;
+  accountNonce: Index;
+  additional: DeriveBalancesAccountData[];
+}
+
 export interface DeriveBalancesAll extends DeriveBalancesAccount {
   isVesting: boolean;
   lockedBalance: Balance;