feat: state-key transformation w/ override (#205)

## Why this should be merged

`ava-labs/coreth` has a partitioned state-address space, achieved by
setting or clearing a specific bit in the hash used to key the space.
This change allows such behaviour to be achieved with pure `libevm`
instead of the `StateDB` wrapping that `coreth` currently uses.

## How this works

Introduction of `state.StateDBHooks` interface, including a
`TransformStateKey()` method that allows for arbitrary change of state
key. If registered, this hook will be honoured by
`StateDB.{Get,GetCommitted,State}Key()` methods unless they receive a
`stateconf.SkipStateKeyTransformation` option.

## How this was tested

Unit test of `SetState() -> GetState() + GetCommittedState()` round trip
with and without options to skip.

Author: ARR4N

core/state/statedb.go

@@ -28,14 +28,16 @@ import (
 	"github.com/ava-labs/libevm/core/state/snapshot"
 	"github.com/ava-labs/libevm/core/types"
 	"github.com/ava-labs/libevm/crypto"
-	"github.com/ava-labs/libevm/libevm/stateconf"
 	"github.com/ava-labs/libevm/log"
 	"github.com/ava-labs/libevm/metrics"
 	"github.com/ava-labs/libevm/params"
 	"github.com/ava-labs/libevm/trie"
 	"github.com/ava-labs/libevm/trie/trienode"
 	"github.com/ava-labs/libevm/trie/triestate"
 	"github.com/holiman/uint256"
+
+	// libevm extra imports
+	"github.com/ava-labs/libevm/libevm/stateconf"
 )
 
 const (
@@ -341,18 +343,20 @@ func (s *StateDB) GetCodeHash(addr common.Address) common.Hash {
 }
 
 // GetState retrieves a value from the given account's storage trie.
-func (s *StateDB) GetState(addr common.Address, hash common.Hash) common.Hash {
+func (s *StateDB) GetState(addr common.Address, hash common.Hash, opts ...stateconf.StateDBStateOption) common.Hash {
 	stateObject := s.getStateObject(addr)
 	if stateObject != nil {
+		hash = transformStateKey(addr, hash, opts...)
 		return stateObject.GetState(hash)
 	}
 	return common.Hash{}
 }
 
 // GetCommittedState retrieves a value from the given account's committed storage trie.
-func (s *StateDB) GetCommittedState(addr common.Address, hash common.Hash) common.Hash {
+func (s *StateDB) GetCommittedState(addr common.Address, hash common.Hash, opts ...stateconf.StateDBStateOption) common.Hash {
 	stateObject := s.getStateObject(addr)
 	if stateObject != nil {
+		hash = transformStateKey(addr, hash, opts...)
 		return stateObject.GetCommittedState(hash)
 	}
 	return common.Hash{}
@@ -412,9 +416,10 @@ func (s *StateDB) SetCode(addr common.Address, code []byte) {
 	}
 }
 
-func (s *StateDB) SetState(addr common.Address, key, value common.Hash) {
+func (s *StateDB) SetState(addr common.Address, key, value common.Hash, opts ...stateconf.StateDBStateOption) {
 	stateObject := s.getOrNewStateObject(addr)
 	if stateObject != nil {
+		key = transformStateKey(addr, key, opts...)
 		stateObject.SetState(key, value)
 	}
 }

core/state/statedb.libevm.go

@@ -21,6 +21,7 @@ import (
 
 	"github.com/ava-labs/libevm/common"
 	"github.com/ava-labs/libevm/core/state/snapshot"
+	"github.com/ava-labs/libevm/libevm/register"
 	"github.com/ava-labs/libevm/libevm/stateconf"
 )
 
@@ -57,3 +58,46 @@ func clearTypedNilPointer(snaps SnapshotTree) SnapshotTree {
 	}
 	return snaps
 }
+
+// StateDBHooks modify the behaviour of [StateDB] instances.
+type StateDBHooks interface {
+	// TransformStateKey receives the arguments passed to [StateDB.GetState],
+	// [StateDB.GetCommittedState] or [StateDB.SetState], and returns the key
+	// that each of those methods will use for accessing state. This method will
+	// not, however, be called if any of the aforementioned [StateDB] methods
+	// receives a [stateconf.SkipStateKeyTransformation] option.
+	//
+	// This method SHOULD NOT be used for anything other than achieving
+	// backwards compatibility with an existing chain. In the event that other
+	// methods are added to the [StateDBHooks] interface and no key
+	// transformation is required, it is acceptable for this method to echo the
+	// [common.Hash], unchanged.
+	TransformStateKey(_ common.Address, key common.Hash) (newKey common.Hash)
+}
+
+// RegisterExtras registers the [StateDBHooks] such that they modify the
+// behaviour of all [StateDB] instances. It is expected to be called in an
+// `init()` function and MUST NOT be called more than once.
+func RegisterExtras(s StateDBHooks) {
+	registeredExtras.MustRegister(s)
+}
+
+// TestOnlyClearRegisteredExtras clears the arguments previously passed to
+// [RegisterExtras]. It panics if called from a non-testing call stack.
+//
+// In tests it SHOULD be called before every call to [RegisterExtras] and then
+// defer-called afterwards, either directly or via testing.TB.Cleanup(). This is
+// a workaround for the single-call limitation on [RegisterExtras].
+func TestOnlyClearRegisteredExtras() {
+	registeredExtras.TestOnlyClear()
+}
+
+var registeredExtras register.AtMostOnce[StateDBHooks]
+
+func transformStateKey(addr common.Address, key common.Hash, opts ...stateconf.StateDBStateOption) common.Hash {
+	r := &registeredExtras
+	if !r.Registered() || !stateconf.ShouldTransformStateKey(opts...) {
+		return key
+	}
+	return r.Get().TransformStateKey(addr, key)
+}

core/state/statedb.libevm_test.go

@@ -138,3 +138,81 @@ func (r *triedbRecorder) Update(
 func (r *triedbRecorder) Reader(_ common.Hash) (database.Reader, error) {
 	return r.Database.Reader(common.Hash{})
 }
+
+type highByteFlipper struct{}
+
+func flipHighByte(h common.Hash) common.Hash {
+	h[0] = ^h[0]
+	return h
+}
+
+func (highByteFlipper) TransformStateKey(_ common.Address, key common.Hash) common.Hash {
+	return flipHighByte(key)
+}
+
+func TestTransformStateKey(t *testing.T) {
+	rawdb := rawdb.NewMemoryDatabase()
+	trie := triedb.NewDatabase(rawdb, nil)
+	db := NewDatabaseWithNodeDB(rawdb, trie)
+	sdb, err := New(types.EmptyRootHash, db, nil)
+	require.NoErrorf(t, err, "New()")
+
+	addr := common.Address{1}
+	regularKey := common.Hash{0, 'k', 'e', 'y'}
+	flippedKey := flipHighByte(regularKey)
+	regularVal := common.Hash{'r', 'e', 'g', 'u', 'l', 'a', 'r'}
+	flippedVal := common.Hash{'f', 'l', 'i', 'p', 'p', 'e', 'd'}
+
+	sdb.SetState(addr, regularKey, regularVal)
+	sdb.SetState(addr, flippedKey, flippedVal)
+
+	assertEq := func(t *testing.T, key, want common.Hash, opts ...stateconf.StateDBStateOption) {
+		t.Helper()
+		assert.Equal(t, want, sdb.GetState(addr, key, opts...))
+	}
+
+	assertEq(t, regularKey, regularVal)
+	assertEq(t, flippedKey, flippedVal)
+
+	root, err := sdb.Commit(0, false)
+	require.NoErrorf(t, err, "state.Commit()")
+
+	err = trie.Commit(root, false)
+	require.NoErrorf(t, err, "trie.Commit()")
+
+	sdb, err = New(root, db, nil)
+	require.NoErrorf(t, err, "New()")
+
+	assertCommittedEq := func(t *testing.T, key, want common.Hash, opts ...stateconf.StateDBStateOption) {
+		t.Helper()
+		assert.Equal(t, want, sdb.GetCommittedState(addr, key, opts...))
+	}
+
+	assertEq(t, regularKey, regularVal)
+	assertEq(t, flippedKey, flippedVal)
+	assertCommittedEq(t, regularKey, regularVal)
+	assertCommittedEq(t, flippedKey, flippedVal)
+
+	// Typically the hook would be registered before any state access or
+	// setting, but doing it here aids testing by showing the before-and-after
+	// effects.
+	RegisterExtras(highByteFlipper{})
+	t.Cleanup(TestOnlyClearRegisteredExtras)
+
+	noTransform := stateconf.SkipStateKeyTransformation()
+	assertEq(t, regularKey, flippedVal)
+	assertEq(t, regularKey, regularVal, noTransform)
+	assertEq(t, flippedKey, regularVal)
+	assertEq(t, flippedKey, flippedVal, noTransform)
+	assertCommittedEq(t, regularKey, flippedVal)
+	assertCommittedEq(t, regularKey, regularVal, noTransform)
+	assertCommittedEq(t, flippedKey, regularVal)
+	assertCommittedEq(t, flippedKey, flippedVal, noTransform)
+
+	updatedVal := common.Hash{'u', 'p', 'd', 'a', 't', 'e', 'd'}
+	sdb.SetState(addr, regularKey, updatedVal)
+	assertEq(t, regularKey, updatedVal)
+	assertEq(t, flippedKey, updatedVal, noTransform)
+	assertCommittedEq(t, regularKey, flippedVal)
+	assertCommittedEq(t, flippedKey, flippedVal, noTransform)
+}

core/vm/interface.go

@@ -23,6 +23,9 @@ import (
 	"github.com/ava-labs/libevm/core/types"
 	"github.com/ava-labs/libevm/params"
 	"github.com/holiman/uint256"
+
+	// libevm extra imports
+	"github.com/ava-labs/libevm/libevm/stateconf"
 )
 
 // StateDB is an EVM database for full state querying.
@@ -45,9 +48,9 @@ type StateDB interface {
 	SubRefund(uint64)
 	GetRefund() uint64
 
-	GetCommittedState(common.Address, common.Hash) common.Hash
-	GetState(common.Address, common.Hash) common.Hash
-	SetState(common.Address, common.Hash, common.Hash)
+	GetCommittedState(common.Address, common.Hash, ...stateconf.StateDBStateOption) common.Hash
+	GetState(common.Address, common.Hash, ...stateconf.StateDBStateOption) common.Hash
+	SetState(common.Address, common.Hash, common.Hash, ...stateconf.StateDBStateOption)
 
 	GetTransientState(addr common.Address, key common.Hash) common.Hash
 	SetTransientState(addr common.Address, key, value common.Hash)

eth/tracers/logger/logger_test.go

@@ -27,6 +27,9 @@ import (
 	"github.com/ava-labs/libevm/core/vm"
 	"github.com/ava-labs/libevm/params"
 	"github.com/holiman/uint256"
+
+	// libevm extra imports
+	"github.com/ava-labs/libevm/libevm/stateconf"
 )
 
 type dummyContractRef struct {
@@ -49,9 +52,12 @@ type dummyStatedb struct {
 	state.StateDB
 }
 
-func (*dummyStatedb) GetRefund() uint64                                       { return 1337 }
-func (*dummyStatedb) GetState(_ common.Address, _ common.Hash) common.Hash    { return common.Hash{} }
-func (*dummyStatedb) SetState(_ common.Address, _ common.Hash, _ common.Hash) {}
+func (*dummyStatedb) GetRefund() uint64 { return 1337 }
+func (*dummyStatedb) GetState(_ common.Address, _ common.Hash, _ ...stateconf.StateDBStateOption) common.Hash {
+	return common.Hash{}
+}
+func (*dummyStatedb) SetState(_ common.Address, _ common.Hash, _ common.Hash, _ ...stateconf.StateDBStateOption) {
+}
 
 func TestStoreCapture(t *testing.T) {
 	var (

libevm/libevm.go

@@ -20,6 +20,7 @@ import (
 	"github.com/holiman/uint256"
 
 	"github.com/ava-labs/libevm/common"
+	"github.com/ava-labs/libevm/libevm/stateconf"
 )
 
 // PrecompiledContract is an exact copy of vm.PrecompiledContract, mirrored here
@@ -43,8 +44,8 @@ type StateReader interface {
 
 	GetRefund() uint64
 
-	GetCommittedState(common.Address, common.Hash) common.Hash
-	GetState(common.Address, common.Hash) common.Hash
+	GetCommittedState(common.Address, common.Hash, ...stateconf.StateDBStateOption) common.Hash
+	GetState(common.Address, common.Hash, ...stateconf.StateDBStateOption) common.Hash
 
 	GetTransientState(addr common.Address, key common.Hash) common.Hash
 

libevm/stateconf/conf.go

@@ -113,3 +113,27 @@ func ExtractTrieDBUpdatePayload(opts ...TrieDBUpdateOption) (common.Hash, common
 	}
 	return *conf.parentBlockHash, *conf.currentBlockHash, true
 }
+
+// A StateDBStateOption configures the behaviour of state.StateDB methods for
+// getting and setting state: GetState(), GetCommittedState(), and SetState().
+type StateDBStateOption = options.Option[stateDBStateConfig]
+
+type stateDBStateConfig struct {
+	skipKeyTransformation bool
+}
+
+// SkipStateKeyTransformation causes any registered state-key transformation
+// hook to be ignored. See state.RegisterExtras() for details.
+func SkipStateKeyTransformation() StateDBStateOption {
+	return options.Func[stateDBStateConfig](func(c *stateDBStateConfig) {
+		c.skipKeyTransformation = true
+	})
+}
+
+// ShouldTransformStateKey parses the options, returning whether or not any
+// registered state-key transformation hook should be used; i.e. it returns
+// `true` i.f.f. there are no [SkipStateKeyTransformation] options in the
+// arguments.
+func ShouldTransformStateKey(opts ...StateDBStateOption) bool {
+	return !options.As(opts...).skipKeyTransformation
+}