Trie's are no longer referenced directly but through State instead

4 files changed, 63 insertions, 74 deletions

diff --git a/ethchain/block.go b/ethchain/block.go
index 8de57cced..b5739102c 100644
--- a/ethchain/block.go
+++ b/ethchain/block.go
@@ -126,54 +126,6 @@ func (block *Block) Transactions() []*Transaction {
 	return block.transactions
 }
 
-/*
-func (block *Block) GetContract(addr []byte) *Contract {
-	data := block.state.Get(string(addr))
-	if data == "" {
-		return nil
-	}
-
-	value := ethutil.NewValueFromBytes([]byte(data))
-	if value.Len() == 2 {
-		return nil
-	}
-
-	contract := &Contract{}
-	contract.RlpDecode([]byte(data))
-
-	cachedState := block.contractStates[string(addr)]
-	if cachedState != nil {
-		contract.state = cachedState
-	} else {
-		block.contractStates[string(addr)] = contract.state
-	}
-
-	return contract
-}
-func (block *Block) UpdateContract(addr []byte, contract *Contract) {
-	// Make sure the state is synced
-	//contract.State().Sync()
-
-	block.state.trie.Update(string(addr), string(contract.RlpEncode()))
-}
-
-func (block *Block) GetAddr(addr []byte) *Address {
-	var address *Address
-
-	data := block.state.trie.Get(string(addr))
-	if data == "" {
-		address = NewAddress(big.NewInt(0))
-	} else {
-		address = NewAddressFromData([]byte(data))
-	}
-
-	return address
-}
-func (block *Block) UpdateAddr(addr []byte, address *Address) {
-	block.state.trie.Update(string(addr), string(address.RlpEncode()))
-}
-*/
-
 func (block *Block) PayFee(addr []byte, fee *big.Int) bool {
 	contract := block.state.GetContract(addr)
 	// If we can't pay the fee return
@@ -210,23 +162,10 @@ func (block *Block) BlockInfo() BlockInfo {
 
 // Sync the block's state and contract respectively
 func (block *Block) Sync() {
-	/*
-		// Sync all contracts currently in cache
-		for _, val := range block.contractStates {
-			val.Sync()
-		}
-	*/
-	// Sync the block state itself
-	block.state.trie.Sync()
+	block.state.Sync()
 }
 
 func (block *Block) Undo() {
-	/*
-		// Sync all contracts currently in cache
-		for _, val := range block.contractStates {
-			val.Undo()
-		}
-	*/
 	// Sync the block state itself
 	block.state.Reset()
 }
@@ -234,7 +173,7 @@ func (block *Block) Undo() {
 func (block *Block) MakeContract(tx *Transaction) {
 	contract := MakeContract(tx, block.state)
 	if contract != nil {
-		block.contractStates[string(tx.Hash()[12:])] = contract.state
+		block.state.states[string(tx.Hash()[12:])] = contract.state
 	}
 }
 
diff --git a/ethchain/contract.go b/ethchain/contract.go
index dbcbb3697..21ac828fe 100644
--- a/ethchain/contract.go
+++ b/ethchain/contract.go
@@ -8,18 +8,19 @@ import (
 type Contract struct {
 	Amount *big.Int
 	Nonce  uint64
-	state  *ethutil.Trie
+	//state  *ethutil.Trie
+	state *State
 }
 
 func NewContract(Amount *big.Int, root []byte) *Contract {
 	contract := &Contract{Amount: Amount, Nonce: 0}
-	contract.state = ethutil.NewTrie(ethutil.Config.Db, string(root))
+	contract.state = NewState(ethutil.NewTrie(ethutil.Config.Db, string(root)))
 
 	return contract
 }
 
 func (c *Contract) RlpEncode() []byte {
-	return ethutil.Encode([]interface{}{c.Amount, c.Nonce, c.state.Root})
+	return ethutil.Encode([]interface{}{c.Amount, c.Nonce, c.state.trie.Root})
 }
 
 func (c *Contract) RlpDecode(data []byte) {
@@ -27,18 +28,18 @@ func (c *Contract) RlpDecode(data []byte) {
 
 	c.Amount = decoder.Get(0).BigInt()
 	c.Nonce = decoder.Get(1).Uint()
-	c.state = ethutil.NewTrie(ethutil.Config.Db, decoder.Get(2).Interface())
+	c.state = NewState(ethutil.NewTrie(ethutil.Config.Db, decoder.Get(2).Interface()))
 }
 
 func (c *Contract) Addr(addr []byte) *ethutil.Value {
-	return ethutil.NewValueFromBytes([]byte(c.state.Get(string(addr))))
+	return ethutil.NewValueFromBytes([]byte(c.state.trie.Get(string(addr))))
 }
 
 func (c *Contract) SetAddr(addr []byte, value interface{}) {
-	c.state.Update(string(addr), string(ethutil.NewValue(value).Encode()))
+	c.state.trie.Update(string(addr), string(ethutil.NewValue(value).Encode()))
 }
 
-func (c *Contract) State() *ethutil.Trie {
+func (c *Contract) State() *State {
 	return c.state
 }
 
@@ -59,7 +60,7 @@ func MakeContract(tx *Transaction, state *State) *Contract {
 		for i, val := range tx.Data {
 			if len(val) > 0 {
 				bytNum := ethutil.BigToBytes(big.NewInt(int64(i)), 256)
-				contract.state.Update(string(bytNum), string(ethutil.Encode(val)))
+				contract.state.trie.Update(string(bytNum), string(ethutil.Encode(val)))
 			}
 		}
 		state.trie.Update(string(addr), string(contract.RlpEncode()))
diff --git a/ethchain/state.go b/ethchain/state.go
index cff192b54..4cd2c58ef 100644
--- a/ethchain/state.go
+++ b/ethchain/state.go
@@ -5,16 +5,46 @@ import (
 	"math/big"
 )
 
+// States within the ethereum protocol are used to store anything
+// within the merkle trie. States take care of caching and storing
+// nested states. It's the general query interface to retrieve:
+// * Contracts
+// * Accounts
 type State struct {
+	// The trie for this structure
 	trie *ethutil.Trie
+	// Nested states
+	states map[string]*State
 }
 
+// Create a new state from a given trie
 func NewState(trie *ethutil.Trie) *State {
-	return &State{trie: trie}
+	return &State{trie: trie, states: make(map[string]*State)}
 }
 
+// Resets the trie and all siblings
 func (s *State) Reset() {
 	s.trie.Undo()
+
+	// Reset all nested states
+	for _, state := range s.states {
+		state.Reset()
+	}
+}
+
+// Syncs the trie and all siblings
+func (s *State) Sync() {
+	s.trie.Sync()
+
+	// Sync all nested states
+	for _, state := range s.states {
+		state.Sync()
+	}
+}
+
+// Purges the current trie.
+func (s *State) Purge() int {
+	return s.trie.NewIterator().Purge()
 }
 
 func (s *State) GetContract(addr []byte) *Contract {
@@ -23,9 +53,28 @@ func (s *State) GetContract(addr []byte) *Contract {
 		return nil
 	}
 
+	// Whet get contract is called the retrieved value might
+	// be an account. The StateManager uses this to check
+	// to see if the address a tx was sent to is a contract
+	// or an account
+	value := ethutil.NewValueFromBytes([]byte(data))
+	if value.Len() == 2 {
+		return nil
+	}
+
+	// build contract
 	contract := &Contract{}
 	contract.RlpDecode([]byte(data))
 
+	// Check if there's a cached state for this contract
+	cachedState := s.states[string(addr)]
+	if cachedState != nil {
+		contract.state = cachedState
+	} else {
+		// If it isn't cached, cache the state
+		s.states[string(addr)] = contract.state
+	}
+
 	return contract
 }
 
diff --git a/ethchain/vm.go b/ethchain/vm.go
index c7a91a9c5..7e119ac99 100644
--- a/ethchain/vm.go
+++ b/ethchain/vm.go
@@ -330,7 +330,7 @@ out:
 			// Load the value in storage and push it on the stack
 			x := vm.stack.Pop()
 			// decode the object as a big integer
-			decoder := ethutil.NewValueFromBytes([]byte(contract.State().Get(x.String())))
+			decoder := contract.Addr(x.Bytes())
 			if !decoder.IsNil() {
 				vm.stack.Push(decoder.BigInt())
 			} else {
@@ -375,7 +375,7 @@ out:
 		case oSUICIDE:
 			recAddr := vm.stack.Pop().Bytes()
 			// Purge all memory
-			deletedMemory := contract.state.NewIterator().Purge()
+			deletedMemory := contract.state.Purge()
 			// Add refunds to the pop'ed address
 			refund := new(big.Int).Mul(StoreFee, big.NewInt(int64(deletedMemory)))
 			account := state.GetAccount(recAddr)