greedy-cpfp-mempool.py

#!/usr/bin/python3

class CMemPoolTx:
    @property
    def feerate(self):
        return self.nFee / self.nSize

    @property
    def sum_feerate(self):
        return self.nSumFees / self.nSumSize

    def __init__(self, nFee, nSize, prevtxs):
        self.nFee = nFee
        self.nSize = nSize
        self.prevtxs = frozenset(prevtxs)

        self.nDepth = 0
        if self.prevtxs:
            self.nDepth = max(tx.nDepth+1 for tx in self.prevtxs)

        # Calculate total fees/size of all CPFP-eligible parents
        def greedy_sum_fees(tx, max_feerate, visited=None):
            if visited is None:
                visited = set()

            if tx in visited:
                return (0, 0)

            visited.add(tx)

            sum_fees = tx.nFee
            sum_size = tx.nSize
            for prevtx in tx.prevtxs:
                if prevtx.sum_feerate < max_feerate:
                    rf, rs = greedy_sum_fees(prevtx, max_feerate, visited)
                    sum_fees += rf
                    sum_size += rs
            return (sum_fees, sum_size)

        (self.nSumFees, self.nSumSize) = greedy_sum_fees(self, self.feerate)


class CMemPool:
    def __init__(self):
        self.txs = {}
        self.leaves = set()
        self.nTotalFees = 0
        self.nTotalSize = 0

    def add(self, tx):
        self.txs[tx] = set()
        self.leaves.add(tx)

        # add ourself to prevtxs children
        for prevtx in tx.prevtxs:
            self.txs[prevtx].add(tx)
            self.leaves.discard(prevtx)

        self.nTotalFees += tx.nFee
        self.nTotalSize += tx.nSize

        return tx

    def remove_leaf(self, tx):
        self.nTotalFees -= tx.nFee
        self.nTotalSize -= tx.nSize

        # delete ourselves from prevtxs
        for prevtx in tx.prevtxs:
            self.txs[prevtx].discard(tx)
            if not self.txs[prevtx]:
                self.leaves.add(prevtx)

        del self.txs[tx]
        self.leaves.remove(tx)

    def make_space(self, nSize, max_fees):
        removed_txs = set()

        removed_fees = 0
        removed_size = 0

        while removed_size < nSize and pool.leaves:
            rtx = sorted(pool.leaves, key=lambda tx:tx.sum_feerate)[0]
            if removed_fees + rtx.nFee > max_fees:
                break
            removed_txs.add(rtx)
            removed_fees += rtx.nFee
            removed_size += rtx.nSize
            self.remove_leaf(rtx)

        return (removed_txs, removed_fees, removed_size)



pool = CMemPool()

tx0 = pool.add(CMemPoolTx(1,1,[]))

tx1 = pool.add(CMemPoolTx(1,1,[tx0]))
tx2 = pool.add(CMemPoolTx(2,1,[tx1, tx0]))

tx1 = pool.add(CMemPoolTx(10,1,[tx0]))
tx2 = pool.add(CMemPoolTx(200,1,[tx1]))

tx0 = pool.add(CMemPoolTx(5,1,[]))
tx1 = pool.add(CMemPoolTx(7,1,[tx0]))

for tx in sorted(pool.txs.keys(), key=lambda t:t.sum_feerate):
    print(tx.nDepth, tx.sum_feerate, (tx.nFee, tx.nSize), (tx.nSumFees, tx.nSumSize))

print(pool.nTotalFees, pool.nTotalSize)

(rtxs, rfees, rsize) = pool.make_space(3, 100)
print('removed fees/size = %f %f' % (rfees, rsize))

print(pool.nTotalFees, pool.nTotalSize)

for tx in sorted(pool.txs.keys(), key=lambda t:t.sum_feerate):
    print(tx.nDepth, tx.sum_feerate, (tx.nFee, tx.nSize), (tx.nSumFees, tx.nSumSize))