Implemented a manage state for keeping track of nonces

3 files changed, 63 insertions, 6 deletions

diff --git a/core/chain_manager.go b/core/chain_manager.go
index df3d6d1c3..7fed7b8b1 100644
--- a/core/chain_manager.go
+++ b/core/chain_manager.go
@@ -86,7 +86,7 @@ type ChainManager struct {
 	lastBlockHash []byte
 
 	transState *state.StateDB
-	txState    *state.StateDB
+	txState    *state.ManagedState
 
 	quit chan struct{}
 }
@@ -95,7 +95,8 @@ func NewChainManager(blockDb, stateDb ethutil.Database, mux *event.TypeMux) *Cha
 	bc := &ChainManager{blockDb: blockDb, stateDb: stateDb, genesisBlock: GenesisBlock(stateDb), eventMux: mux, quit: make(chan struct{})}
 	bc.setLastBlock()
 	bc.transState = bc.State().Copy()
-	bc.txState = bc.State().Copy()
+	// Take ownership of this particular state
+	bc.txState = state.ManageState(bc.State().Copy())
 	go bc.update()
 
 	return bc
@@ -144,17 +145,17 @@ func (self *ChainManager) TransState() *state.StateDB {
 	return self.transState
 }
 
-func (self *ChainManager) TxState() *state.StateDB {
+func (self *ChainManager) TxState() *state.ManagedState {
 	self.tsmu.RLock()
 	defer self.tsmu.RUnlock()
 
 	return self.txState
 }
 
-func (self *ChainManager) setTxState(state *state.StateDB) {
+func (self *ChainManager) setTxState(statedb *state.StateDB) {
 	self.tsmu.Lock()
 	defer self.tsmu.Unlock()
-	self.txState = state
+	self.txState = state.ManageState(statedb)
 }
 
 func (self *ChainManager) setTransState(statedb *state.StateDB) {
diff --git a/state/managed_state.go b/state/managed_state.go
new file mode 100644
index 000000000..b0a129af0
--- /dev/null
+++ b/state/managed_state.go
@@ -0,0 +1,56 @@
+package state
+
+import "sync"
+
+type ManagedState struct {
+	*StateDB
+
+	mu sync.RWMutex
+
+	accounts map[string]*StateObject
+}
+
+func ManageState(statedb *StateDB) *ManagedState {
+	return &ManagedState{
+		StateDB:  statedb,
+		accounts: make(map[string]*StateObject),
+	}
+}
+
+func (ms *ManagedState) IncrementNonce(addr []byte) {
+	ms.mu.Lock()
+	defer ms.mu.Unlock()
+
+	ms.getAccount(addr).nonce++
+}
+
+func (ms *ManagedState) DecrementNonce(addr []byte) {
+	// Decrementing a nonce does not mean we are interested in the account
+	// incrementing only happens if you control the account, therefor
+	// incrementing  behaves differently from decrementing
+	if ms.hasAccount(addr) {
+		ms.mu.Lock()
+		defer ms.mu.Unlock()
+
+		ms.getAccount(addr).nonce--
+	}
+}
+
+func (ms *ManagedState) GetNonce(addr []byte) uint64 {
+	ms.mu.RLock()
+	defer ms.mu.RUnlock()
+	return ms.getAccount(addr).nonce
+}
+
+func (ms *ManagedState) hasAccount(addr []byte) bool {
+	_, ok := ms.accounts[string(addr)]
+	return ok
+}
+
+func (ms *ManagedState) getAccount(addr []byte) *StateObject {
+	if _, ok := ms.accounts[string(addr)]; !ok {
+		ms.accounts[string(addr)] = ms.GetOrNewStateObject(addr)
+	}
+
+	return ms.accounts[string(addr)]
+}
diff --git a/xeth/xeth.go b/xeth/xeth.go
index 4c28caf86..4ce1ee6a5 100644
--- a/xeth/xeth.go
+++ b/xeth/xeth.go
@@ -362,7 +362,7 @@ func (self *XEth) Transact(fromStr, toStr, valueStr, gasStr, gasPriceStr, codeSt
 	if err := self.eth.TxPool().Add(tx); err != nil {
 		return "", err
 	}
-	state.SetNonce(from, nonce+1)
+	state.IncrementNonce(from)
 
 	if contractCreation {
 		addr := core.AddressFromMessage(tx)