From 0c0326bacf1ad4329d049ef5b88c48dfebb47df4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bastian=20M=C3=BCller?= <bastian@turbolent.com>
Date: Fri, 24 Jan 2025 13:51:16 -0800
Subject: [PATCH] optionally dump all read and written registers

---
 cmd/util/cmd/debug-tx/cmd.go  | 41 +++++++++++++++++++++++++++++++----
 utils/debug/README.md         |  6 ++---
 utils/debug/remoteDebugger.go | 12 +++++++---
 3 files changed, 49 insertions(+), 10 deletions(-)

diff --git a/cmd/util/cmd/debug-tx/cmd.go b/cmd/util/cmd/debug-tx/cmd.go
index c4e0393b9c3..c940218de04 100644
--- a/cmd/util/cmd/debug-tx/cmd.go
+++ b/cmd/util/cmd/debug-tx/cmd.go
@@ -1,12 +1,15 @@
 package debug_tx
 
 import (
+	"cmp"
 	"context"
+	"encoding/hex"
 
 	"github.com/onflow/flow/protobuf/go/flow/execution"
 	"github.com/onflow/flow/protobuf/go/flow/executiondata"
 	"github.com/rs/zerolog/log"
 	"github.com/spf13/cobra"
+	"golang.org/x/exp/slices"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials/insecure"
 
@@ -29,6 +32,7 @@ var (
 	flagComputeLimit        uint64
 	flagProposalKeySeq      uint64
 	flagUseExecutionDataAPI bool
+	flagDumpRegisters       bool
 )
 
 var Cmd = &cobra.Command{
@@ -61,6 +65,8 @@ func init() {
 	Cmd.Flags().Uint64Var(&flagProposalKeySeq, "proposal-key-seq", 0, "proposal key sequence number")
 
 	Cmd.Flags().BoolVar(&flagUseExecutionDataAPI, "use-execution-data-api", false, "use the execution data API")
+
+	Cmd.Flags().BoolVar(&flagDumpRegisters, "dump-registers", false, "dump registers")
 }
 
 func run(*cobra.Command, []string) {
@@ -187,11 +193,38 @@ func run(*cobra.Command, []string) {
 		proposalKeySequenceNumber,
 	)
 
-	txErr, processErr := debugger.RunTransaction(txBody, snap, header)
-	if txErr != nil {
-		log.Err(txErr).Msg("transaction error")
-	}
+	resultSnapshot, txErr, processErr := debugger.RunTransaction(txBody, snap, header)
 	if processErr != nil {
 		log.Fatal().Err(processErr).Msg("process error")
 	}
+
+	if flagDumpRegisters {
+		log.Info().Msg("Read registers:")
+		readRegisterIDs := resultSnapshot.ReadRegisterIDs()
+		sortRegisters(readRegisterIDs)
+		for _, registerID := range readRegisterIDs {
+			log.Info().Msgf("\t%s", registerID)
+		}
+
+		log.Info().Msg("Written registers:")
+		for _, updatedRegister := range resultSnapshot.UpdatedRegisters() {
+			log.Info().Msgf(
+				"\t%s, %s",
+				updatedRegister.Key,
+				hex.EncodeToString(updatedRegister.Value),
+			)
+		}
+	}
+	if txErr != nil {
+		log.Err(txErr).Msg("transaction error")
+	}
+}
+
+func sortRegisters(registerIDs []flow.RegisterID) {
+	slices.SortFunc(registerIDs, func(a, b flow.RegisterID) int {
+		return cmp.Or(
+			cmp.Compare(a.Owner, b.Owner),
+			cmp.Compare(a.Key, b.Key),
+		)
+	})
 }
diff --git a/utils/debug/README.md b/utils/debug/README.md
index 6d85a94a6c1..6bdc182bba0 100644
--- a/utils/debug/README.md
+++ b/utils/debug/README.md
@@ -89,7 +89,7 @@ func TestDebugger_RunTransactionAgainstExecutionNodeAtBlockID(t *testing.T) {
 
 	txBody := getTransaction(chain)
 
-	txErr, err := debugger.RunTransaction(txBody, snapshot, header)
+	_, txErr, err := debugger.RunTransaction(txBody, snapshot, header)
 	require.NoError(t, txErr)
 	require.NoError(t, err)
 }
@@ -137,12 +137,12 @@ func TestDebugger_RunTransactionAgainstAccessNodeAtBlockIDWithFileCache(t *testi
 	txBody := getTransaction(chain)
 
 	// the first run will cache the results
-	txErr, err := debugger.RunTransaction(txBody, snapshot, header)
+	_, txErr, err := debugger.RunTransaction(txBody, snapshot, header)
 	require.NoError(t, txErr)
 	require.NoError(t, err)
 
 	// the second run should only use the cache.
-	txErr, err = debugger.RunTransaction(txBody, snapshot, header)
+	_, txErr, err = debugger.RunTransaction(txBody, snapshot, header)
 	require.NoError(t, txErr)
 	require.NoError(t, err)
 }
diff --git a/utils/debug/remoteDebugger.go b/utils/debug/remoteDebugger.go
index 4bb4d35408c..524f7ed5dc3 100644
--- a/utils/debug/remoteDebugger.go
+++ b/utils/debug/remoteDebugger.go
@@ -5,6 +5,7 @@ import (
 	"github.com/rs/zerolog"
 
 	"github.com/onflow/flow-go/fvm"
+	"github.com/onflow/flow-go/fvm/storage/snapshot"
 	"github.com/onflow/flow-go/model/flow"
 )
 
@@ -43,6 +44,7 @@ func (d *RemoteDebugger) RunTransaction(
 	snapshot StorageSnapshot,
 	blockHeader *flow.Header,
 ) (
+	resultSnapshot *snapshot.ExecutionSnapshot,
 	txErr error,
 	processError error,
 ) {
@@ -52,11 +54,15 @@ func (d *RemoteDebugger) RunTransaction(
 
 	tx := fvm.Transaction(txBody, 0)
 
-	_, output, err := d.vm.Run(blockCtx, tx, snapshot)
+	var (
+		output fvm.ProcedureOutput
+		err    error
+	)
+	resultSnapshot, output, err = d.vm.Run(blockCtx, tx, snapshot)
 	if err != nil {
-		return nil, err
+		return resultSnapshot, nil, err
 	}
-	return output.Err, nil
+	return resultSnapshot, output.Err, nil
 }
 
 // RunScript runs the script using the given storage snapshot.