chore(core/types): restructure libevm hooks code (#135)

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

Author: qdm12

core/types/block.libevm.go

@@ -18,10 +18,8 @@ package types
 
 import (
 	"encoding/json"
-	"fmt"
 	"io"
 
-	"github.com/ava-labs/libevm/libevm/pseudo"
 	"github.com/ava-labs/libevm/rlp"
 )
 
@@ -35,19 +33,6 @@ type HeaderHooks interface {
 	PostCopy(dst *Header)
 }
 
-// hooks returns the Header's registered HeaderHooks, if any, otherwise a
-// [NOOPHeaderHooks] suitable for running default behaviour.
-func (h *Header) hooks() HeaderHooks {
-	if r := registeredExtras; r.Registered() {
-		return r.Get().hooks.hooksFromHeader(h)
-	}
-	return new(NOOPHeaderHooks)
-}
-
-func (e ExtraPayloads[HPtr, BPtr, SA]) hooksFromHeader(h *Header) HeaderHooks {
-	return e.Header.Get(h)
-}
-
 var _ interface {
 	rlp.Encoder
 	rlp.Decoder
@@ -75,18 +60,6 @@ func (h *Header) DecodeRLP(s *rlp.Stream) error {
 	return h.hooks().DecodeRLP(h, s)
 }
 
-func (h *Header) 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 h.extra == nil {
-		h.extra = r.Get().newHeader()
-	}
-	return h.extra
-}
-
 // NOOPHeaderHooks implements [HeaderHooks] such that they are equivalent to
 // no type having been registered.
 type NOOPHeaderHooks struct{}
@@ -133,13 +106,6 @@ type BodyHooks interface {
 	RLPFieldPointersForDecoding(*Body) *rlp.Fields
 }
 
-func (b *Body) hooks() BodyHooks {
-	if r := registeredExtras; r.Registered() {
-		return r.Get().hooks.hooksFromBody(b)
-	}
-	return NOOPBodyHooks{}
-}
-
 // NOOPBodyHooks implements [BodyHooks] such that they are equivalent to no type
 // having been registered.
 type NOOPBodyHooks struct{}
@@ -163,19 +129,3 @@ 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/rlp_payload.libevm.go

@@ -74,11 +74,10 @@ func RegisterExtras[
 		// 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,
+		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,
+		hooks:           extra,
 	})
 	return extra
 }
@@ -96,23 +95,65 @@ func TestOnlyClearRegisteredExtras() {
 var registeredExtras register.AtMostOnce[*extraConstructors]
 
 type extraConstructors struct {
-	stateAccountType  string
-	newHeader         func() *pseudo.Type
-	newBody           func() *pseudo.Type
-	newStateAccount   func() *pseudo.Type
-	cloneStateAccount func(*StateAccountExtra) *StateAccountExtra
-	hooks             interface {
+	stateAccountType string
+	newHeader        func() *pseudo.Type
+	newBody          func() *pseudo.Type
+	newStateAccount  func() *pseudo.Type
+	hooks            interface {
 		hooksFromHeader(*Header) HeaderHooks
 		hooksFromBody(*Body) BodyHooks
+		cloneStateAccount(*StateAccountExtra) *StateAccountExtra
 	}
 }
 
+func extraPayloadOrSetDefault(field **pseudo.Type, construct func(*extraConstructors) *pseudo.Type) *pseudo.Type {
+	r := registeredExtras
+	if !r.Registered() {
+		// See params.ChainConfig.extraPayload() for panic rationale.
+		panic("<T>.extraPayload() called before RegisterExtras()")
+	}
+	if *field == nil {
+		*field = construct(r.Get())
+	}
+	return *field
+}
+
+func (h *Header) extraPayload() *pseudo.Type {
+	return extraPayloadOrSetDefault(&h.extra, func(c *extraConstructors) *pseudo.Type {
+		return c.newHeader()
+	})
+}
+
+func (b *Body) extraPayload() *pseudo.Type {
+	return extraPayloadOrSetDefault(&b.extra, func(c *extraConstructors) *pseudo.Type {
+		return c.newBody()
+	})
+}
+
+// hooks returns the [Header]'s registered [HeaderHooks], if any, otherwise a
+// [NOOPHeaderHooks] suitable for running default behaviour.
+func (h *Header) hooks() HeaderHooks {
+	if r := registeredExtras; r.Registered() {
+		return r.Get().hooks.hooksFromHeader(h)
+	}
+	return new(NOOPHeaderHooks)
+}
+
+// hooks returns the [Body]'s registered [BodyHooks], if any, otherwise a
+// [NOOPBodyHooks] suitable for running default behaviour.
+func (b *Body) hooks() BodyHooks {
+	if r := registeredExtras; r.Registered() {
+		return r.Get().hooks.hooksFromBody(b)
+	}
+	return NOOPBodyHooks{}
+}
+
 func (e *StateAccountExtra) clone() *StateAccountExtra {
 	switch r := registeredExtras; {
 	case !r.Registered(), e == nil:
 		return nil
 	default:
-		return r.Get().cloneStateAccount(e)
+		return r.Get().hooks.cloneStateAccount(e)
 	}
 }
 
@@ -125,6 +166,9 @@ type ExtraPayloads[HPtr HeaderHooks, BPtr BodyHooks, SA any] struct {
 	StateAccount pseudo.Accessor[StateOrSlimAccount, SA] // Also provides [SlimAccount] access.
 }
 
+func (e ExtraPayloads[HPtr, BPtr, SA]) hooksFromHeader(h *Header) HeaderHooks { return e.Header.Get(h) }
+func (e ExtraPayloads[HPtr, BPtr, SA]) hooksFromBody(b *Body) BodyHooks       { return e.Body.Get(b) }
+
 func (ExtraPayloads[HPtr, BPtr, SA]) cloneStateAccount(s *StateAccountExtra) *StateAccountExtra {
 	v := pseudo.MustNewValue[SA](s.t)
 	return &StateAccountExtra{