From 79d2327771e5a10e363bdd6542c624b8060bac7b Mon Sep 17 00:00:00 2001
From: Guillaume Ballet <3272758+gballet@users.noreply.github.com>
Date: Mon, 15 Jul 2024 09:05:59 +0200
Subject: [PATCH] trie: add RollBackAccount function to verkle trees (#30135)

---
 trie/verkle.go      | 51 ++++++++++++++++++++++++++++
 trie/verkle_test.go | 82 +++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 133 insertions(+)

diff --git a/trie/verkle.go b/trie/verkle.go
index a457097e95..fb4d81281c 100644
--- a/trie/verkle.go
+++ b/trie/verkle.go
@@ -199,6 +199,57 @@ func (t *VerkleTrie) DeleteAccount(addr common.Address) error {
 	return nil
 }
 
+// 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 {
+	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)
+	}
+	if len(codeSizeBytes) == 0 {
+		return errors.New("rollback: code size is not existent")
+	}
+	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++ {
+		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(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(key, t.nodeResolver)
+		if err != nil {
+			return fmt.Errorf("error rolling back account header: %w", err)
+		}
+	}
+	// Delete all further code
+	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 = t.root.Delete(key[:], t.nodeResolver)
+		if err != nil {
+			return fmt.Errorf("error deleting code chunk (addr=%x) error: %w", addr[:], err)
+		}
+	}
+	return nil
+}
+
 // DeleteStorage implements state.Trie, deleting the specified storage slot from
 // the trie. If the storage slot was not existent in the trie, no error will be
 // returned. If the trie is corrupted, an error will be returned.
diff --git a/trie/verkle_test.go b/trie/verkle_test.go
index 0cbe28bf01..55438d45e1 100644
--- a/trie/verkle_test.go
+++ b/trie/verkle_test.go
@@ -24,6 +24,7 @@ import (
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/core/rawdb"
 	"github.com/ethereum/go-ethereum/core/types"
+	"github.com/ethereum/go-ethereum/crypto"
 	"github.com/ethereum/go-ethereum/trie/utils"
 	"github.com/holiman/uint256"
 )
@@ -89,3 +90,84 @@ func TestVerkleTreeReadWrite(t *testing.T) {
 		}
 	}
 }
+
+func TestVerkleRollBack(t *testing.T) {
+	db := newTestDatabase(rawdb.NewMemoryDatabase(), rawdb.PathScheme)
+	tr, _ := NewVerkleTrie(types.EmptyVerkleHash, db, utils.NewPointCache(100))
+
+	for addr, acct := range accounts {
+		if err := tr.UpdateAccount(addr, acct); err != nil {
+			t.Fatalf("Failed to update account, %v", err)
+		}
+		for key, val := range storages[addr] {
+			if err := tr.UpdateStorage(addr, key.Bytes(), val); err != nil {
+				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 {
+			code[i] = 0x60
+			code[i+1] = byte(i % 256)
+		}
+		hash := crypto.Keccak256Hash(code)
+		if err := tr.UpdateContractCode(addr, hash, code); err != nil {
+			t.Fatalf("Failed to update contract, %v", err)
+		}
+	}
+
+	// Check that things were created
+	for addr, acct := range accounts {
+		stored, err := tr.GetAccount(addr)
+		if err != nil {
+			t.Fatalf("Failed to get account, %v", err)
+		}
+		if !reflect.DeepEqual(stored, acct) {
+			t.Fatal("account is not matched")
+		}
+		for key, val := range storages[addr] {
+			stored, err := tr.GetStorage(addr, key.Bytes())
+			if err != nil {
+				t.Fatalf("Failed to get storage, %v", err)
+			}
+			if !bytes.Equal(stored, val) {
+				t.Fatal("storage is not matched")
+			}
+		}
+	}
+
+	// ensure there is some code in the 2nd group
+	keyOf2ndGroup := []byte{141, 124, 185, 236, 50, 22, 185, 39, 244, 47, 97, 209, 96, 235, 22, 13, 205, 38, 18, 201, 128, 223, 0, 59, 146, 199, 222, 119, 133, 13, 91, 0}
+	chunk, err := tr.root.Get(keyOf2ndGroup, nil)
+	if err != nil {
+		t.Fatalf("Failed to get account, %v", err)
+	}
+	if len(chunk) == 0 {
+		t.Fatal("account was not created ")
+	}
+
+	// Rollback first account and check that it is gone
+	addr1 := common.Address{1}
+	err = tr.RollBackAccount(addr1)
+	if err != nil {
+		t.Fatalf("error rolling back address 1: %v", err)
+	}
+
+	// ensure the account is gone
+	stored, err := tr.GetAccount(addr1)
+	if err != nil {
+		t.Fatalf("Failed to get account, %v", err)
+	}
+	if stored != nil {
+		t.Fatal("account was not deleted")
+	}
+
+	// ensure that the last code chunk is also gone from the tree
+	chunk, err = tr.root.Get(keyOf2ndGroup, nil)
+	if err != nil {
+		t.Fatalf("Failed to get account, %v", err)
+	}
+	if len(chunk) != 0 {
+		t.Fatal("account was not deleted")
+	}
+}