Cache whole objects instead of states only

2 files changed, 117 insertions, 2 deletions

diff --git a/ethchain/state.go b/ethchain/state.go
index 5af748e00..f413fb8ed 100644
--- a/ethchain/state.go
+++ b/ethchain/state.go
@@ -16,14 +16,17 @@ type State struct {
 	// Nested states
 	states map[string]*State
 
+	stateObjects map[string]*StateObject
+
 	manifest *Manifest
 }
 
 // Create a new state from a given trie
 func NewState(trie *ethutil.Trie) *State {
-	return &State{trie: trie, states: make(map[string]*State), manifest: NewManifest()}
+	return &State{trie: trie, states: make(map[string]*State), stateObjects: make(map[string]*StateObject), manifest: NewManifest()}
 }
 
+/*
 // Resets the trie and all siblings
 func (s *State) Reset() {
 	s.trie.Undo()
@@ -43,6 +46,35 @@ func (s *State) Sync() {
 
 	s.trie.Sync()
 }
+*/
+
+// Resets the trie and all siblings
+func (s *State) Reset() {
+	s.trie.Undo()
+
+	// Reset all nested states
+	for _, stateObject := range s.stateObjects {
+		if stateObject.state == nil {
+			continue
+		}
+
+		stateObject.state.Reset()
+	}
+}
+
+// Syncs the trie and all siblings
+func (s *State) Sync() {
+	// Sync all nested states
+	for _, stateObject := range s.stateObjects {
+		if stateObject.state == nil {
+			continue
+		}
+
+		stateObject.state.Sync()
+	}
+
+	s.trie.Sync()
+}
 
 // Purges the current trie.
 func (s *State) Purge() int {
@@ -54,6 +86,7 @@ func (s *State) EachStorage(cb ethutil.EachCallback) {
 	it.Each(cb)
 }
 
+/*
 func (s *State) GetStateObject(addr []byte) *StateObject {
 	data := s.trie.Get(string(addr))
 	if data == "" {
@@ -78,7 +111,6 @@ func (s *State) UpdateStateObject(object *StateObject) {
 
 	if object.state != nil && s.states[string(addr)] == nil {
 		s.states[string(addr)] = object.state
-		//fmt.Printf("update cached #%d %x addr: %x\n", object.state.trie.Cache().Len(), object.state.Root(), addr[0:4])
 	}
 
 	ethutil.Config.Db.Put(ethutil.Sha3Bin(object.Script()), object.Script())
@@ -96,13 +128,66 @@ func (s *State) GetAccount(addr []byte) (account *StateObject) {
 		account = NewStateObjectFromBytes(addr, []byte(data))
 	}
 
+	// Check if there's a cached state for this contract
+	cachedStateObject := s.states[string(addr)]
+	if cachedStateObject != nil {
+		account.state = cachedStateObject
+	}
+
 	return
 }
+*/
+
+func (self *State) UpdateStateObject(stateObject *StateObject) {
+	addr := stateObject.Address()
+
+	if self.stateObjects[string(addr)] == nil {
+		self.stateObjects[string(addr)] = stateObject
+	}
+
+	ethutil.Config.Db.Put(ethutil.Sha3Bin(stateObject.Script()), stateObject.Script())
+
+	self.trie.Update(string(addr), string(stateObject.RlpEncode()))
+
+	self.manifest.AddObjectChange(stateObject)
+}
+
+func (self *State) GetStateObject(addr []byte) *StateObject {
+	stateObject := self.stateObjects[string(addr)]
+	if stateObject != nil {
+		return stateObject
+	}
+
+	data := self.trie.Get(string(addr))
+	if len(data) == 0 {
+		return nil
+	}
+
+	stateObject = NewStateObjectFromBytes(addr, []byte(data))
+	self.stateObjects[string(addr)] = stateObject
+
+	return stateObject
+}
+
+func (self *State) GetOrNewStateObject(addr []byte) *StateObject {
+	stateObject := self.GetStateObject(addr)
+	if stateObject == nil {
+		stateObject = NewStateObject(addr)
+		self.stateObjects[string(addr)] = stateObject
+	}
+
+	return stateObject
+}
+
+func (self *State) GetAccount(addr []byte) *StateObject {
+	return self.GetOrNewStateObject(addr)
+}
 
 func (s *State) Cmp(other *State) bool {
 	return s.trie.Cmp(other.trie)
 }
 
+/*
 func (s *State) Copy() *State {
 	state := NewState(s.trie.Copy())
 	for k, subState := range s.states {
@@ -111,6 +196,15 @@ func (s *State) Copy() *State {
 
 	return state
 }
+*/
+func (self *State) Copy() *State {
+	state := NewState(self.trie.Copy())
+	for k, stateObject := range self.stateObjects {
+		state.stateObjects[k] = stateObject.Copy()
+	}
+
+	return state
+}
 
 func (s *State) Snapshot() *State {
 	return s.Copy()
diff --git a/ethchain/state_object.go b/ethchain/state_object.go
index 12ebf8e9b..3775d436c 100644
--- a/ethchain/state_object.go
+++ b/ethchain/state_object.go
@@ -38,6 +38,10 @@ func MakeContract(tx *Transaction, state *State) *StateObject {
 	return nil
 }
 
+func NewStateObject(addr []byte) *StateObject {
+	return &StateObject{address: addr, Amount: new(big.Int)}
+}
+
 func NewContract(address []byte, Amount *big.Int, root []byte) *StateObject {
 	contract := &StateObject{address: address, Amount: Amount, Nonce: 0}
 	contract.state = NewState(ethutil.NewTrie(ethutil.Config.Db, string(root)))
@@ -146,6 +150,23 @@ func (self *StateObject) BuyGas(gas, price *big.Int) error {
 	return nil
 }
 
+func (self *StateObject) Copy() *StateObject {
+	stCopy := &StateObject{}
+	stCopy.address = make([]byte, len(self.address))
+	copy(stCopy.address, self.address)
+	stCopy.Amount = new(big.Int).Set(self.Amount)
+	stCopy.ScriptHash = make([]byte, len(self.ScriptHash))
+	copy(stCopy.ScriptHash, self.ScriptHash)
+	stCopy.Nonce = self.Nonce
+	stCopy.state = self.state.Copy()
+	stCopy.script = make([]byte, len(self.script))
+	copy(stCopy.script, self.script)
+	stCopy.initScript = make([]byte, len(self.initScript))
+	copy(stCopy.initScript, self.initScript)
+
+	return stCopy
+}
+
 // Returns the address of the contract/account
 func (c *StateObject) Address() []byte {
 	return c.address