Fix Immature balance bug

scripts/__mocks__/network.js

@@ -62,8 +62,8 @@
                     '0100000001458de14f4f4fecfdeebfef09fb16e761bbd15029f37bec0a63b86808cbb8a512010000006b483045022100ef4f4364aea7604d749aaff7a2609e3a51a12f49500b7910b34ced0d0837e1db022012d153d96ebcb94e9b905a609c0ea97cdc99ae961be2848e0e8f2f695379c21201210371eca6799221b82cbba9e880a8a5a0f47d811f3ff5cad346931406ab0a0469eeffffffff0200e1f505000000001976a9148952bf31104625a7b3e6fcf4c79b35c6849ef74d88ac905cfb2f010000001976a9144e8d2fcf6d909c62597e4defd1c26d50842d73df88ac00000000';
                 await wallet.addTransaction(Transaction.fromHex(tx_1));
             } else {
-                throw new Error(
-                    'Add getLatestTxs implementation for this wallet! ' +
+                console.warn(
+                    'getLatestTxs did not find any txs this wallet! ' +
                         wallet.getKeyToExport()
                 );
             }
@@ -85,7 +85,7 @@
          * @param{number} blockHeight
          * @param{boolean} skipCoinstake
          */
         (blockHeight, skipCoinstake = false) => {
             if (!this.#mapBlocks.has(blockHeight)) {
                 throw new Error('Requested block does not exist!');
             }

scripts/mempool.js

@@ -8,7 +8,6 @@ export const OutpointState = {
     P2CS: 1 << 2, // This is a P2CS outpoint
 
     SPENT: 1 << 3, // This outpoint has been spent
-    IMMATURE: 1 << 4, // Coinbase/coinstake that it's not mature (hence not spendable) yet
     LOCKED: 1 << 5, // Coins in the LOCK set
 };
 
@@ -23,10 +22,13 @@ export class Mempool {
     #txmap = new Map();
 
     /**
-     * balance cache, mapping filter -> balance
-     * @type{Map<number, number>}
+     * Object containing balances of the wallet
      */
-    #balances = new Map();
+    #balances = {
+        balance: new CachableBalance(),
+        coldBalance: new CachableBalance(),
+        immatureBalance: new CachableBalance(),
+    };
 
     /**
      * Add a transaction to the mempool
@@ -54,7 +56,7 @@ export class Mempool {
      */
     setOutpointStatus(outpoint, status) {
         this.#outpointStatus.set(outpoint.toUnique(), status);
-        this.#invalidateBalanceCache();
+        this.invalidateBalanceCache();
     }
 
     /**
@@ -65,7 +67,7 @@ export class Mempool {
     addOutpointStatus(outpoint, status) {
         const oldStatus = this.#outpointStatus.get(outpoint.toUnique());
         this.#outpointStatus.set(outpoint.toUnique(), oldStatus | status);
-        this.#invalidateBalanceCache();
+        this.invalidateBalanceCache();
     }
 
     /**
@@ -76,7 +78,7 @@ export class Mempool {
     removeOutpointStatus(outpoint, status) {
         const oldStatus = this.#outpointStatus.get(outpoint.toUnique());
         this.#outpointStatus.set(outpoint.toUnique(), oldStatus & ~status);
-        this.#invalidateBalanceCache();
+        this.invalidateBalanceCache();
     }
 
     /**
@@ -144,64 +146,91 @@ export class Mempool {
             .reduce((acc, u) => acc + u?.value ?? 0, 0);
     }
 
+    /**
+     * Loop through the unspent balance of the wallet
+     * @template T
+     * @param {number} requirement - Requirement that outpoints must have
+     * @param {T} initialValue - initial value of the result
+     * @param {balanceIterator} fn
+     * @returns {T}
+     */
+    loopSpendableBalance(requirement, initialValue, fn) {
+        for (const tx of this.#txmap.values()) {
+            for (const [index, vout] of tx.vout.entries()) {
+                const status = this.getOutpointStatus(
+                    new COutpoint({ txid: tx.txid, n: index })
+                );
+                if (status & (OutpointState.SPENT | OutpointState.LOCKED)) {
+                    continue;
+                }
+                if ((status & requirement) === requirement) {
+                    initialValue = fn(tx, vout, initialValue);
+                }
+            }
+        }
+        return initialValue;
+    }
+
     /**
      * @param {object} o - options
-     * @param {number} [o.filter] - A filter to apply to all UTXOs. For example
-     * `OutpointState.P2CS` will NOT return P2CS transactions.
-     * By default it's `OutpointState.SPENT | OutpointState.IMMATURE | OutpointState.LOCKED`
      * @param {number} [o.requirement] - A requirement to apply to all UTXOs. For example
      * `OutpointState.P2CS` will only return P2CS transactions.
      * @param {number} [o.target] - Number of satoshis needed. The method will return early when the value of the UTXOs has been reached, plus a bit to account for change
      * By default it's MAX_SAFE_INTEGER
+     * @param {boolean} [o.includeImmature] - If set to true immature UTXOs will be included
+     * @param {number} [o.blockCount] - Current number of blocks
      * @returns {UTXO[]} a list of unspent transaction outputs
      */
     getUTXOs({
-        filter = OutpointState.SPENT |
-            OutpointState.IMMATURE |
-            OutpointState.LOCKED,
         requirement = 0,
+        includeImmature = false,
         target = Number.POSITIVE_INFINITY,
+        blockCount,
     } = {}) {
-        const utxos = [];
-        let value = 0;
-        for (const [o, status] of this.#outpointStatus) {
-            const outpoint = COutpoint.fromUnique(o);
-            if (status & filter) {
-                continue;
-            }
-            if ((status & requirement) !== requirement) {
-                continue;
+        return this.loopSpendableBalance(
+            requirement,
+            { utxos: [], bal: 0 },
+            (tx, vout, currentValue) => {
+                if (
+                    (!includeImmature && tx.isImmature(blockCount)) ||
+                    (currentValue.bal >= (target * 11) / 10 &&
+                        currentValue.bal > 0)
+                ) {
+                    return currentValue;
+                }
+                const n = tx.vout.findIndex((element) => element === vout);
+                currentValue.utxos.push(
+                    new UTXO({
+                        outpoint: new COutpoint({ txid: tx.txid, n }),
+                        script: vout.script,
+                        value: vout.value,
+                    })
+                );
+                currentValue.bal += vout.value;
+                return currentValue;
             }
-            utxos.push(this.outpointToUTXO(outpoint));
-            value += utxos.at(-1).value;
-            if (value >= (target * 11) / 10) {
-                break;
+        ).utxos;
+    }
+
+    #balanceInternal(requirement, blockCount, includeImmature = false) {
+        return this.loopSpendableBalance(
+            requirement,
+            0,
+            (tx, vout, currentValue) => {
+                if (!tx.isImmature(blockCount)) {
+                    return currentValue + vout.value;
+                } else if (includeImmature) {
+                    return currentValue + vout.value;
+                }
+                return currentValue;
             }
-        }
-        return utxos;
+        );
     }
 
-    /**
-     * @param {number} filter
-     */
-    getBalance(filter) {
-        if (this.#balances.has(filter)) {
-            return this.#balances.get(filter);
-        }
-        const bal = Array.from(this.#outpointStatus)
-            .filter(([_, status]) => !(status & OutpointState.SPENT))
-            .filter(([_, status]) => status & filter)
-            .reduce((acc, [o]) => {
-                const outpoint = COutpoint.fromUnique(o);
-                const tx = this.#txmap.get(outpoint.txid);
-                return acc + tx.vout[outpoint.n].value;
-            }, 0);
-        this.#balances.set(filter, bal);
-        return bal;
-    }
-
-    #invalidateBalanceCache() {
-        this.#balances = new Map();
+    invalidateBalanceCache() {
+        this.#balances.immatureBalance.invalidate();
+        this.#balances.balance.invalidate();
+        this.#balances.coldBalance.invalidate();
         this.#emitBalanceUpdate();
     }
 
@@ -228,15 +257,76 @@ export class Mempool {
         return Array.from(this.#txmap.values());
     }
 
-    get balance() {
-        return this.getBalance(OutpointState.P2PKH);
+    /**
+     * @param blockCount - chain height
+     */
+    getBalance(blockCount) {
+        return this.#balances.balance.getOrUpdateInvalid(() => {
+            return this.#balanceInternal(
+                OutpointState.OURS | OutpointState.P2PKH,
+                blockCount
+            );
+        });
     }
 
-    get coldBalance() {
-        return this.getBalance(OutpointState.P2CS);
+    /**
+     * @param blockCount - chain height
+     */
+    getColdBalance(blockCount) {
+        return this.#balances.coldBalance.getOrUpdateInvalid(() => {
+            return this.#balanceInternal(
+                OutpointState.OURS | OutpointState.P2CS,
+                blockCount
+            );
+        });
     }
 
-    get immatureBalance() {
-        return this.getBalance(OutpointState.IMMATURE);
+    /**
+     * @param blockCount - chain height
+     */
+    getImmatureBalance(blockCount) {
+        return this.#balances.immatureBalance.getOrUpdateInvalid(() => {
+            return (
+                this.#balanceInternal(OutpointState.OURS, blockCount, true) -
+                this.getBalance(blockCount) -
+                this.getColdBalance(blockCount)
+            );
+        });
     }
 }
