From eede5923723619d47d5c209cd5ad8929fa7f8495 Mon Sep 17 00:00:00 2001
From: Martin Holst Swende <martin@swende.se>
Date: Mon, 4 Dec 2023 11:04:46 +0100
Subject: [PATCH] common/hexutil: define hex wrappers for uint256.Int

---
 common/hexutil/json.go      | 45 ++++++++++++++++++++++++++++
 common/hexutil/json_test.go | 59 +++++++++++++++++++++++++++++++++++++
 2 files changed, 104 insertions(+)

diff --git a/common/hexutil/json.go b/common/hexutil/json.go
index 50db208118..e0ac98f52d 100644
--- a/common/hexutil/json.go
+++ b/common/hexutil/json.go
@@ -23,6 +23,8 @@ import (
 	"math/big"
 	"reflect"
 	"strconv"
+
+	"github.com/holiman/uint256"
 )
 
 var (
@@ -30,6 +32,7 @@ var (
 	bigT    = reflect.TypeOf((*Big)(nil))
 	uintT   = reflect.TypeOf(Uint(0))
 	uint64T = reflect.TypeOf(Uint64(0))
+	u256T   = reflect.TypeOf((*uint256.Int)(nil))
 )
 
 // Bytes marshals/unmarshals as a JSON string with 0x prefix.
@@ -225,6 +228,48 @@ func (b *Big) UnmarshalGraphQL(input interface{}) error {
 	return err
 }
 
+// U256 marshals/unmarshals as a JSON string with 0x prefix.
+// The zero value marshals as "0x0".
+type U256 uint256.Int
+
+// MarshalText implements encoding.TextMarshaler
+func (b U256) MarshalText() ([]byte, error) {
+	u256 := (*uint256.Int)(&b)
+	return []byte(u256.Hex()), nil
+}
+
+// UnmarshalJSON implements json.Unmarshaler.
+func (b *U256) UnmarshalJSON(input []byte) error {
+	// The uint256.Int.UnmarshalJSON method accepts "dec", "0xhex"; we must be
+	// more strict, hence we check string and invoke SetFromHex directly.
+	if !isString(input) {
+		return errNonString(u256T)
+	}
+	// The hex decoder needs to accept empty string ("") as '0', which uint256.Int
+	// would reject.
+	if len(input) == 2 {
+		(*uint256.Int)(b).Clear()
+		return nil
+	}
+	err := (*uint256.Int)(b).SetFromHex(string(input[1 : len(input)-1]))
+	if err != nil {
+		return &json.UnmarshalTypeError{Value: err.Error(), Type: u256T}
+	}
+	return nil
+}
+
+// UnmarshalText implements encoding.TextUnmarshaler
+func (b *U256) UnmarshalText(input []byte) error {
+	// The uint256.Int.UnmarshalText method accepts "dec", "0xhex"; we must be
+	// more strict, hence we check string and invoke SetFromHex directly.
+	return (*uint256.Int)(b).SetFromHex(string(input))
+}
+
+// String returns the hex encoding of b.
+func (b *U256) String() string {
+	return (*uint256.Int)(b).Hex()
+}
+
 // Uint64 marshals/unmarshals as a JSON string with 0x prefix.
 // The zero value marshals as "0x0".
 type Uint64 uint64
diff --git a/common/hexutil/json_test.go b/common/hexutil/json_test.go
index ed7d6fad1a..766a63a4b2 100644
--- a/common/hexutil/json_test.go
+++ b/common/hexutil/json_test.go
@@ -21,6 +21,7 @@ import (
 	"encoding/hex"
 	"encoding/json"
 	"errors"
+	"github.com/holiman/uint256"
 	"math/big"
 	"testing"
 )
@@ -176,6 +177,64 @@ func TestUnmarshalBig(t *testing.T) {
 	}
 }
 
+var unmarshalU256Tests = []unmarshalTest{
+	// invalid encoding
+	{input: "", wantErr: errJSONEOF},
+	{input: "null", wantErr: errNonString(u256T)},
+	{input: "10", wantErr: errNonString(u256T)},
+	{input: `"0"`, wantErr: wrapTypeError(ErrMissingPrefix, u256T)},
+	{input: `"0x"`, wantErr: wrapTypeError(ErrEmptyNumber, u256T)},
+	{input: `"0x01"`, wantErr: wrapTypeError(ErrLeadingZero, u256T)},
+	{input: `"0xx"`, wantErr: wrapTypeError(ErrSyntax, u256T)},
+	{input: `"0x1zz01"`, wantErr: wrapTypeError(ErrSyntax, u256T)},
+	{
+		input:   `"0x10000000000000000000000000000000000000000000000000000000000000000"`,
+		wantErr: wrapTypeError(ErrBig256Range, u256T),
+	},
+
+	// valid encoding
+	{input: `""`, want: big.NewInt(0)},
+	{input: `"0x0"`, want: big.NewInt(0)},
+	{input: `"0x2"`, want: big.NewInt(0x2)},
+	{input: `"0x2F2"`, want: big.NewInt(0x2f2)},
+	{input: `"0X2F2"`, want: big.NewInt(0x2f2)},
+	{input: `"0x1122aaff"`, want: big.NewInt(0x1122aaff)},
+	{input: `"0xbBb"`, want: big.NewInt(0xbbb)},
+	{input: `"0xfffffffff"`, want: big.NewInt(0xfffffffff)},
+	{
+		input: `"0x112233445566778899aabbccddeeff"`,
+		want:  referenceBig("112233445566778899aabbccddeeff"),
+	},
+	{
+		input: `"0xffffffffffffffffffffffffffffffffffff"`,
+		want:  referenceBig("ffffffffffffffffffffffffffffffffffff"),
+	},
+	{
+		input: `"0xffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff"`,
+		want:  referenceBig("ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff"),
+	},
+}
+
+func TestUnmarshalU256(t *testing.T) {
+	for _, test := range unmarshalU256Tests {
+		var v U256
+		err := json.Unmarshal([]byte(test.input), &v)
+		if !checkError(t, test.input, err, test.wantErr) {
+			continue
+		}
+		if test.want == nil {
+			continue
+		}
+		want := new(uint256.Int)
+		want.SetFromBig(test.want.(*big.Int))
+		have := (*uint256.Int)(&v)
+		if want.Cmp(have) != 0 {
+			t.Errorf("input %s: value mismatch: have %x, want %x", test.input, have, want)
+			continue
+		}
+	}
+}
+
 func BenchmarkUnmarshalBig(b *testing.B) {
 	input := []byte(`"0x123456789abcdef123456789abcdef"`)
 	for i := 0; i < b.N; i++ {