refactor: separate the xudt balance calculation of total/pending amount

ShookLyngs · ShookLyngs · commit ed4788030e1e · 2024-08-20T11:18:40.000+08:00
diff --git a/src/routes/rgbpp/address.ts b/src/routes/rgbpp/address.ts
@@ -5,14 +5,23 @@ import { ZodTypeProvider } from 'fastify-type-provider-zod';
 import { CKBTransaction, Cell, IsomorphicTransaction, Script, XUDTBalance } from './types';
 import z from 'zod';
 import { Env } from '../../env';
-import { buildPreLockArgs, getXudtTypeScript, isScriptEqual, isTypeAssetSupported } from '@rgbpp-sdk/ckb';
-import { groupBy } from 'lodash';
+import {
+  isScriptEqual,
+  buildPreLockArgs,
+  getRgbppLockScript,
+  getXudtTypeScript,
+  isTypeAssetSupported,
+} from '@rgbpp-sdk/ckb';
+import { groupBy, uniq } from 'lodash';
 import { BI } from '@ckb-lumos/lumos';
 import { UTXO } from '../../services/bitcoin/schema';
 import { Transaction as BTCTransaction } from '../bitcoin/types';
 import { TransactionWithStatus } from '../../services/ckb';
 import { computeScriptHash } from '@ckb-lumos/lumos/utils';
 import { filterCellsByTypeScript, getTypeScript } from '../../utils/typescript';
+import { unpackRgbppLockArgs } from '@rgbpp-sdk/btc/lib/ckb/molecule';
+import { TestnetTypeMap } from '../../constants';
+import { remove0x } from '@rgbpp-sdk/btc';
 
 const addressRoutes: FastifyPluginCallback<Record<never, never>, Server, ZodTypeProvider> = (fastify, _, done) => {
   const env: Env = fastify.container.resolve('env');
@@ -52,6 +61,18 @@ const addressRoutes: FastifyPluginCallback<Record<never, never>, Server, ZodType
     return cells;
   }
 
