feat(core/types): body RLP hooks registration (#130)

Allow to register body extras in consumers of libevm.

Co-authored-by: Arran Schlosberg <519948+ARR4N@users.noreply.github.com>

Author: qdm12

core/state/state.libevm_test.go

@@ -45,7 +45,11 @@ func TestGetSetExtra(t *testing.T) {
 	t.Cleanup(types.TestOnlyClearRegisteredExtras)
 	// Just as its Data field is a pointer, the registered type is a pointer to
 	// test deep copying.
-	payloads := types.RegisterExtras[types.NOOPHeaderHooks, *types.NOOPHeaderHooks, *accountExtra]().StateAccount
+	payloads := types.RegisterExtras[
+		types.NOOPHeaderHooks, *types.NOOPHeaderHooks,
+		types.NOOPBodyHooks, *types.NOOPBodyHooks,
+		*accountExtra,
+	]().StateAccount
 
 	rng := ethtest.NewPseudoRand(42)
 	addr := rng.Address()

core/state/state_object.libevm_test.go

@@ -46,21 +46,33 @@ func TestStateObjectEmpty(t *testing.T) {
 		{
 			name: "explicit false bool",
 			registerAndSet: func(acc *types.StateAccount) {
-				types.RegisterExtras[types.NOOPHeaderHooks, *types.NOOPHeaderHooks, bool]().StateAccount.Set(acc, false)
+				types.RegisterExtras[
+					types.NOOPHeaderHooks, *types.NOOPHeaderHooks,
+					types.NOOPBodyHooks, *types.NOOPBodyHooks,
+					bool,
+				]().StateAccount.Set(acc, false)
 			},
 			wantEmpty: true,
 		},
 		{
 			name: "implicit false bool",
 			registerAndSet: func(*types.StateAccount) {
-				types.RegisterExtras[types.NOOPHeaderHooks, *types.NOOPHeaderHooks, bool]()
+				types.RegisterExtras[
+					types.NOOPHeaderHooks, *types.NOOPHeaderHooks,
+					types.NOOPBodyHooks, *types.NOOPBodyHooks,
+					bool,
+				]()
 			},
 			wantEmpty: true,
 		},
 		{
 			name: "true bool",
 			registerAndSet: func(acc *types.StateAccount) {
-				types.RegisterExtras[types.NOOPHeaderHooks, *types.NOOPHeaderHooks, bool]().StateAccount.Set(acc, true)
+				types.RegisterExtras[
+					types.NOOPHeaderHooks, *types.NOOPHeaderHooks,
+					types.NOOPBodyHooks, *types.NOOPBodyHooks,
+					bool,
+				]().StateAccount.Set(acc, true)
 			},
 			wantEmpty: false,
 		},

core/types/backwards_compat.libevm_test.go

@@ -21,6 +21,7 @@ import (
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
+	"github.com/google/go-cmp/cmp/cmpopts"
 	"github.com/kr/pretty"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -56,14 +57,18 @@ func TestBodyRLPBackwardsCompatibility(t *testing.T) {
 	for _, tx := range txMatrix {
 		for _, u := range uncleMatrix {
 			for _, w := range withdrawMatrix {
-				bodies = append(bodies, &Body{tx, u, w})
+				bodies = append(bodies, &Body{tx, u, w, nil /* extra field */})
 			}
 		}
 	}
 
 	for _, body := range bodies {
 		t.Run("", func(t *testing.T) {
-			t.Logf("\n%s", pretty.Sprint(body))
+			t.Cleanup(func() {
+				if t.Failed() {
+					t.Logf("\n%s", pretty.Sprint(body))
+				}
+			})
 
 			// The original [Body] doesn't implement [rlp.Encoder] nor
 			// [rlp.Decoder] so we can use a methodless equivalent as the gold
@@ -74,14 +79,15 @@ func TestBodyRLPBackwardsCompatibility(t *testing.T) {
 
 			t.Run("Encode", func(t *testing.T) {
 				got, err := rlp.EncodeToBytes(body)
-				require.NoErrorf(t, err, "rlp.EncodeToBytes(%#v)", body)
-				assert.Equalf(t, wantRLP, got, "rlp.EncodeToBytes(%#v)", body)
+				require.NoErrorf(t, err, "rlp.EncodeToBytes(%T)", body)
+				assert.Equalf(t, wantRLP, got, "rlp.EncodeToBytes(%T)", body)
 			})
 
 			t.Run("Decode", func(t *testing.T) {
 				got := new(Body)
 				err := rlp.DecodeBytes(wantRLP, got)
-				require.NoErrorf(t, err, "rlp.DecodeBytes(%v, %T)", wantRLP, got)
+				require.NoErrorf(t, err, "rlp.DecodeBytes(rlp.EncodeToBytes(%T), %T) resulted in %s",
+					(*withoutMethods)(body), got, pretty.Sprint(got))
 
 				want := body
 				// Regular RLP decoding will never leave these non-optional
@@ -96,9 +102,10 @@ func TestBodyRLPBackwardsCompatibility(t *testing.T) {
 				opts := cmp.Options{
 					cmp.Comparer((*Header).equalHash),
 					cmp.Comparer((*Transaction).equalHash),
+					cmpopts.IgnoreUnexported(Body{}),
 				}
-				if diff := cmp.Diff(body, got, opts); diff != "" {
-					t.Errorf("rlp.DecodeBytes(rlp.EncodeToBytes(%#v)) diff (-want +got):\n%s", body, diff)
+				if diff := cmp.Diff(want, got, opts); diff != "" {
+					t.Errorf("rlp.DecodeBytes(rlp.EncodeToBytes(%T)) diff (-want +got):\n%s", (*withoutMethods)(body), diff)
 				}
 			})
 		})
@@ -148,10 +155,13 @@ func TestBodyRLPCChainCompat(t *testing.T) {
 	// The inputs to this test were used to generate the expected RLP with
 	// ava-labs/coreth. This serves as both an example of how to use [BodyHooks]
 	// and a test of compatibility.
-
-	t.Cleanup(func() {
-		TestOnlyRegisterBodyHooks(NOOPBodyHooks{})
-	})
+	TestOnlyClearRegisteredExtras()
+	t.Cleanup(TestOnlyClearRegisteredExtras)
+	extras := RegisterExtras[
+		NOOPHeaderHooks, *NOOPHeaderHooks,
+		cChainBodyExtras, *cChainBodyExtras,
+		struct{},
+	]()
 
 	body := &Body{
 		Transactions: []*Transaction{
@@ -194,24 +204,24 @@ func TestBodyRLPCChainCompat(t *testing.T) {
 			require.NoErrorf(t, err, "hex.DecodeString(%q)", tt.wantRLPHex)
 
 			t.Run("Encode", func(t *testing.T) {
-				TestOnlyRegisterBodyHooks(tt.extra)
+				extras.Body.Set(body, tt.extra)
 				got, err := rlp.EncodeToBytes(body)
 				require.NoErrorf(t, err, "rlp.EncodeToBytes(%+v)", body)
 				assert.Equalf(t, wantRLP, got, "rlp.EncodeToBytes(%+v)", body)
 			})
 
 			t.Run("Decode", func(t *testing.T) {
 				var extra cChainBodyExtras
-				TestOnlyRegisterBodyHooks(&extra)
-
 				got := new(Body)
+				extras.Body.Set(got, &extra)
 				err := rlp.DecodeBytes(wantRLP, got)
 				require.NoErrorf(t, err, "rlp.DecodeBytes(%#x, %T)", wantRLP, got)
 				assert.Equal(t, tt.extra, &extra, "rlp.DecodeBytes(%#x, [%T as registered extra in %T carrier])", wantRLP, &extra, got)
 
 				opts := cmp.Options{
 					cmp.Comparer((*Header).equalHash),
 					cmp.Comparer((*Transaction).equalHash),
+					cmpopts.IgnoreUnexported(Body{}),
 				}
 				if diff := cmp.Diff(body, got, opts); diff != "" {
 					t.Errorf("rlp.DecodeBytes(%#x, [%T while carrying registered %T extra payload]) diff (-want +got):\n%s", wantRLP, got, &extra, diff)

core/types/backwards_compat_diffpkg.libevm_test.go

@@ -40,7 +40,11 @@ func TestHeaderRLPBackwardsCompatibility(t *testing.T) {
 		{
 			name: "no-op header hooks",
 			register: func() {
-				RegisterExtras[NOOPHeaderHooks, *NOOPHeaderHooks, struct{}]()
+				RegisterExtras[
+					NOOPHeaderHooks, *NOOPHeaderHooks,
+					NOOPBodyHooks, *NOOPBodyHooks,
+					struct{},
+				]()
 			},
 		},
 	}

core/types/block.go

@@ -176,6 +176,8 @@ type Body struct {
 	Transactions []*Transaction
 	Uncles       []*Header
 	Withdrawals  []*Withdrawal `rlp:"optional"`
+
+	extra *pseudo.Type // See [RegisterExtras]
 }
 
 // Block represents an Ethereum block.
@@ -338,7 +340,7 @@ func (b *Block) EncodeRLP(w io.Writer) error {
 // Body returns the non-header content of the block.
 // Note the returned data is not an independent copy.
 func (b *Block) Body() *Body {
-	return &Body{b.transactions, b.uncles, b.withdrawals}
+	return &Body{b.transactions, b.uncles, b.withdrawals, nil /* unexported extras field */}
 }
 
 // Accessors for body data. These do not return a copy because the content

core/types/block.libevm.go

@@ -22,7 +22,6 @@ import (
 	"io"
 
 	"github.com/ava-labs/libevm/libevm/pseudo"
-	"github.com/ava-labs/libevm/libevm/testonly"
 	"github.com/ava-labs/libevm/rlp"
 )
 
@@ -45,7 +44,7 @@ func (h *Header) hooks() HeaderHooks {
 	return new(NOOPHeaderHooks)
 }
 
-func (e ExtraPayloads[HPtr, SA]) hooksFromHeader(h *Header) HeaderHooks {
+func (e ExtraPayloads[HPtr, BPtr, SA]) hooksFromHeader(h *Header) HeaderHooks {
 	return e.Header.Get(h)
 }
 
@@ -134,22 +133,11 @@ type BodyHooks interface {
 	RLPFieldPointersForDecoding(*Body) *rlp.Fields
 }
 
-// TestOnlyRegisterBodyHooks is a temporary means of "registering" BodyHooks for
-// the purpose of testing. It will panic if called outside of a test.
-func TestOnlyRegisterBodyHooks(h BodyHooks) {
-	testonly.OrPanic(func() {
-		todoRegisteredBodyHooks = h
-	})
-}
-
-// todoRegisteredBodyHooks is a temporary placeholder for "registering"
-// BodyHooks, before they are included in [RegisterExtras].
-var todoRegisteredBodyHooks BodyHooks = NOOPBodyHooks{}
-
 func (b *Body) hooks() BodyHooks {
-	// TODO(arr4n): when incorporating BodyHooks into [RegisterExtras], the
-	// [todoRegisteredBodyHooks] variable MUST be removed.
-	return todoRegisteredBodyHooks
+	if r := registeredExtras; r.Registered() {
+		return r.Get().hooks.hooksFromBody(b)
+	}
+	return NOOPBodyHooks{}
 }
 
 // NOOPBodyHooks implements [BodyHooks] such that they are equivalent to no type
@@ -160,7 +148,7 @@ type NOOPBodyHooks struct{}
 // fields and their order, which we lock in here as a change detector. If this
 // breaks then it MUST be updated and the RLP methods reviewed + new
 // backwards-compatibility tests added.
-var _ = &Body{[]*Transaction{}, []*Header{}, []*Withdrawal{}}
+var _ = &Body{[]*Transaction{}, []*Header{}, []*Withdrawal{}, nil /* extra unexported type */}
 
 func (NOOPBodyHooks) RLPFieldsForEncoding(b *Body) *rlp.Fields {
 	return &rlp.Fields{
@@ -175,3 +163,19 @@ func (NOOPBodyHooks) RLPFieldPointersForDecoding(b *Body) *rlp.Fields {
 		Optional: []any{&b.Withdrawals},
 	}
 }
+
+func (e ExtraPayloads[HPtr, BPtr, SA]) hooksFromBody(b *Body) BodyHooks {
+	return e.Body.Get(b)
+}
+
+func (b *Body) extraPayload() *pseudo.Type {
+	r := registeredExtras
+	if !r.Registered() {
+		// See params.ChainConfig.extraPayload() for panic rationale.
+		panic(fmt.Sprintf("%T.extraPayload() called before RegisterExtras()", r))
+	}
+	if b.extra == nil {
+		b.extra = r.Get().newBody()
+	}
+	return b.extra
+}

core/types/block.libevm_test.go

@@ -86,7 +86,11 @@ func TestHeaderHooks(t *testing.T) {
 	TestOnlyClearRegisteredExtras()
 	defer TestOnlyClearRegisteredExtras()
 
-	extras := RegisterExtras[stubHeaderHooks, *stubHeaderHooks, struct{}]()
+	extras := RegisterExtras[
+		stubHeaderHooks, *stubHeaderHooks,
+		NOOPBodyHooks, *NOOPBodyHooks,
+		struct{},
+	]()
 	rng := ethtest.NewPseudoRand(13579)
 
 	suffix := rng.Bytes(8)

core/types/rlp_payload.libevm.go

@@ -46,13 +46,21 @@ func RegisterExtras[
 		HeaderHooks
 		*H
 	},
+	B any, BPtr interface {
+		BodyHooks
+		*B
+	},
 	SA any,
-]() ExtraPayloads[HPtr, SA] {
-	extra := ExtraPayloads[HPtr, SA]{
+]() ExtraPayloads[HPtr, BPtr, SA] {
+	extra := ExtraPayloads[HPtr, BPtr, SA]{
 		Header: pseudo.NewAccessor[*Header, HPtr](
 			(*Header).extraPayload,
 			func(h *Header, t *pseudo.Type) { h.extra = t },
 		),
+		Body: pseudo.NewAccessor[*Body, BPtr](
+			(*Body).extraPayload,
+			func(b *Body, t *pseudo.Type) { b.extra = t },
+		),
 		StateAccount: pseudo.NewAccessor[StateOrSlimAccount, SA](
 			func(a StateOrSlimAccount) *pseudo.Type { return a.extra().payload() },
 			func(a StateOrSlimAccount, t *pseudo.Type) { a.extra().t = t },
@@ -63,10 +71,11 @@ func RegisterExtras[
 			var x SA
 			return fmt.Sprintf("%T", x)
 		}(),
-		// The [ExtraPayloads] that we returns is based on [HPtr,SA], not [H,SA]
-		// so our constructors MUST match that. This guarantees that calls to
-		// the [HeaderHooks] methods will never be performed on a nil pointer.
+		// The [ExtraPayloads] that we returns is based on [HPtr,BPtr,SA], not
+		// [H,B,SA] so our constructors MUST match that. This guarantees that calls to
+		// the [HeaderHooks] and [BodyHooks] methods will never be performed on a nil pointer.
 		newHeader:         pseudo.NewConstructor[H]().NewPointer, // i.e. non-nil HPtr
+		newBody:           pseudo.NewConstructor[B]().NewPointer, // i.e. non-nil BPtr
 		newStateAccount:   pseudo.NewConstructor[SA]().Zero,
 		cloneStateAccount: extra.cloneStateAccount,
 		hooks:             extra,
@@ -87,11 +96,14 @@ func TestOnlyClearRegisteredExtras() {
 var registeredExtras register.AtMostOnce[*extraConstructors]
 
 type extraConstructors struct {
-	stateAccountType           string
-	newHeader, newStateAccount func() *pseudo.Type
-	cloneStateAccount          func(*StateAccountExtra) *StateAccountExtra
-	hooks                      interface {
+	stateAccountType  string
+	newHeader         func() *pseudo.Type
+	newBody           func() *pseudo.Type
+	newStateAccount   func() *pseudo.Type
+	cloneStateAccount func(*StateAccountExtra) *StateAccountExtra
+	hooks             interface {
 		hooksFromHeader(*Header) HeaderHooks
+		hooksFromBody(*Body) BodyHooks
 	}
 }
 
@@ -105,14 +117,15 @@ func (e *StateAccountExtra) clone() *StateAccountExtra {
 }
 
 // ExtraPayloads provides strongly typed access to the extra payload carried by
-// [Header], [StateAccount], and [SlimAccount] structs. The only valid way to
+// [Header], [Body], [StateAccount], and [SlimAccount] structs. The only valid way to
 // construct an instance is by a call to [RegisterExtras].
-type ExtraPayloads[HPtr HeaderHooks, SA any] struct {
+type ExtraPayloads[HPtr HeaderHooks, BPtr BodyHooks, SA any] struct {
 	Header       pseudo.Accessor[*Header, HPtr]
+	Body         pseudo.Accessor[*Body, BPtr]
 	StateAccount pseudo.Accessor[StateOrSlimAccount, SA] // Also provides [SlimAccount] access.
 }
 
-func (ExtraPayloads[HPtr, SA]) cloneStateAccount(s *StateAccountExtra) *StateAccountExtra {
+func (ExtraPayloads[HPtr, BPtr, SA]) cloneStateAccount(s *StateAccountExtra) *StateAccountExtra {
 	v := pseudo.MustNewValue[SA](s.t)
 	return &StateAccountExtra{
 		t: pseudo.From(v.Get()).Type,

core/types/state_account.libevm_test.go

@@ -46,7 +46,11 @@ func TestStateAccountRLP(t *testing.T) {
 	explicitFalseBoolean := test{
 		name: "explicit false-boolean extra",
 		register: func() {
-			RegisterExtras[NOOPHeaderHooks, *NOOPHeaderHooks, bool]()
+			RegisterExtras[
+				NOOPHeaderHooks, *NOOPHeaderHooks,
+				NOOPBodyHooks, *NOOPBodyHooks,
+				bool,
+			]()
 		},
 		acc: &StateAccount{
 			Nonce:    0x444444,
@@ -76,7 +80,11 @@ func TestStateAccountRLP(t *testing.T) {
 		{
 			name: "true-boolean extra",
 			register: func() {
-				RegisterExtras[NOOPHeaderHooks, *NOOPHeaderHooks, bool]()
+				RegisterExtras[
+					NOOPHeaderHooks, *NOOPHeaderHooks,
+					NOOPBodyHooks, *NOOPBodyHooks,
+					bool,
+				]()
 			},
 			acc: &StateAccount{
 				Nonce:    0x444444,