Cleanup VM

3 files changed, 22 insertions, 10 deletions

diff --git a/core/vm/common.go b/core/vm/common.go
index 8d8f4253f..5226f4828 100644
--- a/core/vm/common.go
+++ b/core/vm/common.go
@@ -80,3 +80,13 @@ func getData(data []byte, start, size *big.Int) []byte {
 	e := common.BigMin(new(big.Int).Add(s, size), dlen)
 	return common.RightPadBytes(data[s.Uint64():e.Uint64()], int(size.Uint64()))
 }
+
+func UseGas(gas, amount *big.Int) bool {
+	if gas.Cmp(amount) < 0 {
+		return false
+	}
+
+	// Sub the amount of gas from the remaining
+	gas.Sub(gas, amount)
+	return true
+}
diff --git a/core/vm/context.go b/core/vm/context.go
index e73199b77..e4b94b600 100644
--- a/core/vm/context.go
+++ b/core/vm/context.go
@@ -74,16 +74,12 @@ func (c *Context) Return(ret []byte) []byte {
 /*
  * Gas functions
  */
-func (c *Context) UseGas(gas *big.Int) bool {
-	if c.Gas.Cmp(gas) < 0 {
-		return false
+func (c *Context) UseGas(gas *big.Int) (ok bool) {
+	ok = UseGas(c.Gas, gas)
+	if ok {
+		c.UsedGas.Add(c.UsedGas, gas)
 	}
-
-	// Sub the amount of gas from the remaining
-	c.Gas.Sub(c.Gas, gas)
-	c.UsedGas.Add(c.UsedGas, gas)
-
-	return true
+	return
 }
 
 // Implement the caller interface
diff --git a/core/vm/vm.go b/core/vm/vm.go
index 88fbdf763..eceb6e3a9 100644
--- a/core/vm/vm.go
+++ b/core/vm/vm.go
@@ -24,6 +24,9 @@ type Vm struct {
 	Fn          string
 
 	Recoverable bool
+
+	// Will be called before the vm returns
+	After func(*Context, error)
 }
 
 func New(env Environment) *Vm {
@@ -47,6 +50,10 @@ func (self *Vm) Run(context *Context, callData []byte) (ret []byte, err error) {
 
 	// User defer pattern to check for an error and, based on the error being nil or not, use all gas and return.
 	defer func() {
+		if self.After != nil {
+			self.After(context, err)
+		}
+
 		if err != nil {
 			self.Printf(" %v", err).Endl()
 			// In case of a VM exception (known exceptions) all gas consumed (panics NOT included).
@@ -647,7 +654,6 @@ func (self *Vm) Run(context *Context, callData []byte) (ret []byte, err error) {
 
 				self.Printf(" (*) 0x0 %v", suberr)
 			} else {
-
 				// gas < len(ret) * CreateDataGas == NO_CODE
 				dataGas := big.NewInt(int64(len(ret)))
 				dataGas.Mul(dataGas, GasCreateByte)