+
+class CachableBalance {
+    /**
+     * @type {number}
+     * represents a cachable balance
+     */
+    value = -1;
+
+    isValid() {
+        return this.value != -1;
+    }
+    invalidate() {
+        this.value = -1;
+    }
+
+    /**
+     * Return the cached balance if it's valid, or re-compute and return.
+     * @param {Function} fn - function with which calculate the balance
+     * @returns {number} cached balance
+     */
+    getOrUpdateInvalid(fn) {
+        if (!this.isValid()) {
+            this.value = fn();
+        }
+        return this.value;
+    }
+}
+
+/**
+ * @template T
+ * @typedef {Function} balanceIterator
+ * @param {import('./transaction.js').Transaction} tx
+ * @param {CTxOut} vout
+ * @param {T} currentValue - the current value iterated
+ * @returns {number} amount
+ */

scripts/transaction.js

@@ -4,7 +4,7 @@ import { hexToBytes, bytesToHex, dSHA256 } from './utils.js';
 import { OP } from './script.js';
 import { varIntToNum, deriveAddress } from './encoding.js';
 import * as nobleSecp256k1 from '@noble/secp256k1';
-import { SAPLING_TX_VERSION } from './chain_params.js';
+import { cChainParams, SAPLING_TX_VERSION } from './chain_params.js';
 
 /** An Unspent Transaction Output, used as Inputs of future transactions */
 export class COutpoint {
@@ -167,6 +167,20 @@ export class Transaction {
         );
     }
 
