core/state: track all accounts in canon state

This change introduces a global, per-state cache that keeps account data
in the canon state. Thanks to @karalabe for lots of fixes.

6 files changed, 367 insertions, 295 deletions

diff --git a/core/state/dump.go b/core/state/dump.go
index a328b0537..58ecd852b 100644
--- a/core/state/dump.go
+++ b/core/state/dump.go
@@ -21,9 +21,10 @@ import (
 	"fmt"
 
 	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/rlp"
 )
 
-type Account struct {
+type DumpAccount struct {
 	Balance  string            `json:"balance"`
 	Nonce    uint64            `json:"nonce"`
 	Root     string            `json:"root"`
@@ -32,40 +33,41 @@ type Account struct {
 	Storage  map[string]string `json:"storage"`
 }
 
-type World struct {
-	Root     string             `json:"root"`
-	Accounts map[string]Account `json:"accounts"`
+type Dump struct {
+	Root     string                 `json:"root"`
+	Accounts map[string]DumpAccount `json:"accounts"`
 }
 
-func (self *StateDB) RawDump() World {
-	world := World{
+func (self *StateDB) RawDump() Dump {
+	dump := Dump{
 		Root:     common.Bytes2Hex(self.trie.Root()),
-		Accounts: make(map[string]Account),
+		Accounts: make(map[string]DumpAccount),
 	}
 
 	it := self.trie.Iterator()
 	for it.Next() {
 		addr := self.trie.GetKey(it.Key)
-		stateObject, err := DecodeObject(common.BytesToAddress(addr), self.db, it.Value)
-		if err != nil {
+		var data Account
+		if err := rlp.DecodeBytes(it.Value, &data); err != nil {
 			panic(err)
 		}
 
-		account := Account{
-			Balance:  stateObject.balance.String(),
-			Nonce:    stateObject.nonce,
-			Root:     common.Bytes2Hex(stateObject.Root()),
-			CodeHash: common.Bytes2Hex(stateObject.codeHash),
-			Code:     common.Bytes2Hex(stateObject.Code()),
+		obj := NewObject(common.BytesToAddress(addr), data, nil)
+		account := DumpAccount{
+			Balance:  data.Balance.String(),
+			Nonce:    data.Nonce,
+			Root:     common.Bytes2Hex(data.Root[:]),
+			CodeHash: common.Bytes2Hex(data.CodeHash),
+			Code:     common.Bytes2Hex(obj.Code(self.db)),
 			Storage:  make(map[string]string),
 		}
-		storageIt := stateObject.trie.Iterator()
+		storageIt := obj.getTrie(self.db).Iterator()
 		for storageIt.Next() {
 			account.Storage[common.Bytes2Hex(self.trie.GetKey(storageIt.Key))] = common.Bytes2Hex(storageIt.Value)
 		}
-		world.Accounts[common.Bytes2Hex(addr)] = account
+		dump.Accounts[common.Bytes2Hex(addr)] = account
 	}
-	return world
+	return dump
 }
 
 func (self *StateDB) Dump() []byte {
@@ -76,12 +78,3 @@ func (self *StateDB) Dump() []byte {
 
 	return json
 }
-
-// Debug stuff
-func (self *StateObject) CreateOutputForDiff() {
-	fmt.Printf("%x %x %x %x\n", self.Address(), self.Root(), self.balance.Bytes(), self.nonce)
-	it := self.trie.Iterator()
-	for it.Next() {
-		fmt.Printf("%x %x\n", it.Key, it.Value)
-	}
-}
diff --git a/core/state/managed_state.go b/core/state/managed_state.go
index f8e2f2b87..ad73dc0dc 100644
--- a/core/state/managed_state.go
+++ b/core/state/managed_state.go
@@ -33,14 +33,14 @@ type ManagedState struct {
 
 	mu sync.RWMutex
 
-	accounts map[string]*account
+	accounts map[common.Address]*account
 }
 
 // ManagedState returns a new managed state with the statedb as it's backing layer
 func ManageState(statedb *StateDB) *ManagedState {
 	return &ManagedState{
 		StateDB:  statedb.Copy(),
-		accounts: make(map[string]*account),
+		accounts: make(map[common.Address]*account),
 	}
 }
 
@@ -103,7 +103,7 @@ func (ms *ManagedState) SetNonce(addr common.Address, nonce uint64) {
 	so := ms.GetOrNewStateObject(addr)
 	so.SetNonce(nonce)
 
-	ms.accounts[addr.Str()] = newAccount(so)
+	ms.accounts[addr] = newAccount(so)
 }
 
 // HasAccount returns whether the given address is managed or not
@@ -114,29 +114,28 @@ func (ms *ManagedState) HasAccount(addr common.Address) bool {
 }
 
 func (ms *ManagedState) hasAccount(addr common.Address) bool {
-	_, ok := ms.accounts[addr.Str()]
+	_, ok := ms.accounts[addr]
 	return ok
 }
 
 // populate the managed state
 func (ms *ManagedState) getAccount(addr common.Address) *account {
-	straddr := addr.Str()
-	if account, ok := ms.accounts[straddr]; !ok {
+	if account, ok := ms.accounts[addr]; !ok {
 		so := ms.GetOrNewStateObject(addr)
-		ms.accounts[straddr] = newAccount(so)
+		ms.accounts[addr] = newAccount(so)
 	} else {
 		// Always make sure the state account nonce isn't actually higher
 		// than the tracked one.
 		so := ms.StateDB.GetStateObject(addr)
-		if so != nil && uint64(len(account.nonces))+account.nstart < so.nonce {
-			ms.accounts[straddr] = newAccount(so)
+		if so != nil && uint64(len(account.nonces))+account.nstart < so.Nonce() {
+			ms.accounts[addr] = newAccount(so)
 		}
 
 	}
 
-	return ms.accounts[straddr]
+	return ms.accounts[addr]
 }
 
 func newAccount(so *StateObject) *account {
-	return &account{so, so.nonce, nil}
+	return &account{so, so.Nonce(), nil}
 }
diff --git a/core/state/managed_state_test.go b/core/state/managed_state_test.go
index 0b53a42c5..baa53428f 100644
--- a/core/state/managed_state_test.go
+++ b/core/state/managed_state_test.go
@@ -29,11 +29,12 @@ func create() (*ManagedState, *account) {
 	db, _ := ethdb.NewMemDatabase()
 	statedb, _ := New(common.Hash{}, db)
 	ms := ManageState(statedb)
-	so := &StateObject{address: addr, nonce: 100}
-	ms.StateDB.stateObjects[addr.Str()] = so
-	ms.accounts[addr.Str()] = newAccount(so)
+	so := &StateObject{address: addr}
+	so.SetNonce(100)
+	ms.StateDB.stateObjects[addr] = so
+	ms.accounts[addr] = newAccount(so)
 
-	return ms, ms.accounts[addr.Str()]
+	return ms, ms.accounts[addr]
 }
 
 func TestNewNonce(t *testing.T) {
@@ -92,7 +93,7 @@ func TestRemoteNonceChange(t *testing.T) {
 	account.nonces = append(account.nonces, nn...)
 	nonce := ms.NewNonce(addr)
 
-	ms.StateDB.stateObjects[addr.Str()].nonce = 200
+	ms.StateDB.stateObjects[addr].data.Nonce = 200
 	nonce = ms.NewNonce(addr)
 	if nonce != 200 {
 		t.Error("expected nonce after remote update to be", 201, "got", nonce)
@@ -100,7 +101,7 @@ func TestRemoteNonceChange(t *testing.T) {
 	ms.NewNonce(addr)
 	ms.NewNonce(addr)
 	ms.NewNonce(addr)
-	ms.StateDB.stateObjects[addr.Str()].nonce = 200
+	ms.StateDB.stateObjects[addr].data.Nonce = 200
 	nonce = ms.NewNonce(addr)
 	if nonce != 204 {
 		t.Error("expected nonce after remote update to be", 201, "got", nonce)
diff --git a/core/state/state_object.go b/core/state/state_object.go
index 9f0ce5b4b..3496008a6 100644
--- a/core/state/state_object.go
+++ b/core/state/state_object.go
@@ -57,108 +57,163 @@ func (self Storage) Copy() Storage {
 	return cpy
 }
 
+// StateObject represents an Ethereum account which is being modified.
+//
+// The usage pattern is as follows:
+// First you need to obtain a state object.
+// Account values can be accessed and modified through the object.
+// Finally, call CommitTrie to write the modified storage trie into a database.
 type StateObject struct {
-	db   trie.Database // State database for storing state changes
-	trie *trie.SecureTrie
-
-	// Address belonging to this account
-	address common.Address
-	// The balance of the account
-	balance *big.Int
-	// The nonce of the account
-	nonce uint64
-	// The code hash if code is present (i.e. a contract)
-	codeHash []byte
-	// The code for this account
-	code Code
-	// Cached storage (flushed when updated)
-	storage Storage
-
-	// Mark for deletion
+	address common.Address // Ethereum address of this account
+	data    Account
+
+	// DB error.
+	// State objects are used by the consensus core and VM which are
+	// unable to deal with database-level errors. Any error that occurs
+	// during a database read is memoized here and will eventually be returned
+	// by StateDB.Commit.
+	dbErr error
+
+	// Write caches.
+	trie    *trie.SecureTrie // storage trie, which becomes non-nil on first access
+	code    Code             // contract bytecode, which gets set when code is loaded
+	storage Storage          // Cached storage (flushed when updated)
+
+	// Cache flags.
 	// When an object is marked for deletion it will be delete from the trie
 	// during the "update" phase of the state transition
-	remove  bool
-	deleted bool
-	dirty   bool
+	dirtyCode bool // true if the code was updated
+	remove    bool
+	deleted   bool
+	onDirty   func(addr common.Address) // Callback method to mark a state object newly dirty
 }
 
-func NewStateObject(address common.Address, db trie.Database) *StateObject {
-	object := &StateObject{
-		db:       db,
-		address:  address,
-		balance:  new(big.Int),
-		dirty:    true,
-		codeHash: emptyCodeHash,
-		storage:  make(Storage),
-	}
-	object.trie, _ = trie.NewSecure(common.Hash{}, db)
-	return object
-}
+// Account is the Ethereum consensus representation of accounts.
+// These objects are stored in the main account trie.
+type Account struct {
+	Nonce    uint64
+	Balance  *big.Int
+	Root     common.Hash // merkle root of the storage trie
+	CodeHash []byte
 
-func (self *StateObject) MarkForDeletion() {
-	self.remove = true
-	self.dirty = true
+	codeSize *int
+}
 
-	if glog.V(logger.Core) {
-		glog.Infof("%x: #%d %v X\n", self.Address(), self.nonce, self.balance)
+// NewObject creates a state object.
+func NewObject(address common.Address, data Account, onDirty func(addr common.Address)) *StateObject {
+	if data.Balance == nil {
+		data.Balance = new(big.Int)
+	}
+	if data.CodeHash == nil {
+		data.CodeHash = emptyCodeHash
 	}
+	return &StateObject{address: address, data: data, storage: make(Storage), onDirty: onDirty}
 }
 
-func (c *StateObject) getAddr(addr common.Hash) common.Hash {
-	var ret []byte
-	rlp.DecodeBytes(c.trie.Get(addr[:]), &ret)
-	return common.BytesToHash(ret)
+// EncodeRLP implements rlp.Encoder.
+func (c *StateObject) EncodeRLP(w io.Writer) error {
+	return rlp.Encode(w, c.data)
 }
 
-func (c *StateObject) setAddr(addr, value common.Hash) {
-	v, err := rlp.EncodeToBytes(bytes.TrimLeft(value[:], "\x00"))
-	if err != nil {
-		// if RLPing failed we better panic and not fail silently. This would be considered a consensus issue
-		panic(err)
+// setError remembers the first non-nil error it is called with.
+func (self *StateObject) setError(err error) {
+	if self.dbErr == nil {
+		self.dbErr = err
 	}
-	c.trie.Update(addr[:], v)
 }
 
-func (self *StateObject) Storage() Storage {
-	return self.storage
+func (self *StateObject) MarkForDeletion() {
+	self.remove = true
+	if self.onDirty != nil {
+		self.onDirty(self.Address())
+		self.onDirty = nil
+	}
+	if glog.V(logger.Core) {
+		glog.Infof("%x: #%d %v X\n", self.Address(), self.Nonce(), self.Balance())
+	}
 }
 
-func (self *StateObject) GetState(key common.Hash) common.Hash {
-	value, exists := self.storage[key]
-	if !exists {
-		value = self.getAddr(key)
-		if (value != common.Hash{}) {
-			self.storage[key] = value
+func (c *StateObject) getTrie(db trie.Database) *trie.SecureTrie {
+	if c.trie == nil {
+		var err error
+		c.trie, err = trie.NewSecure(c.data.Root, db)
+		if err != nil {
+			c.trie, _ = trie.NewSecure(common.Hash{}, db)
+			c.setError(fmt.Errorf("can't create storage trie: %v", err))
 		}
 	}
+	return c.trie
+}
 
+// GetState returns a value in account storage.
+func (self *StateObject) GetState(db trie.Database, key common.Hash) common.Hash {
+	value, exists := self.storage[key]
+	if exists {
+		return value
+	}
+	// Load from DB in case it is missing.
+	tr := self.getTrie(db)
+	var ret []byte
+	rlp.DecodeBytes(tr.Get(key[:]), &ret)
+	value = common.BytesToHash(ret)
+	if (value != common.Hash{}) {
+		self.storage[key] = value
+	}
 	return value
 }
 
+// SetState updates a value in account storage.
 func (self *StateObject) SetState(key, value common.Hash) {
 	self.storage[key] = value
-	self.dirty = true
+	if self.onDirty != nil {
+		self.onDirty(self.Address())
+		self.onDirty = nil
+	}
 }
 
-// Update updates the current cached storage to the trie
-func (self *StateObject) Update() {
+// updateTrie writes cached storage modifications into the object's storage trie.
+func (self *StateObject) updateTrie(db trie.Database) {
+	tr := self.getTrie(db)
 	for key, value := range self.storage {
 		if (value == common.Hash{}) {
-			self.trie.Delete(key[:])
+			tr.Delete(key[:])
 			continue
 		}
-		self.setAddr(key, value)
+		// Encoding []byte cannot fail, ok to ignore the error.
+		v, _ := rlp.EncodeToBytes(bytes.TrimLeft(value[:], "\x00"))
+		tr.Update(key[:], v)
 	}
 }
 
+// UpdateRoot sets the trie root to the current root hash of
+func (self *StateObject) UpdateRoot(db trie.Database) {
+	self.updateTrie(db)
+	self.data.Root = self.trie.Hash()
+}
+
+// CommitTrie the storage trie of the object to dwb.
+// This updates the trie root.
+func (self *StateObject) CommitTrie(db trie.Database, dbw trie.DatabaseWriter) error {
+	self.updateTrie(db)
+	if self.dbErr != nil {
+		fmt.Println("dbErr:", self.dbErr)
+		return self.dbErr
+	}
+	root, err := self.trie.CommitTo(dbw)
+	if err == nil {
+		self.data.Root = root
+	}
+	return err
+}
+
 func (c *StateObject) AddBalance(amount *big.Int) {
 	if amount.Cmp(common.Big0) == 0 {
 		return
 	}
-	c.SetBalance(new(big.Int).Add(c.balance, amount))
+	c.SetBalance(new(big.Int).Add(c.Balance(), amount))
 
 	if glog.V(logger.Core) {
-		glog.Infof("%x: #%d %v (+ %v)\n", c.Address(), c.nonce, c.balance, amount)
+		glog.Infof("%x: #%d %v (+ %v)\n", c.Address(), c.Nonce(), c.Balance(), amount)
 	}
 }
 
@@ -166,37 +221,32 @@ func (c *StateObject) SubBalance(amount *big.Int) {
 	if amount.Cmp(common.Big0) == 0 {
 		return
 	}
-	c.SetBalance(new(big.Int).Sub(c.balance, amount))
+	c.SetBalance(new(big.Int).Sub(c.Balance(), amount))
 
 	if glog.V(logger.Core) {
-		glog.Infof("%x: #%d %v (- %v)\n", c.Address(), c.nonce, c.balance, amount)
+		glog.Infof("%x: #%d %v (- %v)\n", c.Address(), c.Nonce(), c.Balance(), amount)
 	}
 }
 
-func (c *StateObject) SetBalance(amount *big.Int) {
-	c.balance = amount
-	c.dirty = true
-}
-
-func (c *StateObject) St() Storage {
-	return c.storage
+func (self *StateObject) SetBalance(amount *big.Int) {
+	self.data.Balance = amount
+	if self.onDirty != nil {
+		self.onDirty(self.Address())
+		self.onDirty = nil
+	}
 }
 
 // Return the gas back to the origin. Used by the Virtual machine or Closures
 func (c *StateObject) ReturnGas(gas, price *big.Int) {}
 
-func (self *StateObject) Copy() *StateObject {
-	stateObject := NewStateObject(self.Address(), self.db)
-	stateObject.balance.Set(self.balance)
-	stateObject.codeHash = common.CopyBytes(self.codeHash)
-	stateObject.nonce = self.nonce
+func (self *StateObject) Copy(db trie.Database, onDirty func(addr common.Address)) *StateObject {
+	stateObject := NewObject(self.address, self.data, onDirty)
 	stateObject.trie = self.trie
 	stateObject.code = self.code
 	stateObject.storage = self.storage.Copy()
 	stateObject.remove = self.remove
-	stateObject.dirty = self.dirty
+	stateObject.dirtyCode = self.dirtyCode
 	stateObject.deleted = self.deleted
-
 	return stateObject
 }
 
@@ -204,40 +254,66 @@ func (self *StateObject) Copy() *StateObject {
 // Attribute accessors
 //
 
-func (self *StateObject) Balance() *big.Int {
-	return self.balance
-}
-
 // Returns the address of the contract/account
 func (c *StateObject) Address() common.Address {
 	return c.address
 }
 
-func (self *StateObject) Trie() *trie.SecureTrie {
-	return self.trie
-}
-
-func (self *StateObject) Root() []byte {
-	return self.trie.Root()
+// Code returns the contract code associated with this object, if any.
+func (self *StateObject) Code(db trie.Database) []byte {
+	if self.code != nil {
+		return self.code
+	}
+	if bytes.Equal(self.CodeHash(), emptyCodeHash) {
+		return nil
+	}
+	code, err := db.Get(self.CodeHash())
+	if err != nil {
+		self.setError(fmt.Errorf("can't load code hash %x: %v", self.CodeHash(), err))
+	}
+	self.code = code
+	return code
 }
 
-func (self *StateObject) Code() []byte {
-	return self.code
+// CodeSize returns the size of the contract code associated with this object.
+func (self *StateObject) CodeSize(db trie.Database) int {
+	if self.data.codeSize == nil {
+		self.data.codeSize = new(int)
+		*self.data.codeSize = len(self.Code(db))
+	}
+	return *self.data.codeSize
 }
 
 func (self *StateObject) SetCode(code []byte) {
 	self.code = code
-	self.codeHash = crypto.Keccak256(code)
-	self.dirty = true
+	self.data.CodeHash = crypto.Keccak256(code)
+	self.data.codeSize = new(int)
+	*self.data.codeSize = len(code)
+	self.dirtyCode = true
+	if self.onDirty != nil {
+		self.onDirty(self.Address())
+		self.onDirty = nil
+	}
 }
 
 func (self *StateObject) SetNonce(nonce uint64) {
-	self.nonce = nonce
-	self.dirty = true
+	self.data.Nonce = nonce
+	if self.onDirty != nil {
+		self.onDirty(self.Address())
+		self.onDirty = nil
+	}
+}
+
+func (self *StateObject) CodeHash() []byte {
+	return self.data.CodeHash
+}
+
+func (self *StateObject) Balance() *big.Int {
+	return self.data.Balance
 }
 
 func (self *StateObject) Nonce() uint64 {
-	return self.nonce
+	return self.data.Nonce
 }
 
 // Never called, but must be present to allow StateObject to be used
@@ -262,39 +338,3 @@ func (self *StateObject) ForEachStorage(cb func(key, value common.Hash) bool) {
 		}
 	}
 }
-
-type extStateObject struct {
-	Nonce    uint64
-	Balance  *big.Int
-	Root     common.Hash
-	CodeHash []byte
-}
-
-// EncodeRLP implements rlp.Encoder.
-func (c *StateObject) EncodeRLP(w io.Writer) error {
-	return rlp.Encode(w, []interface{}{c.nonce, c.balance, c.Root(), c.codeHash})
-}
-
-// DecodeObject decodes an RLP-encoded state object.
-func DecodeObject(address common.Address, db trie.Database, data []byte) (*StateObject, error) {
-	var (
-		obj = &StateObject{address: address, db: db, storage: make(Storage)}
-		ext extStateObject
-		err error
-	)
-	if err = rlp.DecodeBytes(data, &ext); err != nil {
-		return nil, err
-	}
-	if obj.trie, err = trie.NewSecure(ext.Root, db); err != nil {
-		return nil, err
-	}
-	if !bytes.Equal(ext.CodeHash, emptyCodeHash) {
-		if obj.code, err = db.Get(ext.CodeHash); err != nil {
-			return nil, fmt.Errorf("can't get code for hash %x: %v", ext.CodeHash, err)
-		}
-	}
-	obj.nonce = ext.Nonce
-	obj.balance = ext.Balance
-	obj.codeHash = ext.CodeHash
-	return obj, nil
-}
diff --git a/core/state/state_test.go b/core/state/state_test.go
index 69cf083cf..fcdc38588 100644
--- a/core/state/state_test.go
+++ b/core/state/state_test.go
@@ -146,23 +146,23 @@ func TestSnapshot2(t *testing.T) {
 
 	// db, trie are already non-empty values
 	so0 := state.GetStateObject(stateobjaddr0)
-	so0.balance = big.NewInt(42)
-	so0.nonce = 43
+	so0.SetBalance(big.NewInt(42))
+	so0.SetNonce(43)
 	so0.SetCode([]byte{'c', 'a', 'f', 'e'})
 	so0.remove = false
 	so0.deleted = false
-	so0.dirty = true
 	state.SetStateObject(so0)
-	state.Commit()
+
+	root, _ := state.Commit()
+	state.Reset(root)
 
 	// and one with deleted == true
 	so1 := state.GetStateObject(stateobjaddr1)
-	so1.balance = big.NewInt(52)
-	so1.nonce = 53
+	so1.SetBalance(big.NewInt(52))
+	so1.SetNonce(53)
 	so1.SetCode([]byte{'c', 'a', 'f', 'e', '2'})
 	so1.remove = true
 	so1.deleted = true
-	so1.dirty = true
 	state.SetStateObject(so1)
 
 	so1 = state.GetStateObject(stateobjaddr1)
@@ -174,41 +174,50 @@ func TestSnapshot2(t *testing.T) {
 	state.Set(snapshot)
 
 	so0Restored := state.GetStateObject(stateobjaddr0)
-	so0Restored.GetState(storageaddr)
-	so1Restored := state.GetStateObject(stateobjaddr1)
+	// Update lazily-loaded values before comparing.
+	so0Restored.GetState(db, storageaddr)
+	so0Restored.Code(db)
 	// non-deleted is equal (restored)
 	compareStateObjects(so0Restored, so0, t)
+
 	// deleted should be nil, both before and after restore of state copy
+	so1Restored := state.GetStateObject(stateobjaddr1)
 	if so1Restored != nil {
-		t.Fatalf("deleted object not nil after restoring snapshot")
+		t.Fatalf("deleted object not nil after restoring snapshot: %+v", so1Restored)
 	}
 }
 
 func compareStateObjects(so0, so1 *StateObject, t *testing.T) {
-	if so0.address != so1.address {
+	if so0.Address() != so1.Address() {
 		t.Fatalf("Address mismatch: have %v, want %v", so0.address, so1.address)
 	}
-	if so0.balance.Cmp(so1.balance) != 0 {
-		t.Fatalf("Balance mismatch: have %v, want %v", so0.balance, so1.balance)
+	if so0.Balance().Cmp(so1.Balance()) != 0 {
+		t.Fatalf("Balance mismatch: have %v, want %v", so0.Balance(), so1.Balance())
+	}
+	if so0.Nonce() != so1.Nonce() {
+		t.Fatalf("Nonce mismatch: have %v, want %v", so0.Nonce(), so1.Nonce())
 	}
-	if so0.nonce != so1.nonce {
-		t.Fatalf("Nonce mismatch: have %v, want %v", so0.nonce, so1.nonce)
+	if so0.data.Root != so1.data.Root {
+		t.Errorf("Root mismatch: have %x, want %x", so0.data.Root[:], so1.data.Root[:])
 	}
-	if !bytes.Equal(so0.codeHash, so1.codeHash) {
-		t.Fatalf("CodeHash mismatch: have %v, want %v", so0.codeHash, so1.codeHash)
+	if !bytes.Equal(so0.CodeHash(), so1.CodeHash()) {
+		t.Fatalf("CodeHash mismatch: have %v, want %v", so0.CodeHash(), so1.CodeHash())
 	}
 	if !bytes.Equal(so0.code, so1.code) {
 		t.Fatalf("Code mismatch: have %v, want %v", so0.code, so1.code)
 	}
 
+	if len(so1.storage) != len(so0.storage) {
+		t.Errorf("Storage size mismatch: have %d, want %d", len(so1.storage), len(so0.storage))
+	}
 	for k, v := range so1.storage {
 		if so0.storage[k] != v {
-			t.Fatalf("Storage key %s mismatch: have %v, want %v", k, so0.storage[k], v)
+			t.Errorf("Storage key %x mismatch: have %v, want %v", k, so0.storage[k], v)
 		}
 	}
 	for k, v := range so0.storage {
 		if so1.storage[k] != v {
-			t.Fatalf("Storage key %s mismatch: have %v, want none.", k, v)
+			t.Errorf("Storage key %x mismatch: have %v, want none.", k, v)
 		}
 	}
 
@@ -218,7 +227,4 @@ func compareStateObjects(so0, so1 *StateObject, t *testing.T) {
 	if so0.deleted != so1.deleted {
 		t.Fatalf("Deleted mismatch: have %v, want %v", so0.deleted, so1.deleted)
 	}
-	if so0.dirty != so1.dirty {
-		t.Fatalf("Dirty mismatch: have %v, want %v", so0.dirty, so1.dirty)
-	}
 }
diff --git a/core/state/statedb.go b/core/state/statedb.go
index 8ba81613d..10f3f4652 100644
--- a/core/state/statedb.go
+++ b/core/state/statedb.go
@@ -43,8 +43,14 @@ type StateDB struct {
 	db   ethdb.Database
 	trie *trie.SecureTrie
 
-	stateObjects map[string]*StateObject
+	// This map caches canon state accounts.
+	all map[common.Address]Account
 
+	// This map holds 'live' objects, which will get modified while processing a state transition.
+	stateObjects      map[common.Address]*StateObject
+	stateObjectsDirty map[common.Address]struct{}
+
+	// The refund counter, also used by state transitioning.
 	refund *big.Int
 
 	thash, bhash common.Hash
@@ -60,32 +66,36 @@ func New(root common.Hash, db ethdb.Database) (*StateDB, error) {
 		return nil, err
 	}
 	return &StateDB{
-		db:           db,
-		trie:         tr,
-		stateObjects: make(map[string]*StateObject),
-		refund:       new(big.Int),
-		logs:         make(map[common.Hash]vm.Logs),
+		db:                db,
+		trie:              tr,
+		all:               make(map[common.Address]Account),
+		stateObjects:      make(map[common.Address]*StateObject),
+		stateObjectsDirty: make(map[common.Address]struct{}),
+		refund:            new(big.Int),
+		logs:              make(map[common.Hash]vm.Logs),
 	}, nil
 }
 
 // Reset clears out all emphemeral state objects from the state db, but keeps
 // the underlying state trie to avoid reloading data for the next operations.
 func (self *StateDB) Reset(root common.Hash) error {
-	var (
-		err error
-		tr  = self.trie
-	)
+	tr, err := trie.NewSecure(root, self.db)
+	if err != nil {
+		return err
+	}
+	all := self.all
 	if self.trie.Hash() != root {
-		if tr, err = trie.NewSecure(root, self.db); err != nil {
-			return err
-		}
+		// The root has changed, invalidate canon state.
+		all = make(map[common.Address]Account)
 	}
 	*self = StateDB{
-		db:           self.db,
-		trie:         tr,
-		stateObjects: make(map[string]*StateObject),
-		refund:       new(big.Int),
-		logs:         make(map[common.Hash]vm.Logs),
+		db:                self.db,
+		trie:              tr,
+		all:               all,
+		stateObjects:      make(map[common.Address]*StateObject),
+		stateObjectsDirty: make(map[common.Address]struct{}),
+		refund:            new(big.Int),
+		logs:              make(map[common.Hash]vm.Logs),
 	}
 	return nil
 }
@@ -137,7 +147,7 @@ func (self *StateDB) GetAccount(addr common.Address) vm.Account {
 func (self *StateDB) GetBalance(addr common.Address) *big.Int {
 	stateObject := self.GetStateObject(addr)
 	if stateObject != nil {
-		return stateObject.balance
+		return stateObject.Balance()
 	}
 
 	return common.Big0
@@ -146,7 +156,7 @@ func (self *StateDB) GetBalance(addr common.Address) *big.Int {
 func (self *StateDB) GetNonce(addr common.Address) uint64 {
 	stateObject := self.GetStateObject(addr)
 	if stateObject != nil {
-		return stateObject.nonce
+		return stateObject.Nonce()
 	}
 
 	return StartingNonce
@@ -155,18 +165,24 @@ func (self *StateDB) GetNonce(addr common.Address) uint64 {
 func (self *StateDB) GetCode(addr common.Address) []byte {
 	stateObject := self.GetStateObject(addr)
 	if stateObject != nil {
-		return stateObject.code
+		return stateObject.Code(self.db)
 	}
-
 	return nil
 }
 
+func (self *StateDB) GetCodeSize(addr common.Address) int {
+	stateObject := self.GetStateObject(addr)
+	if stateObject != nil {
+		return stateObject.CodeSize(self.db)
+	}
+	return 0
+}
+
 func (self *StateDB) GetState(a common.Address, b common.Hash) common.Hash {
 	stateObject := self.GetStateObject(a)
 	if stateObject != nil {
-		return stateObject.GetState(b)
+		return stateObject.GetState(self.db, b)
 	}
-
 	return common.Hash{}
 }
 
@@ -214,8 +230,7 @@ func (self *StateDB) Delete(addr common.Address) bool {
 	stateObject := self.GetStateObject(addr)
 	if stateObject != nil {
 		stateObject.MarkForDeletion()
-		stateObject.balance = new(big.Int)
-
+		stateObject.data.Balance = new(big.Int)
 		return true
 	}
 
@@ -242,35 +257,47 @@ func (self *StateDB) DeleteStateObject(stateObject *StateObject) {
 
 	addr := stateObject.Address()
 	self.trie.Delete(addr[:])
-	//delete(self.stateObjects, addr.Str())
 }
 
-// Retrieve a state object given my the address. Nil if not found
+// Retrieve a state object given my the address. Returns nil if not found.
 func (self *StateDB) GetStateObject(addr common.Address) (stateObject *StateObject) {
-	stateObject = self.stateObjects[addr.Str()]
-	if stateObject != nil {
-		if stateObject.deleted {
-			stateObject = nil
+	// Prefer 'live' objects.
+	if obj := self.stateObjects[addr]; obj != nil {
+		if obj.deleted {
+			return nil
 		}
+		return obj
+	}
 
-		return stateObject
+	// Use cached account data from the canon state if possible.
+	if data, ok := self.all[addr]; ok {
+		obj := NewObject(addr, data, self.MarkStateObjectDirty)
+		self.SetStateObject(obj)
+		return obj
 	}
 
-	data := self.trie.Get(addr[:])
-	if len(data) == 0 {
+	// Load the object from the database.
+	enc := self.trie.Get(addr[:])
+	if len(enc) == 0 {
 		return nil
 	}
-	stateObject, err := DecodeObject(addr, self.db, data)
-	if err != nil {
+	var data Account
+	if err := rlp.DecodeBytes(enc, &data); err != nil {
 		glog.Errorf("can't decode object at %x: %v", addr[:], err)
 		return nil
 	}
-	self.SetStateObject(stateObject)
-	return stateObject
+	// Update the all cache. Content in DB always corresponds
+	// to the current head state so this is ok to do here.
+	// The object we just loaded has no storage trie and code yet.
+	self.all[addr] = data
+	// Insert into the live set.
+	obj := NewObject(addr, data, self.MarkStateObjectDirty)
+	self.SetStateObject(obj)
+	return obj
 }
 
 func (self *StateDB) SetStateObject(object *StateObject) {
-	self.stateObjects[object.Address().Str()] = object
+	self.stateObjects[object.Address()] = object
 }
 
 // Retrieve a state object or create a new state object if nil
@@ -288,15 +315,19 @@ func (self *StateDB) newStateObject(addr common.Address) *StateObject {
 	if glog.V(logger.Core) {
 		glog.Infof("(+) %x\n", addr)
 	}
+	obj := NewObject(addr, Account{}, self.MarkStateObjectDirty)
+	obj.SetNonce(StartingNonce) // sets the object to dirty
+	self.stateObjects[addr] = obj
+	return obj
+}
 
-	stateObject := NewStateObject(addr, self.db)
-	stateObject.SetNonce(StartingNonce)
-	self.stateObjects[addr.Str()] = stateObject
-
-	return stateObject
+// MarkStateObjectDirty adds the specified object to the dirty map to avoid costly
+// state object cache iteration to find a handful of modified ones.
+func (self *StateDB) MarkStateObjectDirty(addr common.Address) {
+	self.stateObjectsDirty[addr] = struct{}{}
 }
 
-// Creates creates a new state object and takes ownership. This is different from "NewStateObject"
+// Creates creates a new state object and takes ownership.
 func (self *StateDB) CreateStateObject(addr common.Address) *StateObject {
 	// Get previous (if any)
 	so := self.GetStateObject(addr)
@@ -305,7 +336,7 @@ func (self *StateDB) CreateStateObject(addr common.Address) *StateObject {
 
 	// If it existed set the balance to the new account
 	if so != nil {
-		newSo.balance = so.balance
+		newSo.data.Balance = so.data.Balance
 	}
 
 	return newSo
@@ -320,29 +351,34 @@ func (self *StateDB) CreateAccount(addr common.Address) vm.Account {
 //
 
 func (self *StateDB) Copy() *StateDB {
-	// ignore error - we assume state-to-be-copied always exists
-	state, _ := New(common.Hash{}, self.db)
-	state.trie = self.trie
-	for k, stateObject := range self.stateObjects {
-		if stateObject.dirty {
-			state.stateObjects[k] = stateObject.Copy()
-		}
+	// Copy all the basic fields, initialize the memory ones
+	state := &StateDB{
+		db:                self.db,
+		trie:              self.trie,
+		all:               self.all,
+		stateObjects:      make(map[common.Address]*StateObject, len(self.stateObjectsDirty)),
+		stateObjectsDirty: make(map[common.Address]struct{}, len(self.stateObjectsDirty)),
+		refund:            new(big.Int).Set(self.refund),
+		logs:              make(map[common.Hash]vm.Logs, len(self.logs)),
+		logSize:           self.logSize,
+	}
+	// Copy the dirty states and logs
+	for addr, _ := range self.stateObjectsDirty {
+		state.stateObjects[addr] = self.stateObjects[addr].Copy(self.db, state.MarkStateObjectDirty)
+		state.stateObjectsDirty[addr] = struct{}{}
 	}
-
-	state.refund.Set(self.refund)
-
 	for hash, logs := range self.logs {
 		state.logs[hash] = make(vm.Logs, len(logs))
 		copy(state.logs[hash], logs)
 	}
-	state.logSize = self.logSize
-
 	return state
 }
 
 func (self *StateDB) Set(state *StateDB) {
 	self.trie = state.trie
 	self.stateObjects = state.stateObjects
+	self.stateObjectsDirty = state.stateObjectsDirty
+	self.all = state.all
 
 	self.refund = state.refund
 	self.logs = state.logs
@@ -358,14 +394,13 @@ func (self *StateDB) GetRefund() *big.Int {
 // goes into transaction receipts.
 func (s *StateDB) IntermediateRoot() common.Hash {
 	s.refund = new(big.Int)
-	for _, stateObject := range s.stateObjects {
-		if stateObject.dirty {
-			if stateObject.remove {
-				s.DeleteStateObject(stateObject)
-			} else {
-				stateObject.Update()
-				s.UpdateStateObject(stateObject)
-			}
+	for addr, _ := range s.stateObjectsDirty {
+		stateObject := s.stateObjects[addr]
+		if stateObject.remove {
+			s.DeleteStateObject(stateObject)
+		} else {
+			stateObject.UpdateRoot(s.db)
+			s.UpdateStateObject(stateObject)
 		}
 	}
 	return s.trie.Hash()
@@ -380,15 +415,15 @@ func (s *StateDB) DeleteSuicides() {
 	// Reset refund so that any used-gas calculations can use
 	// this method.
 	s.refund = new(big.Int)
-	for _, stateObject := range s.stateObjects {
-		if stateObject.dirty {
-			// If the object has been removed by a suicide
-			// flag the object as deleted.
-			if stateObject.remove {
-				stateObject.deleted = true
-			}
-			stateObject.dirty = false
+	for addr, _ := range s.stateObjectsDirty {
+		stateObject := s.stateObjects[addr]
+
+		// If the object has been removed by a suicide
+		// flag the object as deleted.
+		if stateObject.remove {
+			stateObject.deleted = true
 		}
+		delete(s.stateObjectsDirty, addr)
 	}
 }
 
@@ -407,46 +442,44 @@ func (s *StateDB) CommitBatch() (root common.Hash, batch ethdb.Batch) {
 	return root, batch
 }
 
-func (s *StateDB) commit(db trie.DatabaseWriter) (common.Hash, error) {
+func (s *StateDB) commit(dbw trie.DatabaseWriter) (root common.Hash, err error) {
 	s.refund = new(big.Int)
+	defer func() {
+		if err != nil {
+			// Committing failed, any updates to the canon state are invalid.
+			s.all = make(map[common.Address]Account)
+		}
+	}()
 
-	for _, stateObject := range s.stateObjects {
+	// Commit objects to the trie.
+	for addr, stateObject := range s.stateObjects {
 		if stateObject.remove {
 			// If the object has been removed, don't bother syncing it
 			// and just mark it for deletion in the trie.
 			s.DeleteStateObject(stateObject)
-		} else {
+			delete(s.all, addr)
+		} else if _, ok := s.stateObjectsDirty[addr]; ok {
 			// Write any contract code associated with the state object
-			if len(stateObject.code) > 0 {
-				if err := db.Put(stateObject.codeHash, stateObject.code); err != nil {
+			if stateObject.code != nil && stateObject.dirtyCode {
+				if err := dbw.Put(stateObject.CodeHash(), stateObject.code); err != nil {
 					return common.Hash{}, err
 				}
+				stateObject.dirtyCode = false
 			}
-			// Write any storage changes in the state object to its trie.
-			stateObject.Update()
-
-			// Commit the trie of the object to the batch.
-			// This updates the trie root internally, so
-			// getting the root hash of the storage trie
-			// through UpdateStateObject is fast.
-			if _, err := stateObject.trie.CommitTo(db); err != nil {
+			// Write any storage changes in the state object to its storage trie.
+			if err := stateObject.CommitTrie(s.db, dbw); err != nil {
 				return common.Hash{}, err
 			}
-			// Update the object in the account trie.
+			// Update the object in the main account trie.
 			s.UpdateStateObject(stateObject)
+			s.all[addr] = stateObject.data
 		}
-		stateObject.dirty = false
+		delete(s.stateObjectsDirty, addr)
 	}
-	return s.trie.CommitTo(db)
+	// Write trie changes.
+	return s.trie.CommitTo(dbw)
 }
 
 func (self *StateDB) Refunds() *big.Int {
 	return self.refund
 }
-
-// Debug stuff
-func (self *StateDB) CreateOutputForDiff() {
-	for _, stateObject := range self.stateObjects {
-		stateObject.CreateOutputForDiff()
-	}
-}