Merge branch '87-feat-blockchain-add-re-org-and-tx-dropping-capabilities-to-blockchaininmemory' into 'development'

Resolve "feat(Blockchain): Add re-org and tx dropping capabilities to BlockchainInMemory"

Closes #87

See merge request GeneralProtocols/xo/stack!123

packages/stack/src/blockchain/blockchain-inmemory.ts

@@ -6,6 +6,8 @@ import {
   type BlockchainUTXO,
   type BlockchainUTXOs,
   type BlockchainChainTip,
+  AddressStatusPayload,
+  BlockHeightPayload,
 } from './types.js';
 
 import {
@@ -15,6 +17,7 @@ import {
   Outpoint,
   Transaction,
   UTXO,
+  Bytes,
 } from '@xocash/primitives';
 
 import {
@@ -26,7 +29,10 @@ import {
 /**
  * @remarks This provider DOES NOT enforce the 100 confirmations Coinbase Transaction rule.
  **/
-export class BlockchainInMemory extends BaseBlockchain {
+export class BlockchainInMemory extends BaseBlockchain<
+  AddressStatusPayload,
+  BlockHeightPayload
+> {
   static async from<T extends typeof BlockchainInMemory>(this: T) {
     const blockchain = new this();
     await blockchain.start();
@@ -35,7 +41,10 @@ export class BlockchainInMemory extends BaseBlockchain {
 
   private blockHeaders: Array<BlockHeader>;
   private transactions: BlockchainTransactions;
-  private addressSubscriptions: Map<string, Set<BlockchainCallback>>;
+  private addressSubscriptions: Map<
+    string,
+    Set<BlockchainCallback<AddressStatusPayload>>
+  >;
   private transactionSubscriptions: Map<string, Set<BlockchainCallback>>;
 
   constructor() {
@@ -78,6 +87,8 @@ export class BlockchainInMemory extends BaseBlockchain {
   }
 
   async fetchAddressHistory(address: string) {
+    const lockingBytecode = Address.from(address).toLockscriptHex();
+
     return this.transactions.filter((tx) => {
       // Check if any input or output involves the address
       // This is a simplified check and might need to be more sophisticated in a real implementation
@@ -85,17 +96,19 @@ export class BlockchainInMemory extends BaseBlockchain {
         tx.transaction
           .getInputs()
           .some((input) =>
-            binToHex(input.unlockingBytecode).includes(address),
+            binToHex(input.unlockingBytecode).includes(lockingBytecode),
           ) ||
         tx.transaction
           .getOutputs()
-          .some((output) => binToHex(output.lockingBytecode).includes(address))
+          .some((output) =>
+            binToHex(output.lockingBytecode).includes(lockingBytecode),
+          )
       );
     });
   }
 
   async fetchBlockHeader(blockHeight: number) {
-    const blockHeader = this.blockHeaders[blockHeight];
+    const blockHeader = this.blockHeaders[blockHeight - 1];
 
     if (!blockHeader) {
       throw new Error(`No block available for height ${blockHeight}`);
@@ -164,7 +177,7 @@ export class BlockchainInMemory extends BaseBlockchain {
 
   async subscribeLockscript(
     lockscript: string | Uint8Array,
-    callback: BlockchainCallback,
+    callback: BlockchainCallback<AddressStatusPayload>,
   ) {
     // Fulcrum does not support subscribing to lockscripts, so we cast it to an address.
     const address = Address.fromLockscript(lockscript);
@@ -178,7 +191,7 @@ export class BlockchainInMemory extends BaseBlockchain {
 
   async unsubscribeLockscript(
     lockscript: string | Uint8Array,
-    callback: BlockchainCallback,
+    callback: BlockchainCallback<AddressStatusPayload>,
   ): Promise<boolean> {
     // Fulcrum does not support subscribing to lockscripts, so we cast it to an address.
     const address = Address.fromLockscript(lockscript);
@@ -190,7 +203,10 @@ export class BlockchainInMemory extends BaseBlockchain {
     return await this.unsubscribeAddress(address.toCashAddr(), callback);
   }
 
-  async subscribeAddress(address: string, callback: BlockchainCallback) {
+  async subscribeAddress(
+    address: string,
+    callback: BlockchainCallback<AddressStatusPayload>,
+  ) {
     // Convert the address to it's lockscript hex so that it's consistently represented.
     const lockscriptHex = Address.from(address).toLockscriptHex();
 
@@ -206,7 +222,10 @@ export class BlockchainInMemory extends BaseBlockchain {
     return async () => await this.unsubscribeAddress(address, callback);
   }
 
-  async unsubscribeAddress(address: string, callback: BlockchainCallback) {
+  async unsubscribeAddress(
+    address: string,
+    callback: BlockchainCallback<AddressStatusPayload>,
+  ) {
     // Convert the address to it's lockscript hex so that it's consistently represented.
     const lockscriptHex = Address.from(address).toLockscriptHex();
 
@@ -263,29 +282,32 @@ export class BlockchainInMemory extends BaseBlockchain {
   }
 
   private notifyAddressSubscribers(tx: Transaction) {
-    // Derive the UTXO state so that we can lookup source outputs.
+    const payload = {
-    const utxoState = this.deriveUtxoSet();
+      type: 'address-status',
+      // 32 random characters. Since we only need this to be unique per transaction, we can just randomise the first 16 bytes and pad the rest with 0s.
+      status: Math.random().toString(36).substring(2).padEnd(32, '0'),
+    } as const;
 
     // First, notify all source output listeners.
     tx.getInputs().forEach((input) => {
-      const sourceOutput = utxoState.get(
+      // Skip coinbase-like inputs
-        `${input.outpointTransactionHash};${input.outpointIndex}`,
+      if (input.outpointTransactionHash.length === 0) return;
-      );
 
-      // If there is no corresponding source output, skip this notification.
+      const outpointTxHashHex = binToHex(input.outpointTransactionHash);
-      // NOTE: This can occur if skipValidation is set to true in broadcast (which would be used for "coinbase" transactions).
+      const sourceTx = this.transactions.get(outpointTxHashHex)?.transaction;
-      if (!sourceOutput) {
+      if (!sourceTx) return;
-        return;
+
-      }
+      const sourceOutput = sourceTx.getOutputs()[input.outpointIndex];
+      if (!sourceOutput) return;
 
       const address = Address.fromLockscriptBytes(
-        sourceOutput.utxo.output.lockingBytecode,
+        sourceOutput.lockingBytecode,
       ).toLockscriptHex();
 
       const callbacks = this.addressSubscriptions.get(address);
 
       if (callbacks) {
-        callbacks.forEach((callback) => callback(address));
+        callbacks.forEach((callback) => callback(address, payload));
       }
     });
 
@@ -298,7 +320,7 @@ export class BlockchainInMemory extends BaseBlockchain {
       const callbacks = this.addressSubscriptions.get(address);
 
       if (callbacks) {
-        callbacks.forEach((callback) => callback(address));
+        callbacks.forEach((callback) => callback(address, payload));
       }
     });
   }
@@ -351,16 +373,26 @@ export class BlockchainInMemory extends BaseBlockchain {
     return utxoMap;
   }
 
-  async mineBlock(outputs: Array<Output>): Promise<void> {
+  async mineBlock(
+    outputs: Array<Output>,
+    includeTransactionHashes?: Array<string>,
+  ): Promise<void> {
     // Create a transaction that pays to our subscribed address.
     const transaction = Transaction.fromRaw({
       version: 2,
       locktime: 0,
       inputs: [
+        // There are a few requirements for a coinbase tx: https://learnmeabitcoin.com/technical/mining/coinbase-transaction/
         {
           outpointIndex: 0,
-          outpointTransactionHash: new Uint8Array(),
+          // The TxID for a coinbase tx input should be all zeros
-          unlockingBytecode: new Uint8Array(),
+          outpointTransactionHash: Bytes.fromHex(
+            '0000000000000000000000000000000000000000000000000000000000000000',
+          ),
+          // Unlocking bytecode for a coinbase tx should start with the block height (this also helps us not have txid collisions when hashing this coinbase tx)
+          unlockingBytecode: Bytes.fromUtf8(
+            `${this.blockHeaders.length + 1}-coinbase-memory-blockchain`,
+          ),
           sequenceNumber: 0,
         },
       ],
@@ -384,9 +416,123 @@ export class BlockchainInMemory extends BaseBlockchain {
 
     // Add it to our list of block headers.
     this.blockHeaders.push(newBlockHeader);
+    const newHeight = this.blockHeaders.length;
 
     // Broadcast that transaction.
     await this.broadcastTransaction(transaction, { skipValidation: true });
+
+    // Confirm the coinbase transaction in this new block height.
+    const coinbaseHex = transaction.toHashHex();
+    const coinbaseEntry = this.transactions.get(coinbaseHex);
+    if (coinbaseEntry) {
+      coinbaseEntry.height = newHeight;
+      this.transactions.set(coinbaseHex, coinbaseEntry);
+    }
+
+    // Select which mempool txs to include. By default include all height <= 0.
+    const toInclude: Array<string> = Array.isArray(includeTransactionHashes)
+      ? includeTransactionHashes.filter((h) => {
+          const entry = this.transactions.get(h);
+          return !!entry && entry.height <= 0 && h !== coinbaseHex;
+        })
+      : (Array.from(this.transactions.values()) as BlockchainTransaction[])
+          .filter(
+            (entry: BlockchainTransaction) =>
+              entry.height <= 0 && entry.hash.toHex() !== coinbaseHex,
+          )
+          .map((entry: BlockchainTransaction) => entry.hash.toHex());
+
+    // Mark selected txs as confirmed in this block.
+    for (const h of toInclude) {
+      const entry = this.transactions.get(h);
+      if (entry) {
+        entry.height = newHeight;
+        this.transactions.set(h, entry);
+
+        // Notify the address subscribers
+        this.notifyAddressSubscribers(entry.transaction);
+        this.notifyTransactionSubscribers(entry.transaction);
+      }
+    }
+
+    this.emit('blockHeightUpdated', newHeight);
+  }
+
+  /**
+   * Replace only the last N blocks (N = blocks.length) of the chain with the
+   * provided blocks. Transactions confirmed in the replaced blocks are first
+   * unconfirmed (set to height 0). Transactions listed in the new blocks are
+   * then confirmed at their new heights. Earlier blocks are preserved.
+   */
+  replaceBlocks(
+    blocks: Array<{
+      header: BlockHeader;
+      transactions: Array<Transaction | Uint8Array | string>;
+    }>,
+    options: Partial<{
+      dropMempool: boolean;
+    }> = {},
+  ): void {
+    const { dropMempool } = {
+      dropMempool: false,
+      ...options,
+    };
+
+    if (blocks.length === 0) return;
+
+    const currentHeight = this.blockHeaders.length;
+    const preserved = Math.max(0, currentHeight - blocks.length);
+
+    // 1) Unconfirm any transactions that were in the range being replaced
+    this.transactions.forEach((entry) => {
+      if (entry.height > preserved) {
+        entry.height = 0;
+        this.transactions.set(entry.hash.toHex(), entry);
+      }
+    });
+
+    // 2) Replace headers: keep older headers, then append new headers
+    this.blockHeaders = this.blockHeaders.slice(0, preserved);
+    for (let i = 0; i < blocks.length; i++) {
+      this.blockHeaders.push(blocks[i].header);
+    }
+
+    // 3) Confirm the transactions in the new blocks at their new heights
+    for (let i = 0; i < blocks.length; i++) {
+      const height = preserved + i + 1;
+      for (const t of blocks[i].transactions) {
+        const tx = t instanceof Transaction ? t : Transaction.from(t as any);
+        const txHashHex = tx.toHashHex();
+
+        // Ensure the tx is tracked
+        if (!this.transactions.get(txHashHex)) {
+          this.addTransaction({
+            hash: tx.toHash(),
+            transaction: tx,
+            height: 0,
+          });
+        }
+
+        const entry = this.transactions.get(txHashHex)!;
+        entry.height = height;
+        this.transactions.set(txHashHex, entry);
+
+        // Notify the address subscribers
+        this.notifyAddressSubscribers(tx);
+        this.notifyTransactionSubscribers(tx);
+      }
+    }
+
+    // If set, delete any transactions that are not in the new blocks
+    if (dropMempool) {
+      for (const tx of this.transactions.values()) {
+        if (tx.height <= 0) {
+          this.transactions.delete(tx.hash.toHex());
+        }
+      }
+    }
+
+    this.emit('blockHeightUpdated', this.blockHeaders.length);
   }
 
   /**
@@ -409,4 +555,16 @@ export class BlockchainInMemory extends BaseBlockchain {
       throw new Error(virtualMachineResult);
     }
   }
+
+  /**
+   * Convert the in-memory blockchain transactions to an array of blocks.
+   */
+  public toBlocks() {
+    return this.blockHeaders.map((header, index) => ({
+      header,
+      transactions: Array.from(this.transactions.values()).filter(
+        (tx) => tx.height === index + 1,
+      ),
+    }));
+  }
 }

packages/stack/src/wallet/wallet-p2pkh.ts

@@ -24,6 +24,7 @@ import {
   binToHex,
   walletTemplateP2pkhNonHd,
   walletTemplateToCompilerBCH,
+  Output,
 } from '@bitauth/libauth';
 
 export type WalletP2PKHGenesisData = {
@@ -228,9 +229,16 @@ export class WalletP2PKH extends BaseWallet {
     const getSourceOutputs = async (tx: BlockchainTransaction) => {
       const sourceOutputPromises = tx.transaction
         .getInputs()
-        .map((input) => getSourceOutput(input));
+        // Coinbase transactions do not have source output, so we return null
+        .map((input) =>
+          getSourceOutput(input).catch(() => {
+            return undefined;
+          }),
+        );
 
-      return await Promise.all(sourceOutputPromises);
+      return (await Promise.all(sourceOutputPromises)).filter(
+        (output): output is Output => output !== undefined,
+      );
     };
 
     // Get the source outputs for each of the transactions.