Improve validator code wrapper update

core/genesis.go

@@ -248,7 +248,7 @@ func (g *Genesis) ToBlock(db ethdb.Database) *types.Block {
 	statedb, _ := state.New(common.Hash{}, state.NewDatabase(db), nil)
 	for addr, account := range g.Alloc {
 		statedb.AddBalance(addr, account.Balance)
-		statedb.SetCode(addr, account.Code, false)
+		statedb.SetCode(addr, account.Code)
 		statedb.SetNonce(addr, account.Nonce)
 		for key, value := range account.Storage {
 			statedb.SetState(addr, key, value)

core/state/iterator.go

@@ -120,14 +120,14 @@ func (it *NodeIterator) step() error {
 	if !bytes.Equal(account.CodeHash, types.EmptyCodeHash.Bytes()) {
 		it.codeHash = common.BytesToHash(account.CodeHash)
 		addrHash := common.BytesToHash(it.stateIt.LeafKey())
-		it.code, err = it.state.db.ContractCode(addrHash, common.BytesToHash(account.CodeHash))
-		if err != nil {
-			return fmt.Errorf("code %x: %v", account.CodeHash, err)
-		}
-		if it.code == nil || len(it.code) == 0 {
-			it.code, err = it.state.db.ValidatorCode(addrHash, common.BytesToHash(account.CodeHash))
-			if err != nil {
-				return fmt.Errorf("code %x: %v", account.CodeHash, err)
+
+		if code, errCC := it.state.db.ContractCode(addrHash, common.BytesToHash(account.CodeHash)); code != nil && len(code) > 0 && errCC == nil {
+			it.code = code
+		} else {
+			if code, errVC := it.state.db.ValidatorCode(addrHash, common.BytesToHash(account.CodeHash)); code != nil && len(code) > 0 && errVC == nil {
+				it.code = code
+			} else {
+				return fmt.Errorf("code %x: contract code error: %v, validator code error: %v", account.CodeHash, errCC, errVC)
 			}
 		}
 	}

core/state/iterator_test.go

@@ -54,7 +54,7 @@ func makeTestState() (ethdb.Database, Database, common.Hash, []*testAccount) {
 		acc.nonce = uint64(42 * i)
 
 		if i%3 == 0 {
-			obj.SetCode(crypto.Keccak256Hash([]byte{i, i, i, i, i}), []byte{i, i, i, i, i}, false)
+			obj.SetCode(crypto.Keccak256Hash([]byte{i, i, i, i, i}), []byte{i, i, i, i, i})
 			acc.code = []byte{i, i, i, i, i}
 		}
 		if i%5 == 0 {

core/state/journal.go

@@ -218,7 +218,7 @@ func (ch nonceChange) dirtied() *common.Address {
 }
 
 func (ch codeChange) revert(s *DB) {
-	s.getStateObject(*ch.account).setCode(common.BytesToHash(ch.prevhash), ch.prevcode, false)
+	s.getStateObject(*ch.account).setCode(common.BytesToHash(ch.prevhash), ch.prevcode)
 }
 
 func (ch codeChange) dirtied() *common.Address {

core/state/state_object.go

@@ -99,10 +99,9 @@ type Object struct {
 	// Cache flags.
 	// When an object is marked suicided it will be delete from the trie
 	// during the "update" phase of the state transition.
-	validatorWrapper bool // true if the code belongs to validator wrapper
-	dirtyCode        bool // true if the code was updated
-	suicided         bool
-	deleted          bool
+	dirtyCode bool // true if the code was updated
+	suicided  bool
+	deleted   bool
 }
 
 // empty returns whether the account is considered empty.
@@ -514,24 +513,23 @@ func (s *Object) Code(db Database) []byte {
 	}
 	var err error
 	code := []byte{}
-	// if it's not set for validator wrapper, then it may be either contract code or validator wrapper (old version of db
-	// don't have any prefix to differentiate between them)
-	// so, if it's not set for validator wrapper, we need to check contract code as well
-	if !s.validatorWrapper {
+
+	isValidatorWrapper := s.IsValidator(db)
+	if !isValidatorWrapper {
 		code, err = db.ContractCode(s.addrHash, common.BytesToHash(s.CodeHash()))
-	}
-	// if it couldn't load contract code or it is set to validator wrapper, then it tries to fetch validator wrapper code
-	if s.validatorWrapper || err != nil {
-		vCode, errVCode := db.ValidatorCode(s.addrHash, common.BytesToHash(s.CodeHash()))
-		if errVCode == nil && vCode != nil {
-			s.code = vCode
-			return vCode
+		if err != nil {
+			s.setError(fmt.Errorf("can't load contract code hash %x for account address hash %x : contract code error: %v",
+				s.CodeHash(), s.addrHash, err))
 		}
-		if s.validatorWrapper {
-			s.setError(fmt.Errorf("can't load validator code hash %x for account address hash %x : %v", s.CodeHash(), s.addrHash, err))
-		} else {
-			s.setError(fmt.Errorf("can't load contract/validator code hash %x for account address hash %x : contract code error: %v, validator code error: %v",
-				s.CodeHash(), s.addrHash, err, errVCode))
+	} else {
+		code, err = db.ValidatorCode(s.addrHash, common.BytesToHash(s.CodeHash()))
+		if err != nil {
+			// old version of db don't have any prefix to differentiate between them
+			// so, if it's not set for validator wrapper, we need to check contract code as well
+			code, err = db.ContractCode(s.addrHash, common.BytesToHash(s.CodeHash()))
+			if err != nil {
+				s.setError(fmt.Errorf("can't load validator/contract code hash %x for account address hash %x : %v", s.CodeHash(), s.addrHash, err))
+			}
 		}
 	}
 	s.code = code
@@ -551,44 +549,42 @@ func (s *Object) CodeSize(db Database) int {
 	var err error
 	size := int(0)
 
-	// if it's not set for validator wrapper, then it may be either contract code or validator wrapper (old version of db
-	// don't have any prefix to differentiate between them)
-	// so, if it's not set for validator wrapper, we need to check contract code as well
-	if !s.validatorWrapper {
+	isValidatorWrapper := s.IsValidator(db)
+	if !isValidatorWrapper {
 		size, err = db.ContractCodeSize(s.addrHash, common.BytesToHash(s.CodeHash()))
-	}
-	// if it couldn't get contract code or it is set to validator wrapper, then it tries to retrieve validator wrapper code
-	if s.validatorWrapper || err != nil {
-		vcSize, errVCSize := db.ValidatorCodeSize(s.addrHash, common.BytesToHash(s.CodeHash()))
-		if errVCSize == nil && vcSize > 0 {
-			return vcSize
+		if err != nil {
+			s.setError(fmt.Errorf("can't load contract code size %x for account address hash %x : contract code size error: %v",
+				s.CodeHash(), s.addrHash, err))
 		}
-		if s.validatorWrapper {
-			s.setError(fmt.Errorf("can't load validator code size %x for account address hash %x : %v", s.CodeHash(), s.addrHash, err))
-		} else {
-			s.setError(fmt.Errorf("can't load contract/validator code size %x for account address hash %x : contract code size error: %v, validator code size error: %v",
-				s.CodeHash(), s.addrHash, err, errVCSize))
+	} else {
+		size, err = db.ValidatorCodeSize(s.addrHash, common.BytesToHash(s.CodeHash()))
+		if err != nil {
+			// old version of db don't have any prefix to differentiate between them
+			// so, if it's not set for validator wrapper, we need to check contract code as well
+			size, err = db.ContractCodeSize(s.addrHash, common.BytesToHash(s.CodeHash()))
+			if err != nil {
+				s.setError(fmt.Errorf("can't load validator/contract code size %x for account address hash %x : code error: %v", s.CodeHash(), s.addrHash, err))
+			}
 		}
-		s.setError(fmt.Errorf("can't load code size %x (validator wrapper: %t): %v", s.CodeHash(), s.validatorWrapper, err))
 	}
+
 	return size
 }
 
-func (s *Object) SetCode(codeHash common.Hash, code []byte, isValidatorCode bool) {
+func (s *Object) SetCode(codeHash common.Hash, code []byte) {
 	prevcode := s.Code(s.db.db)
 	s.db.journal.append(codeChange{
 		account:  &s.address,
 		prevhash: s.CodeHash(),
 		prevcode: prevcode,
 	})
-	s.setCode(codeHash, code, isValidatorCode)
+	s.setCode(codeHash, code)
 }
 
-func (s *Object) setCode(codeHash common.Hash, code []byte, isValidatorCode bool) {
+func (s *Object) setCode(codeHash common.Hash, code []byte) {
 	s.code = code
 	s.data.CodeHash = codeHash[:]
 	s.dirtyCode = true
-	s.validatorWrapper = isValidatorCode
 }
 
 func (s *Object) SetNonce(nonce uint64) {

core/state/state_test.go

@@ -48,7 +48,7 @@ func TestDump(t *testing.T) {
 	obj1 := s.state.GetOrNewStateObject(common.BytesToAddress([]byte{0x01}))
 	obj1.AddBalance(big.NewInt(22))
 	obj2 := s.state.GetOrNewStateObject(common.BytesToAddress([]byte{0x01, 0x02}))
-	obj2.SetCode(crypto.Keccak256Hash([]byte{3, 3, 3, 3, 3, 3, 3}), []byte{3, 3, 3, 3, 3, 3, 3}, false)
+	obj2.SetCode(crypto.Keccak256Hash([]byte{3, 3, 3, 3, 3, 3, 3}), []byte{3, 3, 3, 3, 3, 3, 3})
 	obj3 := s.state.GetOrNewStateObject(common.BytesToAddress([]byte{0x02}))
 	obj3.SetBalance(big.NewInt(44))
 
@@ -166,7 +166,7 @@ func TestSnapshot2(t *testing.T) {
 	so0 := state.getStateObject(stateobjaddr0)
 	so0.SetBalance(big.NewInt(42))
 	so0.SetNonce(43)
-	so0.SetCode(crypto.Keccak256Hash([]byte{'c', 'a', 'f', 'e'}), []byte{'c', 'a', 'f', 'e'}, false)
+	so0.SetCode(crypto.Keccak256Hash([]byte{'c', 'a', 'f', 'e'}), []byte{'c', 'a', 'f', 'e'})
 	so0.suicided = false
 	so0.deleted = false
 	state.setStateObject(so0)
@@ -178,7 +178,7 @@ func TestSnapshot2(t *testing.T) {
 	so1 := state.getStateObject(stateobjaddr1)
 	so1.SetBalance(big.NewInt(52))
 	so1.SetNonce(53)
-	so1.SetCode(crypto.Keccak256Hash([]byte{'c', 'a', 'f', 'e', '2'}), []byte{'c', 'a', 'f', 'e', '2'}, false)
+	so1.SetCode(crypto.Keccak256Hash([]byte{'c', 'a', 'f', 'e', '2'}), []byte{'c', 'a', 'f', 'e', '2'})
 	so1.suicided = true
 	so1.deleted = true
 	state.setStateObject(so1)

core/state/statedb.go

@@ -475,10 +475,10 @@ func (db *DB) SetNonce(addr common.Address, nonce uint64) {
 	}
 }
 
-func (db *DB) SetCode(addr common.Address, code []byte, isValidatorCode bool) {
+func (db *DB) SetCode(addr common.Address, code []byte) {
 	Object := db.GetOrNewStateObject(addr)
 	if Object != nil {
-		Object.SetCode(crypto.Keccak256Hash(code), code, isValidatorCode)
+		Object.SetCode(crypto.Keccak256Hash(code), code)
 	}
 }
 
@@ -1053,7 +1053,8 @@ func (db *DB) Commit(deleteEmptyObjects bool) (common.Hash, error) {
 		if obj := db.stateObjects[addr]; !obj.deleted {
 			// Write any contract code associated with the state object
 			if obj.code != nil && obj.dirtyCode {
-				if obj.validatorWrapper {
+				isValidator := obj.IsValidator(db.db)
+				if isValidator {
 					rawdb.WriteValidatorCode(codeWriter, common.BytesToHash(obj.CodeHash()), obj.code)
 				} else {
 					rawdb.WriteCode(codeWriter, common.BytesToHash(obj.CodeHash()), obj.code)
@@ -1290,7 +1291,7 @@ func (db *DB) UpdateValidatorWrapper(
 		return err
 	}
 	// has revert in-built for the code field
-	db.SetCode(addr, by, true)
+	db.SetCode(addr, by)
 	// update cache
 	db.stateValidators[addr] = val
 	return nil

core/state/statedb_test.go

@@ -50,7 +50,7 @@ func TestUpdateLeaks(t *testing.T) {
 			state.SetState(addr, common.BytesToHash([]byte{i, i, i}), common.BytesToHash([]byte{i, i, i, i}))
 		}
 		if i%3 == 0 {
-			state.SetCode(addr, []byte{i, i, i, i, i}, false)
+			state.SetCode(addr, []byte{i, i, i, i, i})
 		}
 	}
 
@@ -84,7 +84,7 @@ func TestIntermediateLeaks(t *testing.T) {
 			state.SetState(addr, common.Hash{i, i, i, tweak}, common.Hash{i, i, i, i, tweak})
 		}
 		if i%3 == 0 {
-			state.SetCode(addr, []byte{i, i, i, i, i, tweak}, false)
+			state.SetCode(addr, []byte{i, i, i, i, i, tweak})
 		}
 	}
 
@@ -286,7 +286,7 @@ func newTestAction(addr common.Address, r *rand.Rand) testAction {
 				code := make([]byte, 16)
 				binary.BigEndian.PutUint64(code, uint64(a.args[0]))
 				binary.BigEndian.PutUint64(code[8:], uint64(a.args[1]))
-				s.SetCode(addr, code, false)
+				s.SetCode(addr, code)
 			},
 			args: make([]int64, 2),
 		},
@@ -525,9 +525,9 @@ func TestCopyCommitCopy(t *testing.T) {
 	skey := common.HexToHash("aaa")
 	sval := common.HexToHash("bbb")
 
-	state.SetBalance(addr, big.NewInt(42))      // Change the account trie
-	state.SetCode(addr, []byte("hello"), false) // Change an external metadata
-	state.SetState(addr, skey, sval)            // Change the storage trie
+	state.SetBalance(addr, big.NewInt(42)) // Change the account trie
+	state.SetCode(addr, []byte("hello"))   // Change an external metadata
+	state.SetState(addr, skey, sval)       // Change the storage trie
 
 	if balance := state.GetBalance(addr); balance.Cmp(big.NewInt(42)) != 0 {
 		t.Fatalf("initial balance mismatch: have %v, want %v", balance, 42)
@@ -597,9 +597,9 @@ func TestCopyCopyCommitCopy(t *testing.T) {
 	skey := common.HexToHash("aaa")
 	sval := common.HexToHash("bbb")
 
-	state.SetBalance(addr, big.NewInt(42))      // Change the account trie
-	state.SetCode(addr, []byte("hello"), false) // Change an external metadata
-	state.SetState(addr, skey, sval)            // Change the storage trie
+	state.SetBalance(addr, big.NewInt(42)) // Change the account trie
+	state.SetCode(addr, []byte("hello"))   // Change an external metadata
+	state.SetState(addr, skey, sval)       // Change the storage trie
 
 	if balance := state.GetBalance(addr); balance.Cmp(big.NewInt(42)) != 0 {
 		t.Fatalf("initial balance mismatch: have %v, want %v", balance, 42)
@@ -717,10 +717,10 @@ func TestMissingTrieNodes(t *testing.T) {
 	addr := common.BytesToAddress([]byte("so"))
 	{
 		state.SetBalance(addr, big.NewInt(1))
-		state.SetCode(addr, []byte{1, 2, 3}, false)
+		state.SetCode(addr, []byte{1, 2, 3})
 		a2 := common.BytesToAddress([]byte("another"))
 		state.SetBalance(a2, big.NewInt(100))
-		state.SetCode(a2, []byte{1, 2, 4}, false)
+		state.SetCode(a2, []byte{1, 2, 4})
 		root, _ = state.Commit(false)
 		t.Logf("root: %x", root)
 		// force-flush

core/state/trie_prefetcher_test.go

@@ -33,9 +33,9 @@ func filledStateDB() *DB {
 	skey := common.HexToHash("aaa")
 	sval := common.HexToHash("bbb")
 
-	state.SetBalance(addr, big.NewInt(42))      // Change the account trie
-	state.SetCode(addr, []byte("hello"), false) // Change an external metadata
-	state.SetState(addr, skey, sval)            // Change the storage trie
+	state.SetBalance(addr, big.NewInt(42)) // Change the account trie
+	state.SetCode(addr, []byte("hello"))   // Change an external metadata
+	state.SetState(addr, skey, sval)       // Change the storage trie
 	for i := 0; i < 100; i++ {
 		sk := common.BigToHash(big.NewInt(int64(i)))
 		state.SetState(addr, sk, sk) // Change the storage trie

core/vm/evm.go

@@ -553,7 +553,7 @@ func (evm *EVM) create(caller ContractRef, codeAndHash *codeAndHash, gas uint64,
 	if err == nil && !maxCodeSizeExceeded {
 		createDataGas := uint64(len(ret)) * params.CreateDataGas
 		if contract.UseGas(createDataGas) {
-			evm.StateDB.SetCode(address, ret, false)
+			evm.StateDB.SetCode(address, ret)
 		} else {
 			err = ErrCodeStoreOutOfGas
 		}

core/vm/gas_table_test.go

@@ -85,7 +85,7 @@ func TestEIP2200(t *testing.T) {
 
 		statedb, _ := state.New(common.Hash{}, state.NewDatabase(rawdb.NewMemoryDatabase()), nil)
 		statedb.CreateAccount(address)
-		statedb.SetCode(address, hexutil.MustDecode(tt.input), false)
+		statedb.SetCode(address, hexutil.MustDecode(tt.input))
 		statedb.SetState(address, common.Hash{}, common.BytesToHash([]byte{tt.original}))
 		statedb.Finalise(true) // Push the state into the "original" slot
 

core/vm/interface.go

@@ -39,7 +39,7 @@ type StateDB interface {
 
 	GetCodeHash(common.Address) common.Hash
 	GetCode(common.Address) []byte
-	SetCode(common.Address, []byte, bool)
+	SetCode(common.Address, []byte)
 	GetCodeSize(common.Address) int
 
 	ValidatorWrapper(common.Address, bool, bool) (*staking.ValidatorWrapper, error)

core/vm/logger_test.go

@@ -28,10 +28,10 @@ type dummyContractRef struct {
 	calledForEach bool
 }
 
-func (dummyContractRef) ReturnGas(*big.Int)                {}
-func (dummyContractRef) Address() common.Address           { return common.Address{} }
-func (dummyContractRef) Value() *big.Int                   { return new(big.Int) }
-func (dummyContractRef) SetCode(common.Hash, []byte, bool) {}
+func (dummyContractRef) ReturnGas(*big.Int)          {}
+func (dummyContractRef) Address() common.Address     { return common.Address{} }
+func (dummyContractRef) Value() *big.Int             { return new(big.Int) }
+func (dummyContractRef) SetCode(common.Hash, []byte) {}
 func (d *dummyContractRef) ForEachStorage(callback func(key, value common.Hash) bool) {
 	d.calledForEach = true
 }

core/vm/runtime/runtime.go

@@ -112,7 +112,7 @@ func Execute(code, input []byte, cfg *Config) ([]byte, *state.DB, error) {
 	)
 	cfg.State.CreateAccount(address)
 	// set the receiver's (the executing contract) code for execution.
-	cfg.State.SetCode(address, code, false)
+	cfg.State.SetCode(address, code)
 	// Call the code with the given configuration.
 	ret, _, err := vmenv.Call(
 		sender,

core/vm/runtime/runtime_test.go

@@ -105,7 +105,7 @@ func TestCall(t *testing.T) {
 		byte(vm.PUSH1), 32,
 		byte(vm.PUSH1), 0,
 		byte(vm.RETURN),
-	}, false)
+	})
 
 	ret, _, err := Call(address, nil, &Config{State: state})
 	if err != nil {
@@ -158,7 +158,7 @@ func benchmarkEVMCreate(bench *testing.B, code string) {
 	)
 
 	statedb.CreateAccount(sender)
-	statedb.SetCode(receiver, common.FromHex(code), false)
+	statedb.SetCode(receiver, common.FromHex(code))
 	runtimeConfig := Config{
 		Origin:      sender,
 		State:       statedb,

staking/availability/measure_test.go

@@ -434,8 +434,8 @@ func (ctx *incStateTestCtx) checkAddrIncStateByType(addr common.Address, typeInc
 // checkHmyNodeStateChangeByAddr checks the state change for hmy nodes. Since hmy nodes does not
 // have wrapper, it is supposed to be unchanged in code field
 func (ctx *incStateTestCtx) checkHmyNodeStateChangeByAddr(addr common.Address) error {
-	snapCode := ctx.snapState.GetCode(addr, false)
-	curCode := ctx.state.GetCode(addr, false)
+	snapCode := ctx.snapState.GetCode(addr)
+	curCode := ctx.state.GetCode(addr)
 	if !reflect.DeepEqual(snapCode, curCode) {
 		return errors.New("code not expected")
 	}
@@ -618,7 +618,7 @@ func (state testStateDB) UpdateValidatorWrapper(addr common.Address, wrapper *st
 	return nil
 }
 
-func (state testStateDB) GetCode(addr common.Address, isValidatorCode bool) []byte {
+func (state testStateDB) GetCode(addr common.Address) []byte {
 	wrapper, ok := state[addr]
 	if !ok {
 		return nil