feat!: types.ExtraPayloads supports SlimAccount

Author: ARR4N

core/state/state.libevm.go

@@ -27,7 +27,7 @@ import (
 func GetExtra[SA any](s *StateDB, p types.ExtraPayloads[SA], addr common.Address) SA {
 	stateObject := s.getStateObject(addr)
 	if stateObject != nil {
-		return p.FromStateAccount(&stateObject.data)
+		return p.FromPayloadCarrier(&stateObject.data)
 	}
 	var zero SA
 	return zero
@@ -45,9 +45,9 @@ func setExtraOnObject[SA any](s *stateObject, p types.ExtraPayloads[SA], addr co
 	s.db.journal.append(extraChange[SA]{
 		payloads: p,
 		account:  &addr,
-		prev:     p.FromStateAccount(&s.data),
+		prev:     p.FromPayloadCarrier(&s.data),
 	})
-	p.SetOnStateAccount(&s.data, extra)
+	p.SetOnPayloadCarrier(&s.data, extra)
 }
 
 // extraChange is a [journalEntry] for [SetExtra] / [setExtraOnObject].
@@ -60,5 +60,5 @@ type extraChange[SA any] struct {
 func (e extraChange[SA]) dirtied() *common.Address { return e.account }
 
 func (e extraChange[SA]) revert(s *StateDB) {
-	e.payloads.SetOnStateAccount(&s.getStateObject(*e.account).data, e.prev)
+	e.payloads.SetOnPayloadCarrier(&s.getStateObject(*e.account).data, e.prev)
 }

core/state/state.libevm_test.go

@@ -87,7 +87,7 @@ func TestGetSetExtra(t *testing.T) {
 				Root:     types.EmptyRootHash,
 				CodeHash: types.EmptyCodeHash[:],
 			}
-			payloads.SetOnStateAccount(want, extra)
+			payloads.SetOnPayloadCarrier(want, extra)
 
 			if diff := cmp.Diff(want, got); diff != "" {
 				t.Errorf("types.FullAccount(%T.Account()) diff (-want +got):\n%s", iter, diff)

core/state/state_object.libevm_test.go

@@ -46,7 +46,7 @@ func TestStateObjectEmpty(t *testing.T) {
 		{
 			name: "explicit false bool",
 			registerAndSet: func(acc *types.StateAccount) {
-				types.RegisterExtras[bool]().SetOnStateAccount(acc, false)
+				types.RegisterExtras[bool]().SetOnPayloadCarrier(acc, false)
 			},
 			wantEmpty: true,
 		},
@@ -60,7 +60,7 @@ func TestStateObjectEmpty(t *testing.T) {
 		{
 			name: "true bool",
 			registerAndSet: func(acc *types.StateAccount) {
-				types.RegisterExtras[bool]().SetOnStateAccount(acc, true)
+				types.RegisterExtras[bool]().SetOnPayloadCarrier(acc, true)
 			},
 			wantEmpty: false,
 		},

core/types/rlp_payload.libevm.go

@@ -26,16 +26,16 @@ import (
 )
 
 // RegisterExtras registers the type `SA` to be carried as an extra payload in
-// [StateAccount] structs. It is expected to be called in an `init()` function
-// and MUST NOT be called more than once.
+// [StateAccount] and [SlimAccount] structs. It is expected to be called in an
+// `init()` function and MUST NOT be called more than once.
 //
 // The payload will be treated as an extra struct field for the purposes of RLP
 // encoding and decoding. RLP handling is plumbed through to the `SA` via the
 // [StateAccountExtra] that holds it such that it acts as if there were a field
-// of type `SA` in all StateAccount structs.
+// of type `SA` in all StateAccount and SlimAccount structs.
 //
-// The payload can be acced via the [ExtraPayloads.FromStateAccount] method of
-// the accessor returned by RegisterExtras.
+// The payload can be accessed via the [ExtraPayloads.FromPayloadCarrier] method
+// of the accessor returned by RegisterExtras.
 func RegisterExtras[SA any]() ExtraPayloads[SA] {
 	if registeredExtras != nil {
 		panic("re-registration of Extras")
@@ -82,8 +82,8 @@ func (e *StateAccountExtra) clone() *StateAccountExtra {
 }
 
 // ExtraPayloads provides strongly typed access to the extra payload carried by
-// [StateAccount] structs. The only valid way to construct an instance is by a
-// call to [RegisterExtras].
+// [StateAccount] and [SlimAccount] structs. The only valid way to construct an
+// instance is by a call to [RegisterExtras].
 type ExtraPayloads[SA any] struct {
 	_ struct{} // make godoc show unexported fields so nobody tries to make their own instance ;)
 }
@@ -95,24 +95,36 @@ func (ExtraPayloads[SA]) cloneStateAccount(s *StateAccountExtra) *StateAccountEx
 	}
 }
 
-// FromStateAccount returns the StateAccount's payload.
-func (ExtraPayloads[SA]) FromStateAccount(a *StateAccount) SA {
+// ExtraPayloadCarrier is implemented by both [StateAccount] and [SlimAccount],
+// allowing for their [StateAccountExtra] payloads to be accessed in a type-safe
+// manner by [ExtraPayloads] instances.
+type ExtraPayloadCarrier interface {
+	extra() *StateAccountExtra
+}
+
+var _ = []ExtraPayloadCarrier{
+	(*StateAccount)(nil),
+	(*SlimAccount)(nil),
+}
+
+// FromPayloadCarrier returns the carriers's payload.
+func (ExtraPayloads[SA]) FromPayloadCarrier(a ExtraPayloadCarrier) SA {
 	return pseudo.MustNewValue[SA](a.extra().payload()).Get()
 }
 
-// PointerFromStateAccount returns a pointer to the StateAccounts's extra
-// payload. This is guaranteed to be non-nil.
+// PointerFromPayloadCarrier returns a pointer to the carriers's extra payload.
+// This is guaranteed to be non-nil.
 //
-// Note that copying a StateAccount by dereferencing a pointer will result in a
-// shallow copy and that the *SA returned here will therefore be shared by all
-// copies. If this is not the desired behaviour, use
-// [StateAccount.Copy] or [ExtraPayloads.SetOnStateAccount].
-func (ExtraPayloads[SA]) PointerFromStateAccount(a *StateAccount) *SA {
+// Note that copying a [StateAccount] or [SlimAccount] by dereferencing a
+// pointer will result in a shallow copy and that the *SA returned here will
+// therefore be shared by all copies. If this is not the desired behaviour, use
+// [StateAccount.Copy] or [ExtraPayloads.SetOnPayloadCarrier].
+func (ExtraPayloads[SA]) PointerFromPayloadCarrier(a ExtraPayloadCarrier) *SA {
 	return pseudo.MustPointerTo[SA](a.extra().payload()).Value.Get()
 }
 
-// SetOnStateAccount sets the StateAccount's payload.
-func (ExtraPayloads[SA]) SetOnStateAccount(a *StateAccount, val SA) {
+// SetOnPayloadCarrier sets the carriers's payload.
+func (ExtraPayloads[SA]) SetOnPayloadCarrier(a ExtraPayloadCarrier, val SA) {
 	a.extra().t = pseudo.From(val).Type
 }
 
@@ -124,12 +136,20 @@ type StateAccountExtra struct {
 }
 
 func (a *StateAccount) extra() *StateAccountExtra {
-	if a.Extra == nil {
-		a.Extra = &StateAccountExtra{
+	return getOrSetNewStateAccountExtra(&a.Extra)
+}
+
+func (a *SlimAccount) extra() *StateAccountExtra {
+	return getOrSetNewStateAccountExtra(&a.Extra)
+}
+
+func getOrSetNewStateAccountExtra(curr **StateAccountExtra) *StateAccountExtra {
+	if *curr == nil {
+		*curr = &StateAccountExtra{
 			t: registeredExtras.newStateAccount(),
 		}
 	}
-	return a.Extra
+	return *curr
 }
 
 func (e *StateAccountExtra) payload() *pseudo.Type {
@@ -143,7 +163,7 @@ func (e *StateAccountExtra) payload() *pseudo.Type {
 // of tests.
 //
 // Equal MUST NOT be used in production. Instead, compare values returned by
-// [ExtraPayloads.FromStateAccount].
+// [ExtraPayloads.FromPayloadCarrier].
 func (e *StateAccountExtra) Equal(f *StateAccountExtra) bool {
 	if false {
 		// TODO(arr4n): calling this results in an error from cmp.Diff():

core/types/state_account.go

@@ -35,7 +35,7 @@ type StateAccount struct {
 	Root     common.Hash // merkle root of the storage trie
 	CodeHash []byte
 
-	Extra *StateAccountExtra
+	Extra *StateAccountExtra // only access via [ExtraPayloads] instances
 }
 
 // NewEmptyStateAccount constructs an empty state account.
@@ -71,7 +71,7 @@ type SlimAccount struct {
 	Root     []byte // Nil if root equals to types.EmptyRootHash
 	CodeHash []byte // Nil if hash equals to types.EmptyCodeHash
 
-	Extra *StateAccountExtra
+	Extra *StateAccountExtra // only access via [ExtraPayloads] instances
 }
 
 // SlimAccountRLP encodes the state account in 'slim RLP' format.

core/types/state_account_storage.libevm_test.go

@@ -71,25 +71,25 @@ func TestStateAccountExtraViaTrieStorage(t *testing.T) {
 		{
 			name: "true-boolean payload",
 			registerAndSetExtra: func(a *types.StateAccount) *types.StateAccount {
-				types.RegisterExtras[bool]().SetOnStateAccount(a, true)
+				types.RegisterExtras[bool]().SetOnPayloadCarrier(a, true)
 				return a
 			},
 			assertExtra: func(t *testing.T, sa *types.StateAccount) {
 				t.Helper()
-				assert.Truef(t, types.ExtraPayloads[bool]{}.FromStateAccount(sa), "")
+				assert.Truef(t, types.ExtraPayloads[bool]{}.FromPayloadCarrier(sa), "")
 			},
 			wantTrieHash: trueBool,
 		},
 		{
 			name: "explicit false-boolean payload",
 			registerAndSetExtra: func(a *types.StateAccount) *types.StateAccount {
 				p := types.RegisterExtras[bool]()
-				p.SetOnStateAccount(a, false) // the explicit part
+				p.SetOnPayloadCarrier(a, false) // the explicit part
 				return a
 			},
 			assertExtra: func(t *testing.T, sa *types.StateAccount) {
 				t.Helper()
-				assert.Falsef(t, types.ExtraPayloads[bool]{}.FromStateAccount(sa), "")
+				assert.Falsef(t, types.ExtraPayloads[bool]{}.FromPayloadCarrier(sa), "")
 			},
 			wantTrieHash: falseBool,
 		},
@@ -102,20 +102,20 @@ func TestStateAccountExtraViaTrieStorage(t *testing.T) {
 			},
 			assertExtra: func(t *testing.T, sa *types.StateAccount) {
 				t.Helper()
-				assert.Falsef(t, types.ExtraPayloads[bool]{}.FromStateAccount(sa), "")
+				assert.Falsef(t, types.ExtraPayloads[bool]{}.FromPayloadCarrier(sa), "")
 			},
 			wantTrieHash: falseBool,
 		},
 		{
 			name: "arbitrary payload",
 			registerAndSetExtra: func(a *types.StateAccount) *types.StateAccount {
 				p := arbitraryPayload{arbitraryData}
-				types.RegisterExtras[arbitraryPayload]().SetOnStateAccount(a, p)
+				types.RegisterExtras[arbitraryPayload]().SetOnPayloadCarrier(a, p)
 				return a
 			},
 			assertExtra: func(t *testing.T, sa *types.StateAccount) {
 				t.Helper()
-				got := types.ExtraPayloads[arbitraryPayload]{}.FromStateAccount(sa)
+				got := types.ExtraPayloads[arbitraryPayload]{}.FromPayloadCarrier(sa)
 				assert.Equalf(t, arbitraryPayload{arbitraryData}, got, "")
 			},
 			wantTrieHash: arbitrary,