1 files changed, 37 insertions, 15 deletions

diff --git a/ethchain/contract.go b/ethchain/contract.go
index 21ac828fe..f7ae01753 100644
--- a/ethchain/contract.go
+++ b/ethchain/contract.go
@@ -9,26 +9,22 @@ type Contract struct {
 	Amount *big.Int
 	Nonce  uint64
 	//state  *ethutil.Trie
-	state *State
+	state   *State
+	address []byte
 }
 
-func NewContract(Amount *big.Int, root []byte) *Contract {
-	contract := &Contract{Amount: Amount, Nonce: 0}
+func NewContract(address []byte, Amount *big.Int, root []byte) *Contract {
+	contract := &Contract{address: address, Amount: Amount, Nonce: 0}
 	contract.state = NewState(ethutil.NewTrie(ethutil.Config.Db, string(root)))
 
 	return contract
 }
 
-func (c *Contract) RlpEncode() []byte {
-	return ethutil.Encode([]interface{}{c.Amount, c.Nonce, c.state.trie.Root})
-}
+func NewContractFromBytes(address, data []byte) *Contract {
+	contract := &Contract{address: address}
+	contract.RlpDecode(data)
 
-func (c *Contract) RlpDecode(data []byte) {
-	decoder := ethutil.NewValueFromBytes(data)
-
-	c.Amount = decoder.Get(0).BigInt()
-	c.Nonce = decoder.Get(1).Uint()
-	c.state = NewState(ethutil.NewTrie(ethutil.Config.Db, decoder.Get(2).Interface()))
+	return contract
 }
 
 func (c *Contract) Addr(addr []byte) *ethutil.Value {
@@ -43,19 +39,45 @@ func (c *Contract) State() *State {
 	return c.state
 }
 
-func (c *Contract) GetMem(num int) *ethutil.Value {
-	nb := ethutil.BigToBytes(big.NewInt(int64(num)), 256)
+func (c *Contract) GetMem(num *big.Int) *ethutil.Value {
+	nb := ethutil.BigToBytes(num, 256)
 
 	return c.Addr(nb)
 }
 
+func (c *Contract) SetMem(num *big.Int, val *ethutil.Value) {
+	addr := ethutil.BigToBytes(num, 256)
+	c.state.trie.Update(string(addr), string(val.Encode()))
+}
+
+// Return the gas back to the origin. Used by the Virtual machine or Closures
+func (c *Contract) ReturnGas(val *big.Int, state *State) {
+	c.Amount.Add(c.Amount, val)
+}
+
+func (c *Contract) Address() []byte {
+	return c.address
+}
+
+func (c *Contract) RlpEncode() []byte {
+	return ethutil.Encode([]interface{}{c.Amount, c.Nonce, c.state.trie.Root})
+}
+
+func (c *Contract) RlpDecode(data []byte) {
+	decoder := ethutil.NewValueFromBytes(data)
+
+	c.Amount = decoder.Get(0).BigInt()
+	c.Nonce = decoder.Get(1).Uint()
+	c.state = NewState(ethutil.NewTrie(ethutil.Config.Db, decoder.Get(2).Interface()))
+}
+
 func MakeContract(tx *Transaction, state *State) *Contract {
 	// Create contract if there's no recipient
 	if tx.IsContract() {
 		addr := tx.Hash()[12:]
 
 		value := tx.Value
-		contract := NewContract(value, []byte(""))
+		contract := NewContract(addr, value, []byte(""))
 		state.trie.Update(string(addr), string(contract.RlpEncode()))
 		for i, val := range tx.Data {
 			if len(val) > 0 {