From 6ff23b33b9cf9d90adf0004cc78877cb06c421a4 Mon Sep 17 00:00:00 2001
From: Jerry <jerrycgh@gmail.com>
Date: Fri, 25 Oct 2024 10:07:56 -0700
Subject: [PATCH] Change balance and nonce in zero's tracer output from decimal
 to hex (#1350)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* Change balance and nonce in zero's tracer output from decimal to hex

* test: add unit test for TxnTrace_MarshalJSON

---------

Co-authored-by: Stefan Negovanović <stefan@ethernal.tech>
---
 core/types/trace.go      | 48 +++++++++++++++++++++++++++++++++++++++
 core/types/trace_test.go | 49 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 97 insertions(+)
 create mode 100644 core/types/trace_test.go

diff --git a/core/types/trace.go b/core/types/trace.go
index 262e8ba21ae..979b7fe6fbc 100644
--- a/core/types/trace.go
+++ b/core/types/trace.go
@@ -2,6 +2,8 @@
 package types
 
 import (
+	"encoding/json"
+
 	"github.com/holiman/uint256"
 	libcommon "github.com/ledgerwatch/erigon-lib/common"
 	"github.com/ledgerwatch/erigon-lib/common/hexutility"
@@ -18,6 +20,17 @@ type ContractCodeUsage struct {
 	Write HexBytes        `json:"write,omitempty"`
 }
 
+type Uint256 struct {
+	uint256.Int
+}
+
+// MarshalJSON implements the json.Marshaler interface
+// It returns the hex representation of the uint256.Int.
+// We need to do this because the uint256.Int will be marshalled in decimal format by default.
+func (u *Uint256) MarshalJSON() ([]byte, error) {
+	return json.Marshal(u.Hex())
+}
+
 type TxnTrace struct {
 	Balance        *uint256.Int                    `json:"balance,omitempty"`
 	Nonce          *uint256.Int                    `json:"nonce,omitempty"`
@@ -28,6 +41,41 @@ type TxnTrace struct {
 	StorageReadMap map[libcommon.Hash]struct{}     `json:"-"`
 }
 
+type TxnTraceHex struct {
+	Balance        *Uint256                    `json:"balance,omitempty"`
+	Nonce          *Uint256                    `json:"nonce,omitempty"`
+	StorageRead    []libcommon.Hash            `json:"storage_read,omitempty"`
+	StorageWritten map[libcommon.Hash]*Uint256 `json:"storage_written,omitempty"`
+	CodeUsage      *ContractCodeUsage          `json:"code_usage,omitempty"`
+	SelfDestructed *bool                       `json:"self_destructed,omitempty"`
+	StorageReadMap map[libcommon.Hash]struct{} `json:"-"`
+}
+
+func (t *TxnTrace) MarshalJSON() ([]byte, error) {
+	tHex := TxnTraceHex{
+		StorageRead:    t.StorageRead,
+		StorageWritten: make(map[libcommon.Hash]*Uint256),
+		CodeUsage:      t.CodeUsage,
+		SelfDestructed: t.SelfDestructed,
+		StorageReadMap: t.StorageReadMap,
+	}
+
+	if t.Balance != nil {
+		tHex.Balance = &Uint256{*t.Balance}
+	}
+	if t.Nonce != nil {
+		tHex.Nonce = &Uint256{*t.Nonce}
+	}
+
+	for k, v := range t.StorageWritten {
+		if v != nil {
+			tHex.StorageWritten[k] = &Uint256{*v}
+		}
+	}
+
+	return json.Marshal(tHex)
+}
+
 type TxnMeta struct {
 	ByteCode           HexBytes `json:"byte_code,omitempty"`
 	NewTxnTrieNode     HexBytes `json:"new_txn_trie_node_byte,omitempty"`
diff --git a/core/types/trace_test.go b/core/types/trace_test.go
new file mode 100644
index 00000000000..76a3240f85c
--- /dev/null
+++ b/core/types/trace_test.go
@@ -0,0 +1,49 @@
+package types
+
+import (
+	"encoding/json"
+	"testing"
+
+	"github.com/holiman/uint256"
+	"github.com/ledgerwatch/erigon-lib/common"
+	"github.com/stretchr/testify/require"
+)
+
+func TestTxnTrace_MarshalJSON(t *testing.T) {
+	balance := uint256.NewInt(32)
+	nonce := uint256.NewInt(64)
+	storageRead := []common.Hash{common.HexToHash("0x1")}
+	selfDestructed := false
+	codeRead := common.HexToHash("0x6")
+
+	trace := TxnTrace{
+		Balance:     balance,
+		Nonce:       nonce,
+		StorageRead: storageRead,
+		StorageWritten: map[common.Hash]*uint256.Int{
+			common.HexToHash("0x7"): uint256.NewInt(7),
+			common.HexToHash("0x8"): uint256.NewInt(8),
+			common.HexToHash("0x9"): uint256.NewInt(9),
+		},
+		CodeUsage: &ContractCodeUsage{
+			Read: &codeRead,
+		},
+		SelfDestructed: &selfDestructed,
+	}
+
+	traceJSON, err := trace.MarshalJSON()
+	require.NoError(t, err)
+
+	// Parse the JSON output into a generic map and assert balance and nonce are hex encoded
+	var genericTraceResult map[string]interface{}
+	err = json.Unmarshal(traceJSON, &genericTraceResult)
+	require.NoError(t, err)
+
+	require.Equal(t, balance.Hex(), genericTraceResult["balance"])
+	require.Equal(t, nonce.Hex(), genericTraceResult["nonce"])
+
+	var resultTrace TxnTrace
+	err = json.Unmarshal(traceJSON, &resultTrace)
+	require.NoError(t, err)
+	require.Equal(t, trace, resultTrace)
+}