From 9d6877b26ea3ca48b4288335afb2f91a25c79b62 Mon Sep 17 00:00:00 2001
From: Gary Rong <garyrong0905@gmail.com>
Date: Fri, 12 Jul 2024 16:10:29 +0800
Subject: [PATCH] trie: polish the code

---
 trie/verkle.go      | 36 +++++++++++++-----------------------
 trie/verkle_test.go |  2 +-
 2 files changed, 14 insertions(+), 24 deletions(-)

diff --git a/trie/verkle.go b/trie/verkle.go
index ab30c0397c..c1e09ad3b1 100644
--- a/trie/verkle.go
+++ b/trie/verkle.go
@@ -202,58 +202,48 @@ func (t *VerkleTrie) DeleteAccount(addr common.Address) error {
 // RollBackAccount removes the account info + code from the tree, unlike DeleteAccount
 // that will overwrite it with 0s. The first 64 storage slots are also removed.
 func (t *VerkleTrie) RollBackAccount(addr common.Address) error {
-	evaluatedAddr := t.cache.Get(addr.Bytes())
-	codesizekey := utils.CodeSizeKeyWithEvaluatedAddress(evaluatedAddr)
-	codesizeBytes, err := t.root.Get(codesizekey, t.nodeResolver)
+	var (
+		evaluatedAddr = t.cache.Get(addr.Bytes())
+		codeSizeKey   = utils.CodeSizeKeyWithEvaluatedAddress(evaluatedAddr)
+	)
+	codeSizeBytes, err := t.root.Get(codeSizeKey, t.nodeResolver)
 	if err != nil {
 		return fmt.Errorf("rollback: error finding code size: %w", err)
 	}
-	codesize := binary.LittleEndian.Uint64(codesizeBytes)
+	codeSize := binary.LittleEndian.Uint64(codeSizeBytes)
 
 	// Delete the account header + first 64 slots + first 128 code chunks
+	key := common.CopyBytes(codeSizeKey)
 	for i := 0; i < verkle.NodeWidth; i++ {
-		codesizekey[31] = byte(i)
+		key[31] = byte(i)
 
 		// this is a workaround to avoid deleting nil leaves, the lib needs to be
 		// fixed to be able to handle that
-		v, err := t.root.Get(codesizekey, t.nodeResolver)
+		v, err := t.root.Get(key, t.nodeResolver)
 		if err != nil {
 			return fmt.Errorf("error rolling back account header: %w", err)
 		}
 		if len(v) == 0 {
 			continue
 		}
-
-		_, err = t.root.Delete(codesizekey, t.nodeResolver)
+		_, err = t.root.Delete(key, t.nodeResolver)
 		if err != nil {
 			return fmt.Errorf("error rolling back account header: %w", err)
 		}
 	}
-
-	var root *verkle.InternalNode
-	switch r := t.root.(type) {
-	case *verkle.InternalNode:
-		root = r
-	default:
-		return errInvalidRootType
-	}
-
 	// Delete all further code
-	var key []byte
-	for i, chunknr := uint64(32*128), uint64(128); i < codesize; i, chunknr = i+32, chunknr+1 {
+	for i, chunknr := uint64(32*128), uint64(128); i < codeSize; i, chunknr = i+32, chunknr+1 {
 		// evaluate group key at the start of a new group
 		groupOffset := (chunknr + 128) % 256
 		if groupOffset == 0 {
 			key = utils.CodeChunkKeyWithEvaluatedAddress(evaluatedAddr, uint256.NewInt(chunknr))
 		}
-
 		key[31] = byte(groupOffset)
-		_, err = root.Delete(key[:], t.nodeResolver)
+		_, err = t.root.Delete(key[:], t.nodeResolver)
 		if err != nil {
-			return fmt.Errorf("RollbackContractCode (addr=%x) error: %w", addr[:], err)
+			return fmt.Errorf("error deleting code chunk (addr=%x) error: %w", addr[:], err)
 		}
 	}
-
 	return nil
 }
 
diff --git a/trie/verkle_test.go b/trie/verkle_test.go
index b83d32325f..55438d45e1 100644
--- a/trie/verkle_test.go
+++ b/trie/verkle_test.go
@@ -104,7 +104,6 @@ func TestVerkleRollBack(t *testing.T) {
 				t.Fatalf("Failed to update account, %v", err)
 			}
 		}
-
 		// create more than 128 chunks of code
 		code := make([]byte, 129*32)
 		for i := 0; i < len(code); i += 2 {
@@ -113,6 +112,7 @@ func TestVerkleRollBack(t *testing.T) {
 		}
 		hash := crypto.Keccak256Hash(code)
 		if err := tr.UpdateContractCode(addr, hash, code); err != nil {
+			t.Fatalf("Failed to update contract, %v", err)
 		}
 	}