blockchain.java

// Block Chain should maintain only limited block nodes to satisfy the functions
// You should not have all the blocks added to the block chain in memory 
// as it would cause a memory overflow.
import java.sql.Array;
import java.util.*;
import java.util.stream.Collectors;
import java.sql.Timestamp;

public class BlockChain {
    public static final int CUT_OFF_AGE = 10;
    private TransactionPool txPool;
    private HashMap<ByteArrayWrapper, BlockInfo> addedBlocksByHash;
    private BlockInfo maxHeightBlockInfo;

    private class BlockInfo {
        public Block block;
        public Timestamp timestamp;
        public Integer height;
        public UTXOPool utxoPool;

        BlockInfo(Block block, Timestamp timestamp, Integer height, UTXOPool utxoPool) {
            this.block = block;
            this.timestamp = timestamp;
            this.height = height;
            this.utxoPool = utxoPool;
        }
    }

    /**
     * create an empty block chain with just a genesis block. Assume {@code genesisBlock} is a valid
     * block
     */
    public BlockChain(Block genesisBlock) {
        this.txPool = new TransactionPool();
        this.addedBlocksByHash = new HashMap<>();
        UTXOPool utxoPool = new UTXOPool();
        addCoinbaseOutputsToUTXOPool(genesisBlock, utxoPool);
        BlockInfo genesisBlockInfo = new BlockInfo(genesisBlock, new Timestamp(System.currentTimeMillis()), 1, utxoPool);
        this.addedBlocksByHash.put(new ByteArrayWrapper(genesisBlock.getHash()), genesisBlockInfo);
        this.maxHeightBlockInfo = genesisBlockInfo;
    }

    /** Get the maximum height block */
    public Block getMaxHeightBlock() {
        return maxHeightBlockInfo.block;
    }

    /** Get the UTXOPool for mining a new block on top of max height block */
    public UTXOPool getMaxHeightUTXOPool() {
        return maxHeightBlockInfo.utxoPool;
    }

    /** Get the transaction pool to mine a new block */
    public TransactionPool getTransactionPool() {
        return txPool;
    }

    /**
     * Add {@code block} to the block chain if it is valid. For validity, all transactions should be
     * valid and block should be at {@code height > (maxHeight - CUT_OFF_AGE)}.
     * 
     * <p>
     * For example, you can try creating a new block over the genesis block (block height 2) if the
     * block chain height is {@code <=
     * CUT_OFF_AGE + 1}. As soon as {@code height > CUT_OFF_AGE + 1}, you cannot create a new block
     * at height 2.
     * 
     * @return true if block is successfully added
     */
    public boolean addBlock(Block block) {
        // Don't add genesis block
        byte[] prevBlockHash = block.getPrevBlockHash();
        if (prevBlockHash == null) {
            return false;
        }

        // Previous block must already exist in block chain
        BlockInfo prevBlockInfo = addedBlocksByHash.get(new ByteArrayWrapper(prevBlockHash));
        if (prevBlockInfo == null) {
            return false;
        }

        //
        Integer blockHeight = prevBlockInfo.height + 1;
        if(blockHeight <= maxHeightBlockInfo.height - CUT_OFF_AGE) {
            return false;
        }

        ArrayList<Transaction> txList = block.getTransactions();
        Transaction[] currentBlockTxs = new Transaction[txList.size()];
        currentBlockTxs = txList.toArray(currentBlockTxs);

        TxHandler txHandler = new TxHandler(prevBlockInfo.utxoPool);
        Transaction[] validTxs = txHandler.handleTxs(currentBlockTxs);
        Set<Transaction> validTxSet = Arrays.stream(validTxs).collect(Collectors.toCollection(HashSet::new));
        Set<Transaction> currentBlockTxSet = Arrays.stream(currentBlockTxs).collect(Collectors.toCollection(HashSet::new));
        if (!validTxSet.equals(currentBlockTxSet)) {
            return false;
        }

        UTXOPool utxoPool = new UTXOPool(txHandler.getUTXOPool());
        addCoinbaseOutputsToUTXOPool(block, utxoPool);
        BlockInfo blockInfo = new BlockInfo(block, new Timestamp(System.currentTimeMillis()),blockHeight, utxoPool);
        addedBlocksByHash.put(new ByteArrayWrapper(block.getHash()),blockInfo);

        if(blockInfo.height > maxHeightBlockInfo.height) {
            maxHeightBlockInfo = blockInfo;
        } else if (blockInfo.height == maxHeightBlockInfo.height && blockInfo.timestamp.compareTo(maxHeightBlockInfo.timestamp) < 0) {
            maxHeightBlockInfo = blockInfo;
        }

        for (Transaction tx: txList) {
            txPool.removeTransaction(tx.getHash());
        }

        return true;
    }

    /** Add a transaction to the transaction pool */
    public void addTransaction(Transaction tx) {
        txPool.addTransaction(tx);
    }

    private static void addCoinbaseOutputsToUTXOPool(Block block, UTXOPool utxoPool) {
        Transaction coinbaseTx = block.getCoinbase();
        byte[] txHash = coinbaseTx.getHash();
        for(int i = 0; i < coinbaseTx.numOutputs(); i++) {
            UTXO utxo = new UTXO(txHash, i);
            utxoPool.addUTXO(utxo, coinbaseTx.getOutput(i));
        }
    }
}