Skip to content

Commit

Permalink
Fix Immature balance bug (#361)
Browse files Browse the repository at this point in the history 
* move wallet balance from mempool to wallet

* remove immature outpoint state

* cache the balance

* move getUTXOs to wallet

* Remove .Immature OutpointState from tests

* test: just warn in case of missing txs in network mock

* fix: a couple of fix in getBalance and getUTXOs

* test: add immature balace integration test

* test: remove old getbalance test and update getUTXOs test

* move balance calculation back to mempool

* jsdoc
  • Loading branch information
@panleone
panleone authored May 21, 2024
1 parent e54fbb5 commit 718323e
Show file tree
Hide file tree
Showing 7 changed files with 290 additions and 104 deletions.
4 changes: 2 additions & 2 deletions scripts/__mocks__/network.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -62,8 +62,8 @@ class TestNetwork {
'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()
);
}
Expand Down
200 changes: 145 additions & 55 deletions scripts/mempool.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
};

Expand All @@ -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
Expand Down Expand Up @@ -54,7 +56,7 @@ export class Mempool {
*/
setOutpointStatus(outpoint, status) {
this.#outpointStatus.set(outpoint.toUnique(), status);
this.#invalidateBalanceCache();
this.invalidateBalanceCache();
}

/**
Expand All @@ -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();
}

/**
Expand All @@ -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();
}

/**
Expand Down Expand Up @@ -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();
}

Expand All @@ -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
*/
16 changes: 15 additions & 1 deletion scripts/transaction.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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);
Expand Down
17 changes: 11 additions & 6 deletions scripts/wallet.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -792,6 +792,7 @@ export class Wallet {
.getUTXOs({
requirement: OutpointState.P2PKH | OutpointState.OURS,
target,
blockCount,
})
.map((u) => {
return {
Expand All @@ -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');
}
Expand Down Expand Up @@ -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');
Expand Down Expand Up @@ -1007,6 +1010,7 @@ export class Wallet {
const utxos = this.#mempool.getUTXOs({
requirement: requirement | OutpointState.OURS,
target: value,
blockCount,
});
transactionBuilder.addUTXOs(utxos);

Expand Down Expand Up @@ -1191,6 +1195,7 @@ export class Wallet {
return this.#mempool
.getUTXOs({
requirement: OutpointState.P2PKH | OutpointState.OURS,
blockCount,
})
.filter((u) => u.value === collateralValue);
}
Expand All @@ -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);
}

/**
Expand Down Expand Up @@ -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
Expand Down
Loading

0 comments on commit 718323e

Please sign in to comment.