1 files changed, 24 insertions, 9 deletions
diff --git a/core/vm/contract.go b/core/vm/contract.go
index 091106d84..dfa93ab18 100644
--- a/core/vm/contract.go
+++ b/core/vm/contract.go
@@ -24,6 +24,7 @@ import (
 
 // ContractRef is a reference to the contract's backing object
 type ContractRef interface {
+	ReturnGas(*big.Int)
 	Address() common.Address
 	Value() *big.Int
 	SetCode(common.Hash, []byte)
@@ -47,8 +48,7 @@ type Contract struct {
 	CodeAddr *common.Address
 	Input    []byte
 
-	Gas   uint64
-	value *big.Int
+	value, Gas, UsedGas *big.Int
 
 	Args []byte
 
@@ -56,7 +56,7 @@ type Contract struct {
 }
 
 // NewContract returns a new contract environment for the execution of EVM.
-func NewContract(caller ContractRef, object ContractRef, value *big.Int, gas uint64) *Contract {
+func NewContract(caller ContractRef, object ContractRef, value, gas *big.Int) *Contract {
 	c := &Contract{CallerAddress: caller.Address(), caller: caller, self: object, Args: nil}
 
 	if parent, ok := caller.(*Contract); ok {
@@ -68,8 +68,9 @@ func NewContract(caller ContractRef, object ContractRef, value *big.Int, gas uin
 
 	// Gas should be a pointer so it can safely be reduced through the run
 	// This pointer will be off the state transition
-	c.Gas = gas
+	c.Gas = gas //new(big.Int).Set(gas)
 	c.value = new(big.Int).Set(value)
+	c.UsedGas = new(big.Int)
 
 	return c
 }
@@ -106,13 +107,27 @@ func (c *Contract) Caller() common.Address {
 	return c.CallerAddress
 }
 
+// Finalise finalises the contract and returning any remaining gas to the original
+// caller.
+func (c *Contract) Finalise() {
+	// Return the remaining gas to the caller
+	c.caller.ReturnGas(c.Gas)
+}
+
 // UseGas attempts the use gas and subtracts it and returns true on success
-func (c *Contract) UseGas(gas uint64) (ok bool) {
-	if c.Gas < gas {
-		return false
+func (c *Contract) UseGas(gas *big.Int) (ok bool) {
+	ok = useGas(c.Gas, gas)
+	if ok {
+		c.UsedGas.Add(c.UsedGas, gas)
 	}
-	c.Gas -= gas
-	return true
+	return
+}
+
+// ReturnGas adds the given gas back to itself.
+func (c *Contract) ReturnGas(gas *big.Int) {
+	// Return the gas to the context
+	c.Gas.Add(c.Gas, gas)
+	c.UsedGas.Sub(c.UsedGas, gas)
 }
 
 // Address returns the contracts address