core/state, core/vm: cleanup refunds

3 files changed, 13 insertions, 21 deletions

diff --git a/core/state/statedb.go b/core/state/statedb.go
index b3050515b..895d9fe8b 100644
--- a/core/state/statedb.go
+++ b/core/state/statedb.go
@@ -21,7 +21,7 @@ type StateDB struct {
 
 	stateObjects map[string]*StateObject
 
-	refund map[string]*big.Int
+	refund *big.Int
 
 	thash, bhash common.Hash
 	txIndex      int
@@ -31,7 +31,7 @@ type StateDB struct {
 // Create a new state from a given trie
 func New(root common.Hash, db common.Database) *StateDB {
 	trie := trie.NewSecure(root[:], db)
-	return &StateDB{db: db, trie: trie, stateObjects: make(map[string]*StateObject), refund: make(map[string]*big.Int), logs: make(map[common.Hash]Logs)}
+	return &StateDB{db: db, trie: trie, stateObjects: make(map[string]*StateObject), refund: new(big.Int), logs: make(map[common.Hash]Logs)}
 }
 
 func (self *StateDB) PrintRoot() {
@@ -63,12 +63,8 @@ func (self *StateDB) Logs() Logs {
 	return logs
 }
 
-func (self *StateDB) Refund(address common.Address, gas *big.Int) {
-	addr := address.Str()
-	if self.refund[addr] == nil {
-		self.refund[addr] = new(big.Int)
-	}
-	self.refund[addr].Add(self.refund[addr], gas)
+func (self *StateDB) Refund(gas *big.Int) {
+	self.refund.Add(self.refund, gas)
 }
 
 /*
@@ -268,9 +264,7 @@ func (self *StateDB) Copy() *StateDB {
 		state.stateObjects[k] = stateObject.Copy()
 	}
 
-	for addr, refund := range self.refund {
-		state.refund[addr] = new(big.Int).Set(refund)
-	}
+	state.refund.Set(self.refund)
 
 	for hash, logs := range self.logs {
 		state.logs[hash] = make(Logs, len(logs))
@@ -330,15 +324,15 @@ func (s *StateDB) Sync() {
 
 func (self *StateDB) Empty() {
 	self.stateObjects = make(map[string]*StateObject)
-	self.refund = make(map[string]*big.Int)
+	self.refund = new(big.Int)
 }
 
-func (self *StateDB) Refunds() map[string]*big.Int {
+func (self *StateDB) Refunds() *big.Int {
 	return self.refund
 }
 
 func (self *StateDB) Update() {
-	self.refund = make(map[string]*big.Int)
+	self.refund = new(big.Int)
 
 	for _, stateObject := range self.stateObjects {
 		if stateObject.dirty {
diff --git a/core/state_transition.go b/core/state_transition.go
index fedea8021..4a8d92375 100644
--- a/core/state_transition.go
+++ b/core/state_transition.go
@@ -241,11 +241,9 @@ func (self *StateTransition) refundGas() {
 	sender.AddBalance(remaining)
 
 	uhalf := new(big.Int).Div(self.gasUsed(), common.Big2)
-	for addr, ref := range self.state.Refunds() {
-		refund := common.BigMin(uhalf, ref)
-		self.gas.Add(self.gas, refund)
-		self.state.AddBalance(common.StringToAddress(addr), refund.Mul(refund, self.msg.GasPrice()))
-	}
+	refund := common.BigMin(uhalf, self.state.Refunds())
+	self.gas.Add(self.gas, refund)
+	self.state.AddBalance(sender.Address(), refund.Mul(refund, self.msg.GasPrice()))
 
 	coinbase.RefundGas(self.gas, self.msg.GasPrice())
 }
diff --git a/core/vm/vm.go b/core/vm/vm.go
index c5ad761f6..0486fbbc7 100644
--- a/core/vm/vm.go
+++ b/core/vm/vm.go
@@ -690,7 +690,7 @@ func (self *Vm) calculateGasAndSize(context *Context, caller ContextRef, op OpCo
 			// 0 => non 0
 			g = params.SstoreSetGas
 		} else if len(val) > 0 && len(y.Bytes()) == 0 {
-			statedb.Refund(self.env.Origin(), params.SstoreRefundGas)
+			statedb.Refund(params.SstoreRefundGas)
 
 			g = params.SstoreClearGas
 		} else {
@@ -700,7 +700,7 @@ func (self *Vm) calculateGasAndSize(context *Context, caller ContextRef, op OpCo
 		gas.Set(g)
 	case SUICIDE:
 		if !statedb.IsDeleted(context.Address()) {
-			statedb.Refund(self.env.Origin(), params.SuicideRefundGas)
+			statedb.Refund(params.SuicideRefundGas)
 		}
 	case MLOAD:
 		newMemSize = calcMemSize(stack.peek(), u256(32))