diff --git a/main.go b/main.go
index 11338f21..4c0c41db 100644
--- a/main.go
+++ b/main.go
@@ -17,9 +17,11 @@ import (
 
 var (
 	// Key for the sub-store with Ethereum accounts
-	AccountsKey= types.NewKVStoreKey("account")
+	AccountsKey = types.NewKVStoreKey("account")
 	// Key for the sub-store with storage data of Ethereum contracts
 	StorageKey = types.NewKVStoreKey("storage")
+	// Key for the sub-store with the code for contracts
+	CodeKey = types.NewKVStoreKey("code")
 )
 
 // This is what stored in the lookupDb
@@ -34,42 +36,58 @@ type OurDatabase struct {
 	lookupDb          dbm.DB // Maping [trie_root_hash] => <version_id>.
 	                         // This mapping exists so that we can implement OpenTrie and OpenStorageTrie functions 
 	                         // of the state.Database interface
+	codeDb            dbm.DB // Mapping [codeHash] -> <code>
 	addrPreimageDb    dbm.DB // Mapping [contract_address_hash] -> <contract_address>
 	cdc               *amino.Codec // Amino codec to encode the values forthe lookupDb
 }
 
-func OurNewDatabase(stateDb, lookupDb dbm.DB) *OurDatabase {
+func OurNewDatabase(stateDb, lookupDb, addrPreimageDb, codeDb dbm.DB) (*OurDatabase, error) {
 	od := &OurDatabase{}
 	od.stateStore = store.NewCommitMultiStore(stateDb)
 	od.stateStore.MountStoreWithDB(AccountsKey, types.StoreTypeIAVL, nil)
 	od.stateStore.MountStoreWithDB(StorageKey, types.StoreTypeIAVL, nil)
+	if err := od.stateStore.LoadLatestVersion(); err != nil {
+		return nil, err
+	}
 	od.lookupDb = lookupDb
+	od.addrPreimageDb = addrPreimageDb
+	od.codeDb = codeDb
 	od.cdc = amino.NewCodec()
-	return od
+	return od, nil
 }
 
 func (od *OurDatabase) OpenTrie(root eth_common.Hash) (eth_state.Trie, error) {
 	// Look up version id to use
-	val := od.lookupDb.Get(root[:])
-	var versionId int64
-	_, err := od.cdc.UnmarshalBinaryReader(bytes.NewBuffer(val), &versionId, 0)
-	if err != nil {
-		return nil, err
+	if root != (eth_common.Hash{}) {
+		val := od.lookupDb.Get(root[:])
+		if val == nil {
+			return nil, fmt.Errorf("Could not find version with root hash %x", root[:])
+		}
+		var versionId int64
+		_, err := od.cdc.UnmarshalBinaryReader(bytes.NewBuffer(val), &versionId, 0)
+		if err != nil {
+			return nil, err
+		}
+		od.stateStore.LoadVersion(versionId)
 	}
-	od.stateStore.LoadVersion(versionId)
 	st := od.stateStore.GetCommitKVStore(AccountsKey)
 	return &OurTrie{od: od, st: st, prefix: nil}, nil
 }
 
 func (od *OurDatabase) OpenStorageTrie(addrHash, root eth_common.Hash) (eth_state.Trie, error) {
-	val := od.lookupDb.Get(root[:])
-	var versionId int64
-	_, err := od.cdc.UnmarshalBinaryReader(bytes.NewBuffer(val), &versionId, 0)
-	if err != nil {
-		return nil, err
+	if root != (eth_common.Hash{}) {
+		val := od.lookupDb.Get(root[:])
+		if val == nil {
+			return nil, fmt.Errorf("Could not find version with root hash %x", root[:])
+		}
+		var versionId int64
+		_, err := od.cdc.UnmarshalBinaryReader(bytes.NewBuffer(val), &versionId, 0)
+		if err != nil {
+			return nil, err
+		}
+		od.stateStore.LoadVersion(versionId)     // This might not be required,
+		                                        // we just need to check that accounts and storage are consistent
 	}
-	od.stateStore.LoadVersion(versionId)     // This might not be required,
-	                                        // we just need to check that accounts and storage are consistent
 	st := od.stateStore.GetCommitKVStore(StorageKey)
 	return &OurTrie{od:od, st: st, prefix: addrHash[:]}, nil
 }
@@ -79,11 +97,13 @@ func (od *OurDatabase) CopyTrie(eth_state.Trie) eth_state.Trie {
 }
 
 func (od *OurDatabase) ContractCode(addrHash, codeHash eth_common.Hash) ([]byte, error) {
-	return nil, nil
+	code := od.codeDb.Get(codeHash[:])
+	return code, nil
 }
 
 func (od *OurDatabase) ContractCodeSize(addrHash, codeHash eth_common.Hash) (int, error) {
-	return 0, nil
+	code := od.codeDb.Get(codeHash[:])
+	return len(code), nil
 }
 
 func (od *OurDatabase) TrieDB() *eth_trie.Database {
@@ -162,7 +182,21 @@ func main() {
 	fmt.Printf("Instantiating state.Database\n")
 	stateDb := dbm.NewDB("state" /* name */, dbm.MemDBBackend, "" /* dir */)
 	lookupDb := dbm.NewDB("lookup" /* name */, dbm.MemDBBackend, "" /* dir */)
+	addrPreimageDb := dbm.NewDB("addrPreimage" /* name */, dbm.MemDBBackend, "" /* dir */)
+	codeDb := dbm.NewDB("code" /* name */, dbm.MemDBBackend, "" /* dir */)
 	var d eth_state.Database
-	d = OurNewDatabase(stateDb, lookupDb)
-	d.OpenTrie(eth_common.Hash{})
+	var err error
+	d, err = OurNewDatabase(stateDb, lookupDb, addrPreimageDb, codeDb)
+	if err != nil {
+		panic(err)
+	}
+	fmt.Printf("Instantiating state.StateDB\n")
+	// With empty root hash, i.e. empty state
+	statedb, err := eth_state.New(eth_common.Hash{}, d)
+	if err != nil {
+		panic(err)
+	}
+	// Try something
+	b := statedb.GetBalance(eth_common.HexToAddress("0x829BD824B016326A401d083B33D092293333A830"))
+	fmt.Printf("Balance: %s\n", b)
 }
\ No newline at end of file