Require block stores to track the best chain head, add for the MemoryBlockStore.

src/com/google/bitcoin/core/BlockChain.java

@@ -69,16 +69,12 @@ public class BlockChain {
 
     public BlockChain(NetworkParameters params, Wallet wallet) {
         // TODO: Let the user pass in a BlockStore object so they can choose how to store the headers.
-        blockStore = new MemoryBlockStore();
         try {
-            // Set up the genesis block. When we start out fresh, it is by definition the top of the chain.
-            Block genesisHeader = params.genesisBlock.cloneAsHeader();
-            chainHead = new StoredBlock(genesisHeader, genesisHeader.getWork(), 0);
-            blockStore.put(chainHead);
+            blockStore = new MemoryBlockStore(params);
+            chainHead = blockStore.getChainHead();
+            LOG("chain head is: " + chainHead.header.toString());
         } catch (BlockStoreException e) {
-            // Cannot happen.
-        } catch (VerificationException e) {
-            // Genesis block always verifies.
+            throw new RuntimeException(e);
         }
 
         this.params = params;
@@ -146,7 +142,7 @@ public class BlockChain {
             blockStore.put(newStoredBlock);
             if (storedPrev.equals(chainHead)) {
                 // This block connects to the best known block, it is a normal continuation of the system.
-                chainHead = newStoredBlock;
+                setChainHead(newStoredBlock);
                 LOG("Received new block, chain is now " + chainHead.height + " blocks high");
             } else {
                 // This block connects to somewhere other than the top of the chain.
@@ -154,7 +150,7 @@ public class BlockChain {
                     // This chain has overtaken the one we currently believe is best. Reorganize is required.
                     wallet.reorganize(chainHead, newStoredBlock);
                     // Update the pointer to the best known block.
-                    chainHead = newStoredBlock;
+                    setChainHead(newStoredBlock);
                 } else {
                     LOG("Received a block which forks the chain, but it did not cause a reorganize.");
                 }
@@ -167,6 +163,15 @@ public class BlockChain {
         return true;
     }
 
+    private void setChainHead(StoredBlock chainHead) {
+        this.chainHead = chainHead;
+        try {
+            blockStore.setChainHead(chainHead);
+        } catch (BlockStoreException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
     /**
      * Calculates the additional fields a StoredBlock holds given the previous block in the chain and the new block.
      */

src/com/google/bitcoin/core/BlockStore.java

@@ -39,4 +39,14 @@ interface BlockStore {
      * parameter. If no such block is found, returns null.
      */
     StoredBlock get(byte[] hash) throws BlockStoreException;
+
+    /**
+     * Returns the {@link StoredBlock} that represents the top of the chain of greatest total work.
+     */
+    StoredBlock getChainHead() throws BlockStoreException;
+
+    /**
+     * Sets the {@link StoredBlock} that represents the top of the chain of greatest total work.
+     */
+    void setChainHead(StoredBlock chainHead) throws BlockStoreException;
 }

src/com/google/bitcoin/core/BlockStoreException.java

@@ -20,4 +20,7 @@ package com.google.bitcoin.core;
  * Thrown when something goes wrong with storing a block. Examples: out of disk space.
  */
 public class BlockStoreException extends Exception {
+    public BlockStoreException(Throwable t) {
+        super(t);
+    }
 }

src/com/google/bitcoin/core/MemoryBlockStore.java

@@ -16,6 +16,7 @@
 
 package com.google.bitcoin.core;
 
+import java.io.*;
 import java.nio.ByteBuffer;
 import java.util.HashMap;
 import java.util.Map;
@@ -27,18 +28,60 @@ class MemoryBlockStore implements BlockStore {
     // We use a ByteBuffer to hold hashes here because the Java array equals()/hashcode() methods do not operate on
     // the contents of the array but just inherit the default Object behavior. ByteBuffer provides the functionality
     // needed to act as a key in a map.
-    private Map<ByteBuffer, StoredBlock> blockMap;
+    //
+    // The StoredBlocks are also stored as serialized objects to ensure we don't have assumptions that would make
+    // things harder for disk based implementations.
+    private Map<ByteBuffer, byte[]> blockMap;
+    private StoredBlock chainHead;
 
-    MemoryBlockStore() {
-        blockMap = new HashMap<ByteBuffer, StoredBlock>();
+    MemoryBlockStore(NetworkParameters params) {
+        blockMap = new HashMap<ByteBuffer, byte[]>();
+        // Insert the genesis block.
+        try {
+            Block genesisHeader = params.genesisBlock.cloneAsHeader();
+            StoredBlock storedGenesis = new StoredBlock(genesisHeader, genesisHeader.getWork(), 0);
+            put(storedGenesis);
+            setChainHead(storedGenesis);
+        } catch (BlockStoreException e) {
+            throw new RuntimeException(e);  // Cannot happen.
+        } catch (VerificationException e) {
+            throw new RuntimeException(e);  // Cannot happen.
+        }
     }
 
     public synchronized void put(StoredBlock block) throws BlockStoreException {
         byte[] hash = block.header.getHash();
-        blockMap.put(ByteBuffer.wrap(hash), block);
+        ByteArrayOutputStream bos = new ByteArrayOutputStream();
+        try {
+            ObjectOutputStream oos = new ObjectOutputStream(bos);
+            oos.writeObject(block);
+            oos.close();
+            blockMap.put(ByteBuffer.wrap(hash), bos.toByteArray());
+        } catch (IOException e) {
+            throw new BlockStoreException(e);
+        }
     }
 
     public synchronized StoredBlock get(byte[] hash) throws BlockStoreException {
-        return blockMap.get(ByteBuffer.wrap(hash));
+        try {
+            byte[] serializedBlock = blockMap.get(ByteBuffer.wrap(hash));
+            if (serializedBlock == null)
+                return null;
+            ObjectInputStream ois = new ObjectInputStream(new ByteArrayInputStream(serializedBlock));
+            StoredBlock storedBlock = (StoredBlock) ois.readObject();
+            return storedBlock;
+        } catch (IOException e) {
+            throw new BlockStoreException(e);
+        } catch (ClassNotFoundException e) {
+            throw new BlockStoreException(e);
+        }
+    }
+
+    public StoredBlock getChainHead() {
+        return chainHead;
+    }
+
+    public void setChainHead(StoredBlock chainHead) throws BlockStoreException {
+        this.chainHead = chainHead;
     }
 }