feat: params.ChainConfig extra payload can use root JSON (#8)

* feat: `params.ChainConfig` extra payload can use root JSON

* refactor: simplify `ChainConfig.UnmarshalJSON()` branches

* fix: change redundant `assert` to `require` for simplicity

Author: ARR4N

params/config.libevm.go

@@ -1,7 +1,6 @@
 package params
 
 import (
-	"encoding/json"
 	"fmt"
 	"math/big"
 	"reflect"
@@ -14,6 +13,15 @@ import (
 // Extras are arbitrary payloads to be added as extra fields in [ChainConfig]
 // and [Rules] structs. See [RegisterExtras].
 type Extras[C ChainConfigHooks, R RulesHooks] struct {
+	// ReuseJSONRoot, if true, signals that JSON unmarshalling of a
+	// [ChainConfig] MUST reuse the root JSON input when unmarshalling the extra
+	// payload. If false, it is assumed that the extra JSON payload is nested in
+	// the "extra" key.
+	//
+	// *NOTE* this requires multiple passes for both marshalling and
+	// unmarshalling of JSON so is inefficient and should be used as a last
+	// resort.
+	ReuseJSONRoot bool
 	// NewRules, if non-nil is called at the end of [ChainConfig.Rules] with the
 	// newly created [Rules] and other context from the method call. Its
 	// returned value will be the extra payload of the [Rules]. If NewRules is
@@ -51,10 +59,11 @@ func RegisterExtras[C ChainConfigHooks, R RulesHooks](e Extras[C, R]) ExtraPaylo
 
 	getter := e.getter()
 	registeredExtras = &extraConstructors{
-		chainConfig: pseudo.NewConstructor[C](),
-		rules:       pseudo.NewConstructor[R](),
-		newForRules: e.newForRules,
-		getter:      getter,
+		chainConfig:   pseudo.NewConstructor[C](),
+		rules:         pseudo.NewConstructor[R](),
+		reuseJSONRoot: e.ReuseJSONRoot,
+		newForRules:   e.newForRules,
+		getter:        getter,
 	}
 	return getter
 }
@@ -87,6 +96,7 @@ var registeredExtras *extraConstructors
 
 type extraConstructors struct {
 	chainConfig, rules pseudo.Constructor
+	reuseJSONRoot      bool
 	newForRules        func(_ *ChainConfig, _ *Rules, blockNum *big.Int, isMerge bool, timestamp uint64) *pseudo.Type
 	// use top-level hooksFrom<X>() functions instead of these as they handle
 	// instances where no [Extras] were registered.
@@ -158,42 +168,6 @@ func (e ExtraPayloadGetter[C, R]) hooksFromRules(r *Rules) RulesHooks {
 	return NOOPHooks{}
 }
 
-// UnmarshalJSON implements the [json.Unmarshaler] interface.
-func (c *ChainConfig) UnmarshalJSON(data []byte) error {
-	type raw ChainConfig // doesn't inherit methods so avoids recursing back here (infinitely)
-	cc := &struct {
-		*raw
-		Extra *pseudo.Type `json:"extra"`
-	}{
-		raw: (*raw)(c), // embedded to achieve regular JSON unmarshalling
-	}
-	if e := registeredExtras; e != nil {
-		cc.Extra = e.chainConfig.NilPointer() // `c.extra` is otherwise unexported
-	}
-
-	if err := json.Unmarshal(data, cc); err != nil {
-		return err
-	}
-	c.extra = cc.Extra
-	return nil
-}
-
-// MarshalJSON implements the [json.Marshaler] interface.
-func (c *ChainConfig) MarshalJSON() ([]byte, error) {
-	// See UnmarshalJSON() for rationale.
-	type raw ChainConfig
-	cc := &struct {
-		*raw
-		Extra *pseudo.Type `json:"extra"`
-	}{raw: (*raw)(c), Extra: c.extra}
-	return json.Marshal(cc)
-}
-
-var _ interface {
-	json.Marshaler
-	json.Unmarshaler
-} = (*ChainConfig)(nil)
-
 // addRulesExtra is called at the end of [ChainConfig.Rules]; it exists to
 // abstract the libevm-specific behaviour outside of original geth code.
 func (c *ChainConfig) addRulesExtra(r *Rules, blockNum *big.Int, isMerge bool, timestamp uint64) {

params/config.libevm_test.go

@@ -104,8 +104,6 @@ func TestRegisterExtras(t *testing.T) {
 			require.NoError(t, json.Unmarshal(buf, got))
 			assert.Equal(t, tt.ccExtra.Interface(), got.extraPayload().Interface())
 			assert.Equal(t, in, got)
-			// TODO: do we need an explicit test of the JSON output, or is a
-			// Marshal-Unmarshal round trip sufficient?
 
 			gotRules := got.Rules(nil, false, 0)
 			assert.Equal(t, tt.wantRulesExtra, gotRules.extraPayload().Interface())

params/json.libevm.go

@@ -0,0 +1,110 @@
+package params
+
+import (
+	"encoding/json"
+	"fmt"
+
+	"github.com/ethereum/go-ethereum/libevm/pseudo"
+)
+
+var _ interface {
+	json.Marshaler
+	json.Unmarshaler
+} = (*ChainConfig)(nil)
+
+// chainConfigWithoutMethods avoids infinite recurion into
+// [ChainConfig.UnmarshalJSON].
+type chainConfigWithoutMethods ChainConfig
+
+// chainConfigWithExportedExtra supports JSON (un)marshalling of a [ChainConfig]
+// while exposing the `extra` field as the "extra" JSON key.
+type chainConfigWithExportedExtra struct {
+	*chainConfigWithoutMethods              // embedded to achieve regular JSON unmarshalling
+	Extra                      *pseudo.Type `json:"extra"` // `c.extra` is otherwise unexported
+}
+
+// UnmarshalJSON implements the [json.Unmarshaler] interface.
+func (c *ChainConfig) UnmarshalJSON(data []byte) error {
+	switch reg := registeredExtras; {
+	case reg != nil && !reg.reuseJSONRoot:
+		return c.unmarshalJSONWithExtra(data)
+
+	case reg != nil && reg.reuseJSONRoot: // although the latter is redundant, it's clearer
+		c.extra = reg.chainConfig.NilPointer()
+		if err := json.Unmarshal(data, c.extra); err != nil {
+			c.extra = nil
+			return err
+		}
+		fallthrough // Important! We've only unmarshalled the extra field.
+	default: // reg == nil
+		return json.Unmarshal(data, (*chainConfigWithoutMethods)(c))
+	}
+}
+
+// unmarshalJSONWithExtra unmarshals JSON under the assumption that the
+// registered [Extras] payload is in the JSON "extra" key. All other
+// unmarshalling is performed as if no [Extras] were registered.
+func (c *ChainConfig) unmarshalJSONWithExtra(data []byte) error {
+	cc := &chainConfigWithExportedExtra{
+		chainConfigWithoutMethods: (*chainConfigWithoutMethods)(c),
+		Extra:                     registeredExtras.chainConfig.NilPointer(),
+	}
+	if err := json.Unmarshal(data, cc); err != nil {
+		return err
+	}
+	c.extra = cc.Extra
+	return nil
+}
+
+// MarshalJSON implements the [json.Marshaler] interface.
+func (c *ChainConfig) MarshalJSON() ([]byte, error) {
+	switch reg := registeredExtras; {
+	case reg == nil:
+		return json.Marshal((*chainConfigWithoutMethods)(c))
+
+	case !reg.reuseJSONRoot:
+		return c.marshalJSONWithExtra()
+
+	default: // reg.reuseJSONRoot == true
+		// The inverse of reusing the JSON root is merging two JSON buffers,
+		// which isn't supported by the native package. So we use
+		// map[string]json.RawMessage intermediates.
+		geth, err := toJSONRawMessages((*chainConfigWithoutMethods)(c))
+		if err != nil {
+			return nil, err
+		}
+		extra, err := toJSONRawMessages(c.extra)
+		if err != nil {
+			return nil, err
+		}
+
+		for k, v := range extra {
+			if _, ok := geth[k]; ok {
+				return nil, fmt.Errorf("duplicate JSON key %q in both %T and registered extra", k, c)
+			}
+			geth[k] = v
+		}
+		return json.Marshal(geth)
+	}
+}
+
+// marshalJSONWithExtra is the inverse of unmarshalJSONWithExtra().
+func (c *ChainConfig) marshalJSONWithExtra() ([]byte, error) {
+	cc := &chainConfigWithExportedExtra{
+		chainConfigWithoutMethods: (*chainConfigWithoutMethods)(c),
+		Extra:                     c.extra,
+	}
+	return json.Marshal(cc)
+}
+
+func toJSONRawMessages(v any) (map[string]json.RawMessage, error) {
+	buf, err := json.Marshal(v)
+	if err != nil {
+		return nil, err
+	}
+	msgs := make(map[string]json.RawMessage)
+	if err := json.Unmarshal(buf, &msgs); err != nil {
+		return nil, err
+	}
+	return msgs, nil
+}

params/json.libevm_test.go

@@ -0,0 +1,98 @@
+package params
+
+import (
+	"bytes"
+	"encoding/json"
+	"math/big"
+	"testing"
+
+	"github.com/ethereum/go-ethereum/libevm/pseudo"
+	"github.com/stretchr/testify/require"
+)
+
+type nestedChainConfigExtra struct {
+	NestedFoo string `json:"foo"`
+
+	NOOPHooks
+}
+
+type rootJSONChainConfigExtra struct {
+	TopLevelFoo string `json:"foo"`
+
+	NOOPHooks
+}
+
+func TestChainConfigJSONRoundTrip(t *testing.T) {
+	tests := []struct {
+		name      string
+		register  func()
+		jsonInput string
+		want      *ChainConfig
+	}{
+		{
+			name:     "no registered extras",
+			register: func() {},
+			jsonInput: `{
+				"chainId": 1234
+			}`,
+			want: &ChainConfig{
+				ChainID: big.NewInt(1234),
+			},
+		},
+		{
+			name: "reuse top-level JSON",
+			register: func() {
+				RegisterExtras(Extras[rootJSONChainConfigExtra, NOOPHooks]{
+					ReuseJSONRoot: true,
+				})
+			},
+			jsonInput: `{
+				"chainId": 5678,
+				"foo": "hello"
+			}`,
+			want: &ChainConfig{
+				ChainID: big.NewInt(5678),
+				extra:   pseudo.From(&rootJSONChainConfigExtra{TopLevelFoo: "hello"}).Type,
+			},
+		},
+		{
+			name: "nested JSON",
+			register: func() {
+				RegisterExtras(Extras[nestedChainConfigExtra, NOOPHooks]{
+					ReuseJSONRoot: false, // explicit zero value only for tests
+				})
+			},
+			jsonInput: `{
+				"chainId": 42,
+				"extra": {"foo": "world"}
+			}`,
+			want: &ChainConfig{
+				ChainID: big.NewInt(42),
+				extra:   pseudo.From(&nestedChainConfigExtra{NestedFoo: "world"}).Type,
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			TestOnlyClearRegisteredExtras()
+			t.Cleanup(TestOnlyClearRegisteredExtras)
+			tt.register()
+
+			t.Run("json.Unmarshal()", func(t *testing.T) {
+				got := new(ChainConfig)
+				require.NoError(t, json.Unmarshal([]byte(tt.jsonInput), got))
+				require.Equal(t, tt.want, got)
+			})
+
+			t.Run("json.Marshal()", func(t *testing.T) {
+				var want bytes.Buffer
+				require.NoError(t, json.Compact(&want, []byte(tt.jsonInput)), "json.Compact()")
+
+				got, err := json.Marshal(tt.want)
+				require.NoError(t, err, "json.Marshal()")
+				require.Equal(t, want.String(), string(got))
+			})
+		})
+	}
+}