wallet ctx init commit

packages/sdk/src/api/WalletCtx.ts

@@ -0,0 +1,98 @@
+import { ApiPromise } from '@polkadot/api';
+import { UnsubscribePromise, VoidFn } from '@polkadot/api-base/types';
+import { u32, Vec } from '@polkadot/types';
+import { ITuple } from '@polkadot/types-codec/types';
+import { OrmlTokensAccountData } from '@polkadot/types/lookup';
+
+import { memoize1 } from '@thi.ng/memoize';
+
+import { AssetClient, BalanceClient } from '../client';
+import { Asset } from '../types';
+import { BigNumber, ZERO } from '../utils/bignumber';
+
+export class WalletCtx extends BalanceClient {
+  protected readonly assetClient: AssetClient;
+
+  protected ctx: Map<string, BigNumber> = new Map([]);
+  protected subs: VoidFn[] = [];
+
+  private assets: Asset[] = [];
+
+  private memBalances = memoize1((account: string) => {
+    console.log('Wallet mem balances', account, '✅');
+    return this.getBalances(account);
+  });
+
+  constructor(api: ApiPromise) {
+    super(api);
+    this.assetClient = new AssetClient(this.api);
+  }
+
+  async getBalancesMem(account: string): Promise<Map<string, BigNumber>> {
+    return this.memBalances(account);
+  }
+
+  private async getBalances(account: string): Promise<Map<string, BigNumber>> {
+    this.unsubscribe();
+
+    const balances = await this.getAccountBalanceData(account);
+    balances.forEach(([token, data]) => {
+      this.ctx.set(token.toString(), this.calculateFreeBalance(data));
+    });
+
+    this.assets = await this.assetClient.getOnChainAssets(true);
+    this.subs = await Promise.all([
+      this.subscribeSystemAccountBalance(account),
+      this.subscribeTokensAccountBalance(account),
+      this.subscribeErc20AccountBalance(account),
+    ]);
+
+    console.log(`Wallet assets(${this.assets.length})`, '✅');
+    console.log(`Wallet subs(${this.subs.length})`, '✅');
+
+    return this.ctx;
+  }
+
+  unsubscribe() {
+    this.subs.forEach((unsub) => {
+      unsub();
+    });
+  }
+
+  private async getAccountBalanceData(account: string) {
+    return await this.api.call.currenciesApi.accounts<
+      Vec<ITuple<[u32, OrmlTokensAccountData]>>
+    >(account.toString());
+  }
+
+  private subscribeTokensAccountBalance(account: string): UnsubscribePromise {
+    const updateCtx = this.updateCtxCallback(this.ctx);
+    return this.subscribeTokenBalance(account, this.assets, (balances) =>
+      balances.forEach(([token, balance]) => updateCtx(token, balance))
+    );
+  }
+
+  private subscribeErc20AccountBalance(account: string): UnsubscribePromise {
+    const updateCtx = this.updateCtxCallback(this.ctx);
+    return this.subscribeErc20Balance(account, this.assets, (balances) =>
+      balances.forEach(([token, balance]) => updateCtx(token, balance))
+    );
+  }
+
+  private subscribeSystemAccountBalance(account: string): UnsubscribePromise {
+    const updateCtx = this.updateCtxCallback(this.ctx);
+    return this.subscribeSystemBalance(account, (token, balance) =>
+      updateCtx(token, balance)
+    );
+  }
+
+  private updateCtxCallback(ctx: Map<string, BigNumber>) {
+    return function (token: string, balance: BigNumber) {
+      if (balance.isGreaterThan(ZERO)) {
+        ctx.set(token, balance);
+      } else {
+        ctx.delete(token);
+      }
+    };
+  }
+}

packages/sdk/src/api/index.ts

@@ -1,2 +1,3 @@
 export { Router } from './Router';
 export { TradeRouter } from './TradeRouter';
+export { WalletCtx } from './WalletCtx';

packages/sdk/src/pool/PoolClient.ts

@@ -46,7 +46,7 @@ export abstract class PoolClient extends BalanceClient {
     this.mem = this.mem + 1;
   }
 
-  async getMemPools(): Promise<PoolBase[]> {
+  async getPoolsMem(): Promise<PoolBase[]> {
     return this.memPools(this.mem);
   }
 

packages/sdk/src/pool/PoolService.ts

@@ -83,15 +83,15 @@ export class PoolService implements IPoolService {
       const pools = await Promise.all(
         this.clients
           .filter((client) => client.isSupported())
-          .map((client) => client.getMemPools())
+          .map((client) => client.getPoolsMem())
       );
       return pools.flat();
     }
 
     const pools = await Promise.all(
       this.clients
         .filter((client) => includeOnly.some((t) => t === client.getPoolType()))
-        .map((client) => client.getMemPools())
+        .map((client) => client.getPoolsMem())
     );
     return pools.flat();
   }

packages/sdk/test/script/examples/router/onBlockPoolChange.ts

@@ -1,12 +1,14 @@
 import { ApiPromise } from '@polkadot/api';
 import { ApiUrl, PolkadotExecutor } from '../../executor';
 import { PoolService } from '../../../../src/pool';
-import { TradeRouter } from '../../../../src/api';
+import { TradeRouter, WalletCtx } from '../../../../src/api';
+import { HYDRADX_OMNIPOOL_ADDRESS } from '../../../../src/consts';
 
 class GetOnBlockPoolChangeExample extends PolkadotExecutor {
   async script(api: ApiPromise): Promise<any> {
     const poolService = new PoolService(api);
     const router = new TradeRouter(poolService);
+    const walletCtx = new WalletCtx(api);
 
     const result = await Promise.all([
       router.getBestSpotPrice('5', '0'),
@@ -17,8 +19,9 @@ class GetOnBlockPoolChangeExample extends PolkadotExecutor {
     ]);
     console.log('Call spot 5x simultaneously', result);
 
-    api.rpc.chain.subscribeNewHeads(async (_lastHeader) => {
-      console.log('==================');
+    api.rpc.chain.subscribeNewHeads(async (lastHeader) => {
+      const block = lastHeader.number.toString();
+      console.log('===============', block, '===============');
       router.getPools().then((p) => {
         p.forEach((o) => {
           o.tokens.forEach((t) => {
@@ -29,6 +32,9 @@ class GetOnBlockPoolChangeExample extends PolkadotExecutor {
           });
         });
       });
+
+      const balance = await walletCtx.getBalancesMem(HYDRADX_OMNIPOOL_ADDRESS);
+      console.log('balance', balance);
     });
 
     return [];