+  /**
+   * Filter RgbppLock cells by cells
+   */
+  function getRgbppLockCellsByCells(cells: Cell[]): Cell[] {
+    const rgbppLockScript = getRgbppLockScript(env.NETWORK === 'mainnet', TestnetTypeMap[env.NETWORK]);
+    return cells.filter(
+      (cell) =>
+        rgbppLockScript.codeHash === cell.cellOutput.lock.codeHash &&
+        rgbppLockScript.hashType === cell.cellOutput.lock.hashType,
+    );
+  }
+
   fastify.get(
     '/:btc_address/assets',
     {
@@ -147,13 +168,14 @@ const addressRoutes: FastifyPluginCallback<Record<never, never>, Server, ZodType
         throw fastify.httpErrors.badRequest('Unsupported type asset');
       }
 
-      const utxos = await getUxtos(btc_address, no_cache);
       const xudtBalances: Record<string, XUDTBalance> = {};
+      const utxos = await getUxtos(btc_address, no_cache);
 
-      let cells = await getRgbppAssetsCells(btc_address, utxos, no_cache);
-      cells = typeScript ? filterCellsByTypeScript(cells, typeScript) : cells;
-
-      const availableXudtBalances = await fastify.rgbppCollector.getRgbppBalanceByCells(cells);
+      // Find confirmed RgbppLock Xudt assets
+      const confirmedUtxos = utxos.filter((utxo) => utxo.status.confirmed);
+      const confirmedCells = await getRgbppAssetsCells(btc_address, confirmedUtxos, no_cache);
+      const confirmedTargetCells = filterCellsByTypeScript(confirmedCells, typeScript);
+      const availableXudtBalances = await fastify.rgbppCollector.getRgbppBalanceByCells(confirmedTargetCells);
       Object.keys(availableXudtBalances).forEach((key) => {
         const { amount, ...xudtInfo } = availableXudtBalances[key];
         xudtBalances[key] = {
@@ -164,6 +186,7 @@ const addressRoutes: FastifyPluginCallback<Record<never, never>, Server, ZodType
         };
       });
 
+      // Find all unconfirmed RgbppLock Xudt outputs
       const pendingUtxos = utxos.filter(
         (utxo) =>
           !utxo.status.confirmed ||
@@ -172,19 +195,14 @@ const addressRoutes: FastifyPluginCallback<Record<never, never>, Server, ZodType
       );
       const pendingUtxosGroup = groupBy(pendingUtxos, (utxo) => utxo.txid);
       const pendingTxids = Object.keys(pendingUtxosGroup);
-
       const pendingOutputCellsGroup = await Promise.all(
         pendingTxids.map(async (txid) => {
           const cells = await fastify.transactionProcessor.getPendingOutputCellsByTxid(txid);
           const lockArgsSet = new Set(pendingUtxosGroup[txid].map((utxo) => buildPreLockArgs(utxo.vout)));
           return cells.filter((cell) => lockArgsSet.has(cell.cellOutput.lock.args));
         }),
       );
-      let pendingOutputCells = pendingOutputCellsGroup.flat();
-      if (typeScript) {
-        pendingOutputCells = filterCellsByTypeScript(pendingOutputCells, typeScript);
-      }
-
+      const pendingOutputCells = filterCellsByTypeScript(pendingOutputCellsGroup.flat(), typeScript);
       const pendingXudtBalances = await fastify.rgbppCollector.getRgbppBalanceByCells(pendingOutputCells);
       Object.values(pendingXudtBalances).forEach(({ amount, type_hash, ...xudtInfo }) => {
         if (!xudtBalances[type_hash]) {
@@ -200,6 +218,49 @@ const addressRoutes: FastifyPluginCallback<Record<never, never>, Server, ZodType
         xudtBalances[type_hash].pending_amount = BI.from(xudtBalances[type_hash].pending_amount)
           .add(BI.from(amount))
           .toHexString();
+      });
+
+      // Find spent RgbppLock Xudt assets in unconfirmed transactions' inputs
+      const allTxs = await fastify.bitcoin.getAddressTxs({ address: btc_address });
+      const unconfirmedTxids = allTxs.filter((tx) => !tx.status.confirmed).map((tx) => tx.txid);
+      const spendingInputCellsGroup = await Promise.all(
+        unconfirmedTxids.map(async (txid) => {
+          const inputCells = await fastify.transactionProcessor.getPendingInputCellsByTxid(txid);
+          const inputRgbppCells = getRgbppLockCellsByCells(filterCellsByTypeScript(inputCells, typeScript));
+          const inputCellLockArgs = inputRgbppCells.map((cell) => unpackRgbppLockArgs(cell.cellOutput.lock.args));
+
+          const txids = uniq(inputCellLockArgs.map((args) => remove0x(args.btcTxid)));
+          const txs = await Promise.all(txids.map((txid) => fastify.bitcoin.getTx({ txid })));
+          const txsMap = txs.reduce(
+            (sum, tx, index) => {
+              const txid = txids[index];
+              sum[txid] = tx ?? null;
+              return sum;
+            },
+            {} as Record<string, BTCTransaction | null>,
+          );
+
+          return inputRgbppCells.filter((cell, index) => {
+            const lockArgs = inputCellLockArgs[index];
+            const tx = txsMap[remove0x(lockArgs.btcTxid)];
+            const utxo = tx?.vout[lockArgs.outIndex];
+            return utxo?.scriptpubkey_address === btc_address;
+          });
+        }),
+      );
+      const spendingInputCells = spendingInputCellsGroup.flat();
+      const spendingXudtBalances = await fastify.rgbppCollector.getRgbppBalanceByCells(spendingInputCells);
+      Object.values(spendingXudtBalances).forEach(({ amount, type_hash, ...xudtInfo }) => {
+        if (!xudtBalances[type_hash]) {
+          xudtBalances[type_hash] = {
+            ...xudtInfo,
+            type_hash,
+            total_amount: '0x0',
+            available_amount: '0x0',
+            pending_amount: '0x0',
+          };
+        }
+
         xudtBalances[type_hash].total_amount = BI.from(xudtBalances[type_hash].total_amount)
           .add(BI.from(amount))
           .toHexString();
@@ -322,18 +383,18 @@ const addressRoutes: FastifyPluginCallback<Record<never, never>, Server, ZodType
             } as const;
           }
 
-          const inputOutpoints = isomorphicTx.ckbRawTx?.inputs || isomorphicTx.ckbTx?.inputs || [];
-          const inputs = await fastify.ckb.getInputCellsByOutPoint(
-            inputOutpoints.map((input) => input.previousOutput) as CKBComponents.OutPoint[],
-          );
+          const inputs = isomorphicTx.ckbRawTx?.inputs || isomorphicTx.ckbTx?.inputs || [];
+          const inputCells = await fastify.ckb.getInputCellsByOutPoint(inputs.map((input) => input.previousOutput!));
+          const inputCellOutputs = inputCells.map((cell) => cell.cellOutput);
+
           const outputs = isomorphicTx.ckbRawTx?.outputs || isomorphicTx.ckbTx?.outputs || [];
 
           return {
             btcTx,
             isRgbpp: true,
             isomorphicTx: {
               ...isomorphicTx,
-              inputs,
+              inputs: inputCellOutputs,
               outputs,
             },
           } as const;
diff --git a/src/services/ckb.ts b/src/services/ckb.ts
@@ -22,7 +22,8 @@ import {
 import { computeScriptHash } from '@ckb-lumos/lumos/utils';
 import DataCache from './base/data-cache';
 import { scriptToHash } from '@nervosnetwork/ckb-sdk-utils';
-import { OutputCell } from '../routes/rgbpp/types';
+import { Cell } from '../routes/rgbpp/types';
+import { uniq } from 'lodash';
 
 export type TransactionWithStatus = Awaited<ReturnType<CKBRPC['getTransaction']>>;
 
@@ -326,14 +327,27 @@ export default class CKBClient {
     return null;
   }
 
-  public async getInputCellsByOutPoint(outPoints: CKBComponents.OutPoint[]): Promise<OutputCell[]> {
-    const batchRequest = this.rpc.createBatchRequest(outPoints.map((outPoint) => ['getTransaction', outPoint.txHash]));
-    const txs = await batchRequest.exec();
-    const inputs = txs.map((tx: TransactionWithStatus, index: number) => {
-      const outPoint = outPoints[index];
-      return tx.transaction.outputs[BI.from(outPoint.index).toNumber()];
+  public async getInputCellsByOutPoint(outPoints: CKBComponents.OutPoint[]): Promise<Cell[]> {
+    const txHashes = uniq(outPoints.map((outPoint) => outPoint.txHash));
+    const batchRequest = this.rpc.createBatchRequest(txHashes.map((txHash) => ['getTransaction', txHash]));
+    const txs: TransactionWithStatus[] = await batchRequest.exec();
+    const txsMap = txs.reduce(
+      (acc, tx: TransactionWithStatus) => {
+        acc[tx.transaction.hash] = tx;
+        return acc;
+      },
+      {} as Record<string, TransactionWithStatus>,
+    );
+    return outPoints.map((outPoint) => {
+      const tx = txsMap[outPoint.txHash];
+      const outPointIndex = BI.from(outPoint.index).toNumber();
+      return Cell.parse({
+        cellOutput: tx.transaction.outputs[outPointIndex],
+        data: tx.transaction.outputsData[outPointIndex],
+        blockHash: tx.txStatus.blockHash,
+        outPoint,
+      });
     });
-    return inputs;
   }
 
   /**
diff --git a/src/services/transaction.ts b/src/services/transaction.ts
@@ -608,14 +608,34 @@ export default class TransactionProcessor
     const { ckbVirtualResult } = job.data;
     const outputs = ckbVirtualResult.ckbRawTx.outputs;
     return outputs.map((output, index) => {
-      const cell: Cell = {
+      return Cell.parse({
         cellOutput: output,
         data: ckbVirtualResult.ckbRawTx.outputsData[index],
-      };
-      return cell;
+      });
     });
   }
 
+  /**
+   * get pending input cells by txid, get ckb input cells from the uncompleted job
+   * @param txid - the transaction id
+   */
+  public async getPendingInputCellsByTxid(txid: string): Promise<Cell[]> {
+    const job = await this.getTransactionRequest(txid);
+    if (!job) {
+      return [];
+    }
+
+    // get ckb input cells from the uncompleted job only
+    const state = await job.getState();
+    if (state === 'completed' || state === 'failed') {
+      return [];
+    }
+
+    const { ckbVirtualResult } = job.data;
+    const inputOutPoints = ckbVirtualResult.ckbRawTx.inputs.map((input) => input.previousOutput!);
+    return await this.cradle.ckb.getInputCellsByOutPoint(inputOutPoints);
+  }
+
   /**
    * Retry all failed jobs in the queue
    * @param maxAttempts - the max attempts to retry
diff --git a/test/routes/__snapshots__/token.test.ts.snap b/test/routes/__snapshots__/token.test.ts.snap
@@ -1,17 +1,5 @@
 // Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
 
-exports[`\`/token/generate\` - 400 1`] = `
-"[
-  {
-    "code": "invalid_type",
-    "expected": "object",
-    "received": "null",
-    "path": [],
-    "message": "Expected object, received null"
-  }
-]"
-`;
-
 exports[`\`/token/generate\` - without params 1`] = `
 "[
   {
diff --git a/test/routes/rgbpp/__snapshots__/address.test.ts.snap b/test/routes/rgbpp/__snapshots__/address.test.ts.snap
@@ -6889,7 +6889,7 @@ exports[`/:btc_address/balance - with pending_amount 1`] = `
       "name": "Unique BBQ",
       "pending_amount": "0x5f5e100",
       "symbol": "",
-      "total_amount": "0xbebc200",
+      "total_amount": "0x5f5e100",
       "type_hash": "0x78e21efcf107e7886eadeadecd1a01cfb88f1e5617f4438685db55b3a540d202",
       "type_script": {
         "args": "0x30d3fbec9ceba691770d57c6d06bdb98cf0f82bef0ca6e87687a118d6ce1e7b7",
@@ -6903,7 +6903,7 @@ exports[`/:btc_address/balance - with pending_amount 1`] = `
       "name": "XUDT Test Token",
       "pending_amount": "0x5f5e100",
       "symbol": "PDD",
-      "total_amount": "0x5f5e100",
+      "total_amount": "0x0",
       "type_hash": "0x10f511f2efb0027191b97ac5b4bd77374ffdac7399e8527d76f5f9bd32e7d35b",
       "type_script": {
         "args": "0x8c556e92974a8dd8237719020a259d606359ac2cc958cb8bda77a1c3bb3cd93b",
