Cleanup VM

5 files changed, 40 insertions, 29 deletions

diff --git a/core/execution.go b/core/execution.go
index 72eb22bd5..24e085e6d 100644
--- a/core/execution.go
+++ b/core/execution.go
@@ -11,14 +11,18 @@ import (
 )
 
 type Execution struct {
-	env               vm.Environment
-	address           *common.Address
-	input             []byte
+	env     vm.Environment
+	address *common.Address
+	input   []byte
+	evm     vm.VirtualMachine
+
 	Gas, price, value *big.Int
 }
 
 func NewExecution(env vm.Environment, address *common.Address, input []byte, gas, gasPrice, value *big.Int) *Execution {
-	return &Execution{env: env, address: address, input: input, Gas: gas, price: gasPrice, value: value}
+	exe := &Execution{env: env, address: address, input: input, Gas: gas, price: gasPrice, value: value}
+	exe.evm = vm.NewVm(env)
+	return exe
 }
 
 func (self *Execution) Call(codeAddr common.Address, caller vm.ContextRef) ([]byte, error) {
@@ -28,11 +32,17 @@ func (self *Execution) Call(codeAddr common.Address, caller vm.ContextRef) ([]by
 	return self.exec(&codeAddr, code, caller)
 }
 
+func (self *Execution) Create(caller vm.ContextRef) (ret []byte, err error, account *state.StateObject) {
+	ret, err = self.exec(nil, self.input, caller)
+	account = self.env.State().GetStateObject(*self.address)
+	return
+}
+
 func (self *Execution) exec(contextAddr *common.Address, code []byte, caller vm.ContextRef) (ret []byte, err error) {
 	start := time.Now()
 
 	env := self.env
-	evm := vm.NewVm(env)
+	evm := self.evm
 	if env.Depth() == vm.MaxCallDepth {
 		caller.ReturnGas(self.Gas, self.price)
 
@@ -70,10 +80,3 @@ func (self *Execution) exec(contextAddr *common.Address, code []byte, caller vm.
 
 	return
 }
-
-func (self *Execution) Create(caller vm.ContextRef) (ret []byte, err error, account *state.StateObject) {
-	ret, err = self.exec(nil, self.input, caller)
-	account = self.env.State().GetStateObject(*self.address)
-
-	return
-}
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)
diff --git a/core/vm_env.go b/core/vm_env.go
index 52e8b20a9..6a604fccd 100644
--- a/core/vm_env.go
+++ b/core/vm_env.go
@@ -54,21 +54,17 @@ func (self *VMEnv) Transfer(from, to vm.Account, amount *big.Int) error {
 	return vm.Transfer(from, to, amount)
 }
 
-func (self *VMEnv) vm(addr *common.Address, data []byte, gas, price, value *big.Int) *Execution {
-	return NewExecution(self, addr, data, gas, price, value)
-}
-
 func (self *VMEnv) Call(me vm.ContextRef, addr common.Address, data []byte, gas, price, value *big.Int) ([]byte, error) {
-	exe := self.vm(&addr, data, gas, price, value)
+	exe := NewExecution(self, &addr, data, gas, price, value)
 	return exe.Call(addr, me)
 }
 func (self *VMEnv) CallCode(me vm.ContextRef, addr common.Address, data []byte, gas, price, value *big.Int) ([]byte, error) {
 	maddr := me.Address()
-	exe := self.vm(&maddr, data, gas, price, value)
+	exe := NewExecution(self, &maddr, data, gas, price, value)
 	return exe.Call(addr, me)
 }
 
 func (self *VMEnv) Create(me vm.ContextRef, data []byte, gas, price, value *big.Int) ([]byte, error, vm.ContextRef) {
-	exe := self.vm(nil, data, gas, price, value)
+	exe := NewExecution(self, nil, data, gas, price, value)
 	return exe.Create(me)
 }