From 045f53e67a07c849806835817a36e87e77b21436 Mon Sep 17 00:00:00 2001
From: Guillaume Ballet <3272758+gballet@users.noreply.github.com>
Date: Fri, 4 Aug 2023 11:16:14 +0200
Subject: [PATCH] move UpdateContractCode to updateStateObject

---
 core/state/statedb.go | 41 ++---------------------------------------
 trie/verkle.go        | 31 ++++++++++++++++++++++++++++++-
 2 files changed, 32 insertions(+), 40 deletions(-)

diff --git a/core/state/statedb.go b/core/state/statedb.go
index 2a61ea7f11..f65dbb21c6 100644
--- a/core/state/statedb.go
+++ b/core/state/statedb.go
@@ -18,7 +18,6 @@
 package state
 
 import (
-	"encoding/binary"
 	"errors"
 	"fmt"
 	"math/big"
@@ -37,9 +36,6 @@ import (
 	"github.com/ethereum/go-ethereum/trie"
 	"github.com/ethereum/go-ethereum/trie/trienode"
 	"github.com/ethereum/go-ethereum/trie/triestate"
-	"github.com/ethereum/go-ethereum/trie/utils"
-	"github.com/gballet/go-verkle"
-	"github.com/holiman/uint256"
 )
 
 type revision struct {
@@ -591,40 +587,8 @@ func (s *StateDB) updateStateObject(obj *stateObject) {
 	if err := s.trie.UpdateAccount(addr, &obj.data); err != nil {
 		s.setError(fmt.Errorf("updateStateObject (%x) error: %v", addr[:], err))
 	}
-	if s.trie.IsVerkle() && obj.dirtyCode {
-		var (
-			chunks = trie.ChunkifyCode(obj.code)
-			values [][]byte
-			key    []byte
-			err    error
-		)
-		for i, chunknr := 0, uint64(0); i < len(chunks); i, chunknr = i+32, chunknr+1 {
-			groupOffset := (chunknr + 128) % 256
-			if groupOffset == 0 /* start of new group */ || chunknr == 0 /* first chunk in header group */ {
-				values = make([][]byte, verkle.NodeWidth)
-				key = utils.GetTreeKeyCodeChunkWithEvaluatedAddress(obj.db.db.(*cachingDB).GetTreeKeyHeader(obj.address[:]), uint256.NewInt(chunknr))
-			}
-			values[groupOffset] = chunks[i : i+32]
-
-			// Reuse the calculated key to also update the code size.
-			if i == 0 {
-				cs := make([]byte, 32)
-				binary.LittleEndian.PutUint64(cs, uint64(len(obj.code)))
-				values[utils.CodeSizeLeafKey] = cs
-			}
-
-			if groupOffset == 255 || len(chunks)-i <= 32 {
-				switch t := s.trie.(type) {
-				case *trie.VerkleTrie:
-					err = t.UpdateStem(key[:31], values)
-				case *trie.TransitionTrie:
-					err = t.UpdateStem(key[:31], values)
-				}
-				if err != nil {
-					s.setError(fmt.Errorf("updateStateObject (%x) error: %w", addr[:], err))
-				}
-			}
-		}
+	if obj.dirtyCode {
+		s.trie.UpdateContractCode(obj.Address(), common.BytesToHash(obj.CodeHash()), obj.code)
 	}
 	// Cache the data until commit. Note, this update mechanism is not symmetric
 	// to the deletion, because whereas it is enough to track account updates
@@ -1293,7 +1257,6 @@ func (s *StateDB) Commit(block uint64, deleteEmptyObjects bool) (common.Hash, er
 		// Write any contract code associated with the state object
 		if obj.code != nil && obj.dirtyCode {
 			rawdb.WriteCode(codeWriter, common.BytesToHash(obj.CodeHash()), obj.code)
-			s.trie.UpdateContractCode(obj.Address(), common.BytesToHash(obj.CodeHash()), obj.code)
 			obj.dirtyCode = false
 		}
 		// Write any storage changes in the state object to its storage trie
diff --git a/trie/verkle.go b/trie/verkle.go
index 4a87bc88ce..385af5287e 100644
--- a/trie/verkle.go
+++ b/trie/verkle.go
@@ -29,6 +29,7 @@ import (
 	"github.com/ethereum/go-ethereum/trie/trienode"
 	"github.com/ethereum/go-ethereum/trie/utils"
 	"github.com/gballet/go-verkle"
+	"github.com/holiman/uint256"
 )
 
 // VerkleTrie is a wrapper around VerkleNode that implements the trie.Trie
@@ -477,6 +478,34 @@ func (t *VerkleTrie) ClearStrorageRootConversion(addr common.Address) {
 }
 
 func (t *VerkleTrie) UpdateContractCode(addr common.Address, codeHash common.Hash, code []byte) error {
-	// XXX a copier depuis statedb/state_object
+	var (
+		chunks = ChunkifyCode(code)
+		values [][]byte
+		key    []byte
+		err    error
+	)
+	for i, chunknr := 0, uint64(0); i < len(chunks); i, chunknr = i+32, chunknr+1 {
+		groupOffset := (chunknr + 128) % 256
+		if groupOffset == 0 /* start of new group */ || chunknr == 0 /* first chunk in header group */ {
+			values = make([][]byte, verkle.NodeWidth)
+			key = utils.GetTreeKeyCodeChunkWithEvaluatedAddress(t.pointCache.GetTreeKeyHeader(addr[:]), uint256.NewInt(chunknr))
+		}
+		values[groupOffset] = chunks[i : i+32]
+
+		// Reuse the calculated key to also update the code size.
+		if i == 0 {
+			cs := make([]byte, 32)
+			binary.LittleEndian.PutUint64(cs, uint64(len(code)))
+			values[utils.CodeSizeLeafKey] = cs
+		}
+
+		if groupOffset == 255 || len(chunks)-i <= 32 {
+			err = t.UpdateStem(key[:31], values)
+
+			if err != nil {
+				return fmt.Errorf("UpdateContractCode (addr=%x) error: %w", addr[:], err)
+			}
+		}
+	}
 	return nil
 }