From cac4f6a4a6f153cb66d05bf015ccb980ec274aca Mon Sep 17 00:00:00 2001
From: Anna Shaleva <shaleva.ann@nspcc.ru>
Date: Wed, 5 Oct 2022 13:56:26 +0300
Subject: [PATCH] cli: allow to dump storage diff for 'storage' VM CLI command

---
 cli/vm/cli.go      | 24 ++++++++++++++++++------
 cli/vm/cli_test.go | 30 ++++++++++++++++++++++++++++++
 2 files changed, 48 insertions(+), 6 deletions(-)

diff --git a/cli/vm/cli.go b/cli/vm/cli.go
index fdcec3c2d..a06c1ed08 100644
--- a/cli/vm/cli.go
+++ b/cli/vm/cli.go
@@ -58,6 +58,7 @@ const (
 	verboseFlagFullName   = "verbose"
 	historicFlagFullName  = "historic"
 	backwardsFlagFullName = "backwards"
+	diffFlagFullName      = "diff"
 )
 
 var historicFlag = cli.IntFlag{
@@ -275,24 +276,30 @@ Example:
 			"Can be used if no script is loaded. " +
 			"Hex-encoded storage items prefix may be specified (empty by default to return the whole set of storage items). " +
 			"If seek prefix is not empty, then it's trimmed from the resulting keys." +
-			"Items are sorted. Backwards seek direction may be specified (false by default, which means forwards storage seek direction).",
-		UsageText: `storage <hash-or-address-or-id> [<prefix>] [--backwards]`,
+			"Items are sorted. Backwards seek direction may be specified (false by default, which means forwards storage seek direction). " +
+			"It is possible to dump only those storage items that were added or changed during current script invocation (use --diff flag for it).",
+		UsageText: `storage <hash-or-address-or-id> [<prefix>] [--backwards] [--diff]`,
 		Flags: []cli.Flag{
 			cli.BoolFlag{
 				Name:  backwardsFlagFullName + ",b",
 				Usage: "Backwards traversal direction",
 			},
+			cli.BoolFlag{
+				Name:  diffFlagFullName + ",d",
+				Usage: "Dump only those storage items that were added or changed during the current script invocation. Note that this call won't show removed storage items.",
+			},
 		},
-		Description: `storage <hash-or-address-or-id> <prefix> --backwards
+		Description: `storage <hash-or-address-or-id> <prefix> [--backwards] [--diff]
 
 Dump storage of the contract with the specified hash, address or ID as is at the current stage of script invocation.
 Can be used if no script is loaded.
 Hex-encoded storage items prefix may be specified (empty by default to return the whole set of storage items).
 If seek prefix is not empty, then it's trimmed from the resulting keys.
 Items are sorted. Backwards seek direction may be specified (false by default, which means forwards storage seek direction).
+It is possible to dump only those storage items that were added or changed during current script invocation (use --diff flag for it).
 
 Example:
-> storage 0x0000000009070e030d0f0e020d0c06050e030c02 030e --backwards`,
+> storage 0x0000000009070e030d0f0e020d0c06050e030c02 030e --backwards --diff`,
 		Action: handleStorage,
 	},
 }
@@ -909,6 +916,7 @@ func handleStorage(c *cli.Context) error {
 		ic        = getInteropContextFromContext(c.App)
 		prefix    []byte
 		backwards bool
+		seekDepth int
 	)
 	h, err := flags.ParseAddress(hashOrID)
 	if err != nil {
@@ -933,9 +941,13 @@ func handleStorage(c *cli.Context) error {
 	if c.Bool(backwardsFlagFullName) {
 		backwards = true
 	}
+	if c.Bool(diffFlagFullName) {
+		seekDepth = 1 // take only upper DAO layer which stores only added or updated items.
+	}
 	ic.DAO.Seek(id, storage.SeekRange{
-		Prefix:    prefix,
-		Backwards: backwards,
+		Prefix:      prefix,
+		Backwards:   backwards,
+		SearchDepth: seekDepth,
 	}, func(k, v []byte) bool {
 		fmt.Fprintf(c.App.Writer, "%s: %v\n", hex.EncodeToString(k), hex.EncodeToString(v))
 		return true
diff --git a/cli/vm/cli_test.go b/cli/vm/cli_test.go
index 8c19b83b6..13df18241 100644
--- a/cli/vm/cli_test.go
+++ b/cli/vm/cli_test.go
@@ -880,3 +880,33 @@ func TestDumpStorage(t *testing.T) {
 	e.checkStorage(t, storage.KeyValue{Key: nil, Value: []byte{2}}) // empty key because search prefix is trimmed
 	e.checkStorage(t, expected[1], expected[0])
 }
+
+func TestDumpStorageDiff(t *testing.T) {
+	e := newTestVMClIWithState(t)
+
+	script := io.NewBufBinWriter()
+	h, err := e.cli.chain.GetContractScriptHash(1) // examples/storage/storage.go
+	require.NoError(t, err)
+	emit.AppCall(script.BinWriter, h, "put", callflag.All, 3, 3)
+
+	expected := []storage.KeyValue{
+		{Key: []byte{1}, Value: []byte{2}},
+		{Key: []byte{2}, Value: []byte{2}},
+	}
+	diff := storage.KeyValue{Key: []byte{3}, Value: []byte{3}}
+	e.runProg(t,
+		"storage 1",
+		"storage 1 --diff",
+		"loadhex "+hex.EncodeToString(script.Bytes()),
+		"run",
+		"storage 1",
+		"storage 1 --diff",
+	)
+
+	e.checkStorage(t, expected...)
+	// no script is executed => no diff
+	e.checkNextLine(t, "READY: loaded 37 instructions")
+	e.checkStack(t, 3)
+	e.checkStorage(t, append(expected, diff)...)
+	e.checkStorage(t, diff)
+}