+    /**
+     * @param {Transaction} tx - transaction we want to check
+     * @returns {boolean}
+     */
+    isImmature(blockCount) {
+        if (this.isCoinStake() || this.isCoinBase()) {
+            return (
+                blockCount - this.blockHeight <
+                cChainParams.current.coinbaseMaturity
+            );
+        }
+        return false;
+    }
+
     static fromHex(hex) {
         const tx = new Transaction();
         return tx.fromHex(hex);

scripts/wallet.js

@@ -792,6 +792,7 @@ export class Wallet {
             .getUTXOs({
                 requirement: OutpointState.P2PKH | OutpointState.OURS,
                 target,
+                blockCount,
             })
             .map((u) => {
                 return {
@@ -809,6 +810,8 @@ export class Wallet {
     subscribeToNetworkEvents() {
         getEventEmitter().on('new-block', async (block) => {
             if (this.#isSynced) {
+                // Invalidate the balance cache to keep immature balance updated
+                this.#mempool.invalidateBalanceCache();
                 await this.getLatestBlocks(block);
                 getEventEmitter().emit('new-tx');
             }
@@ -964,11 +967,11 @@ export class Wallet {
     ) {
         let balance;
         if (useDelegatedInputs) {
-            balance = this.#mempool.coldBalance;
+            balance = this.coldBalance;
         } else if (useShieldInputs) {
             balance = this.#shield.getBalance();
         } else {
-            balance = this.#mempool.balance;
+            balance = this.balance;
         }
         if (balance < value) {
             throw new Error('Not enough balance');
@@ -1007,6 +1010,7 @@ export class Wallet {
             const utxos = this.#mempool.getUTXOs({
                 requirement: requirement | OutpointState.OURS,
                 target: value,
+                blockCount,
             });
             transactionBuilder.addUTXOs(utxos);
 
@@ -1191,6 +1195,7 @@ export class Wallet {
         return this.#mempool
             .getUTXOs({
                 requirement: OutpointState.P2PKH | OutpointState.OURS,
+                blockCount,
             })
             .filter((u) => u.value === collateralValue);
     }
@@ -1203,15 +1208,15 @@ export class Wallet {
     }
 
     get balance() {
-        return this.#mempool.balance;
+        return this.#mempool.getBalance(blockCount);
     }
 
     get immatureBalance() {
-        return this.#mempool.immatureBalance;
+        return this.#mempool.getImmatureBalance(blockCount);
     }
 
     get coldBalance() {
-        return this.#mempool.coldBalance;
+        return this.#mempool.getColdBalance(blockCount);
     }
 
     /**
@@ -1239,7 +1244,7 @@ export class Wallet {
 /**
  * @type{Wallet}
  */
-export const wallet = new Wallet({ nAccountL: 0 }); // For now we are using only the 0-th account, (TODO: update once account system is done)
+export const wallet = new Wallet({ nAccount: 0 }); // For now we are using only the 0-th account, (TODO: update once account system is done)
 
 /**
  * Clean a Seed Phrase string and verify it's integrity