Merge pull request ethereum#17 from OffchainLabs/reuse-recompiled

Changes Needed to Support Reusing Compiled Wasm Programs

core/rawdb/accessors_state_arbitrum.go

@@ -23,7 +23,7 @@ import (
 )
 
 // WriteCompiledWasmCode writes the provided contract compiled wasm code database.
-func WriteCompiledWasmCode(db ethdb.KeyValueWriter, version uint32, hash common.Hash, code []byte) {
+func WriteCompiledWasmCode(db ethdb.KeyValueWriter, version uint16, hash common.Hash, code []byte) {
 	key := CompiledWasmCodeKey(version, hash)
 	if err := db.Put(key[:], code); err != nil {
 		log.Crit("Failed to store compiled wasm contract code", "err", err)

core/rawdb/schema_arbitrum.go

@@ -30,28 +30,28 @@ var (
 )
 
 // CompiledWasmCodeKey = CompiledWasmCodePrefix + version + hash
-const WasmKeyLen = 2 + 4 + 32
+const WasmKeyLen = 2 + 2 + 32
 
 type WasmKey = [WasmKeyLen]byte
 
 // CompiledWasmCodeKey = CompiledWasmCodePrefix + version + hash
-func CompiledWasmCodeKey(version uint32, hash common.Hash) WasmKey {
+func CompiledWasmCodeKey(version uint16, hash common.Hash) WasmKey {
 	var key WasmKey
 	copy(key[:2], CompiledWasmCodePrefix)
-	binary.BigEndian.PutUint32(key[2:6], version)
-	copy(key[6:], hash[:])
+	binary.BigEndian.PutUint16(key[2:4], version)
+	copy(key[4:], hash[:])
 	return key
 }
 
 // IsCompiledWasmCodeKey reports whether the given byte slice is the key of compiled wasm contract code,
 // if so return the raw code hash and version as well.
-func IsCompiledWasmCodeKey(key []byte) (bool, common.Hash, uint32) {
+func IsCompiledWasmCodeKey(key []byte) (bool, common.Hash, uint16) {
 
 	start := len(CompiledWasmCodePrefix)
 
 	if bytes.HasPrefix(key, CompiledWasmCodePrefix) && len(key) == WasmKeyLen {
-		version := binary.BigEndian.Uint32(key[start : start+4])
-		codeHash := common.BytesToHash(key[start+4:])
+		version := binary.BigEndian.Uint16(key[start : start+2])
+		codeHash := common.BytesToHash(key[start+2:])
 		return true, codeHash, version
 	}
 	return false, common.Hash{}, 0

core/state/database.go

@@ -42,7 +42,9 @@ const (
 // Database wraps access to tries and contract code.
 type Database interface {
 	// Arbitrum: CompiledWasmContractCode retrieves a particular contract's user wasm code.
-	CompiledWasmContractCode(version uint32, codeHash common.Hash) ([]byte, error)
+	CompiledWasmContractCode(version uint16, codeHash common.Hash) ([]byte, error)
+	// Arbitrum: SetCompiledWasmContractCode sets a user wasm code.
+	SetCompiledWasmContractCode(version uint16, codeHash common.Hash, code []byte) error
 
 	// OpenTrie opens the main account trie.
 	OpenTrie(root common.Hash) (Trie, error)

core/state/database_arbitrum.go

@@ -7,8 +7,7 @@ import (
 	"github.com/ethereum/go-ethereum/core/rawdb"
 )
 
-// CompiledWasmContractCode retrieves a particular contract's compiled wasm code.
-func (db *cachingDB) CompiledWasmContractCode(version uint32, codeHash common.Hash) ([]byte, error) {
+func (db *cachingDB) CompiledWasmContractCode(version uint16, codeHash common.Hash) ([]byte, error) {
 	wasmKey := rawdb.CompiledWasmCodeKey(version, codeHash)
 	if code, _ := db.compiledWasmCache.Get(wasmKey); len(code) > 0 {
 		return code, nil
@@ -23,3 +22,12 @@ func (db *cachingDB) CompiledWasmContractCode(version uint32, codeHash common.Ha
 	}
 	return nil, errors.New("not found")
 }
+
+func (db *cachingDB) SetCompiledWasmContractCode(version uint16, codeHash common.Hash, code []byte) error {
+	wasmKey := rawdb.CompiledWasmCodeKey(version, codeHash)
+	if code, _ := db.compiledWasmCache.Get(wasmKey); len(code) > 0 {
+		return nil
+	}
+	db.compiledWasmCache.Add(wasmKey, code)
+	return nil
+}

core/state/journal_arbitrum.go

@@ -4,7 +4,7 @@ import "github.com/ethereum/go-ethereum/common"
 
 type wasmCodeChange struct {
 	account *common.Address
-	version uint32
+	version uint16
 }
 
 func (ch wasmCodeChange) revert(s *StateDB) {

core/state/state_object_arbitrum.go

@@ -26,7 +26,7 @@ type CompiledWasmCache struct {
 	code  Code
 	dirty bool
 }
-type CompiledWasms map[uint32]CompiledWasmCache
+type CompiledWasms map[uint16]CompiledWasmCache
 
 func (c CompiledWasms) Copy() CompiledWasms {
 	cpy := make(CompiledWasms, len(c))
@@ -38,7 +38,7 @@ func (c CompiledWasms) Copy() CompiledWasms {
 }
 
 // CompiledWasmCode returns the user wasm contract code associated with this object, if any.
-func (s *stateObject) CompiledWasmCode(db Database, version uint32) []byte {
+func (s *stateObject) CompiledWasmCode(db Database, version uint16) []byte {
 	if wasm, ok := s.compiledWasmCode[version]; ok {
 		return wasm.code
 	}
@@ -56,16 +56,19 @@ func (s *stateObject) CompiledWasmCode(db Database, version uint32) []byte {
 	return compiledWasmCode
 }
 
-func (s *stateObject) SetCompiledWasmCode(code []byte, version uint32) {
+func (s *stateObject) SetCompiledWasmCode(db Database, code []byte, version uint16) {
 	// Can only be compiled once, so if it's being compiled, it was previous empty
 	s.db.journal.append(wasmCodeChange{
 		account: &s.address,
 		version: version,
 	})
 	s.setWASMCode(code, version)
+	if err := db.SetCompiledWasmContractCode(version, common.BytesToHash(s.CodeHash()), code); err != nil {
+		s.db.setError(fmt.Errorf("cannot set compiled wasm contract code %x: %v", s.CodeHash(), err))
+	}
 }
 
-func (s *stateObject) setWASMCode(code []byte, version uint32) {
+func (s *stateObject) setWASMCode(code []byte, version uint16) {
 	s.compiledWasmCode[version] = CompiledWasmCache{
 		code:  code,
 		dirty: true,

core/state/statedb_arbitrum.go

@@ -37,9 +37,9 @@ var (
 	// This allows us to store WASM programs as code in the stateDB side-by-side
 	// with EVM contracts, but match against these prefix bytes when loading code
 	// to execute the WASMs through Stylus rather than the EVM.
-	stylusEOFMagic         = byte(0xEF)
-	stylusEOFMagicSuffix   = byte(0xF0)
-	stylusEOFVersion       = byte(0x00)
+	stylusEOFMagic       = byte(0xEF)
+	stylusEOFMagicSuffix = byte(0xF0)
+	stylusEOFVersion     = byte(0x00)
 
 	StylusPrefix = []byte{stylusEOFMagic, stylusEOFMagicSuffix, stylusEOFVersion}
 )
@@ -62,18 +62,18 @@ func StripStylusPrefix(b []byte) ([]byte, error) {
 	return b[3:], nil
 }
 
-func (s *StateDB) GetCompiledWasmCode(addr common.Address, version uint32) []byte {
+func (s *StateDB) GetCompiledWasmCode(addr common.Address, version uint16) []byte {
 	stateObject := s.getStateObject(addr)
 	if stateObject != nil {
 		return stateObject.CompiledWasmCode(s.db, version)
 	}
 	return nil
 }
 
-func (s *StateDB) SetCompiledWasmCode(addr common.Address, code []byte, version uint32) {
+func (s *StateDB) SetCompiledWasmCode(addr common.Address, code []byte, version uint16) {
 	stateObject := s.GetOrNewStateObject(addr)
 	if stateObject != nil {
-		stateObject.SetCompiledWasmCode(code, version)
+		stateObject.SetCompiledWasmCode(s.db, code, version)
 	}
 }
 
@@ -142,42 +142,54 @@ type UserWasm struct {
 	NoncanonicalHash common.Hash
 	CompressedWasm   []byte
 	Wasm             []byte
+	CodeHash         common.Hash // TODO: remove this field
 }
 type WasmCall struct {
-	Version uint32
-	Address common.Address
+	Version  uint16
+	CodeHash common.Hash
 }
 
 func (s *StateDB) StartRecording() {
 	s.userWasms = make(UserWasms)
 }
 
-func (s *StateDB) RecordProgram(program common.Address, version uint32) {
+func (s *StateDB) RecordProgram(program common.Address, codeHash common.Hash, version uint16) {
 	if s.userWasms != nil {
 		call := WasmCall{
-			Version: version,
-			Address: program,
+			Version:  version,
+			CodeHash: codeHash,
 		}
 		if _, ok := s.userWasms[call]; ok {
 			return
 		}
+		storedCodeHash := s.GetCodeHash(program)
+		if storedCodeHash != codeHash {
+			log.Error(
+				"Wrong recorded codehash for program at addr %#x, got codehash %#x in DB, specified codehash %#x to record",
+				program,
+				storedCodeHash,
+				codeHash,
+			)
+			return
+		}
 		rawCode := s.GetCode(program)
 		compressedWasm, err := StripStylusPrefix(rawCode)
 		if err != nil {
 			log.Error("Could not strip stylus program prefix from raw code: %v", err)
 			return
 		}
 		s.userWasms[call] = &UserWasm{
-			NoncanonicalHash: s.NoncanonicalProgramHash(program, version),
+			NoncanonicalHash: s.NoncanonicalProgramHash(codeHash, version),
 			CompressedWasm:   compressedWasm,
+			CodeHash:         codeHash,
 		}
 	}
 }
 
-func (s *StateDB) NoncanonicalProgramHash(program common.Address, version uint32) common.Hash {
-	prefix := make([]byte, 4)
-	binary.BigEndian.PutUint32(prefix, version)
-	return crypto.Keccak256Hash(prefix, s.GetCodeHash(program).Bytes())
+func (s *StateDB) NoncanonicalProgramHash(codeHash common.Hash, version uint16) common.Hash {
+	prefix := make([]byte, 2)
+	binary.BigEndian.PutUint16(prefix, version)
+	return crypto.Keccak256Hash(prefix, codeHash.Bytes())
 }
 
 func (s *StateDB) UserWasms() UserWasms {

core/vm/interface.go

@@ -28,8 +28,8 @@ import (
 // StateDB is an EVM database for full state querying.
 type StateDB interface {
 	// Arbitrum: manage compiled wasms
-	GetCompiledWasmCode(addr common.Address, version uint32) []byte
-	SetCompiledWasmCode(addr common.Address, code []byte, version uint32)
+	GetCompiledWasmCode(addr common.Address, version uint16) []byte
+	SetCompiledWasmCode(addr common.Address, code []byte, version uint16)
 
 	// Arbitrum: track stylus's memory footprint
 	GetStylusPages() (uint16, uint16)
@@ -38,7 +38,7 @@ type StateDB interface {
 	AddStylusPages(new uint16) (uint16, uint16)
 	AddStylusPagesEver(new uint16)
 
-	NoncanonicalProgramHash(common.Address, uint32) common.Hash
+	NoncanonicalProgramHash(codeHash common.Hash, version uint16) common.Hash
 	Deterministic() bool
 	Database() state.Database
 

light/trie_arbitrum.go

@@ -22,6 +22,10 @@ import (
 	"github.com/ethereum/go-ethereum/common"
 )
 
-func (db *odrDatabase) CompiledWasmContractCode(version uint32, codeHash common.Hash) ([]byte, error) {
+func (db *odrDatabase) CompiledWasmContractCode(version uint16, codeHash common.Hash) ([]byte, error) {
 	return nil, errors.New("retreiving compiled wasm not supported in light client")
 }
+
+func (db *odrDatabase) SetCompiledWasmContractCode(version uint16, codeHash common.Hash, code []byte) error {
+	return errors.New("setting compiled wasm not supported in light client")
+}