From 62e743d5303a3d3bccccac7b90a211122f8565a1 Mon Sep 17 00:00:00 2001
From: Guillaume Ballet <3272758+gballet@users.noreply.github.com>
Date: Fri, 14 Jan 2022 11:27:56 +0100
Subject: [PATCH] store the offsets aside from the code

---
 core/rawdb/accessors_state.go | 12 ++++++++++++
 core/rawdb/schema.go          |  6 ++++++
 core/state/database.go        | 17 +++++++++++++++--
 core/state/statedb.go         |  8 ++++++--
 trie/verkle.go                |  9 +++++++--
 5 files changed, 46 insertions(+), 6 deletions(-)

diff --git a/core/rawdb/accessors_state.go b/core/rawdb/accessors_state.go
index 6112de03ad53..9b320cf342ad 100644
--- a/core/rawdb/accessors_state.go
+++ b/core/rawdb/accessors_state.go
@@ -61,6 +61,11 @@ func ReadCodeWithPrefix(db ethdb.KeyValueReader, hash common.Hash) []byte {
 	return data
 }
 
+func ReadPushDataOffsets(db ethdb.KeyValueReader, hash common.Hash) []byte {
+	data, _ := db.Get(pdKey(hash))
+	return data
+}
+
 // WriteCode writes the provided contract code database.
 func WriteCode(db ethdb.KeyValueWriter, hash common.Hash, code []byte) {
 	if err := db.Put(codeKey(hash), code); err != nil {
@@ -68,6 +73,13 @@ func WriteCode(db ethdb.KeyValueWriter, hash common.Hash, code []byte) {
 	}
 }
 
+// WritePushDataOffsets writes the provided contract code database.
+func WritePushDataOffsets(db ethdb.KeyValueWriter, hash common.Hash, offsets []byte) {
+	if err := db.Put(pdKey(hash), offsets); err != nil {
+		log.Crit("Failed to store pushdata offsets for contract code", "err", err)
+	}
+}
+
 // DeleteCode deletes the specified contract code from the database.
 func DeleteCode(db ethdb.KeyValueWriter, hash common.Hash) {
 	if err := db.Delete(codeKey(hash)); err != nil {
diff --git a/core/rawdb/schema.go b/core/rawdb/schema.go
index b35fcba45f79..38238e2cb18d 100644
--- a/core/rawdb/schema.go
+++ b/core/rawdb/schema.go
@@ -220,6 +220,12 @@ func codeKey(hash common.Hash) []byte {
 	return append(CodePrefix, hash.Bytes()...)
 }
 
+// pdKey = CodePrefix + hash + "pd", so that it's right after the
+// code in the database.
+func pdKey(hash common.Hash) []byte {
+	return append(codeKey(hash), "pd"...)
+}
+
 // IsCodeKey reports whether the given byte slice is the key of contract code,
 // if so return the raw code hash as well.
 func IsCodeKey(key []byte) (bool, []byte) {
diff --git a/core/state/database.go b/core/state/database.go
index ee7d2c9a37a0..6968b2087a3c 100644
--- a/core/state/database.go
+++ b/core/state/database.go
@@ -38,6 +38,10 @@ const (
 	codeCacheSize = 64 * 1024 * 1024
 )
 
+var (
+	errCodeNotFound = errors.New("no code found")
+)
+
 // Database wraps access to tries and contract code.
 type Database interface {
 	// OpenTrie opens the main account trie.
@@ -181,7 +185,7 @@ func (db *cachingDB) ContractCode(addrHash, codeHash common.Hash) ([]byte, error
 		db.codeSizeCache.Add(codeHash, len(code))
 		return code, nil
 	}
-	return nil, errors.New("not found")
+	return nil, errCodeNotFound
 }
 
 // ContractCodeWithPrefix retrieves a particular contract's code. If the
@@ -265,7 +269,16 @@ func (db *VerkleDB) ContractCode(addrHash, codeHash common.Hash) ([]byte, error)
 		db.codeSizeCache.Add(codeHash, len(code))
 		return code, nil
 	}
-	return nil, errors.New("not found")
+	return nil, errCodeNotFound
+}
+
+func (db *VerkleDB) ContractCodePushData(codeHash common.Hash) ([]byte, error) {
+	pdoffsets := rawdb.ReadPushDataOffsets(db.db.DiskDB(), codeHash)
+	if len(pdoffsets) == 0 {
+		return nil, errCodeNotFound
+	}
+
+	return pdoffsets, nil
 }
 
 // ContractCodeSize retrieves a particular contracts code's size.
diff --git a/core/state/statedb.go b/core/state/statedb.go
index 49773021a76a..f6ed328dbf13 100644
--- a/core/state/statedb.go
+++ b/core/state/statedb.go
@@ -494,7 +494,9 @@ func (s *StateDB) updateStateObject(obj *stateObject) {
 		}
 
 		if obj.dirtyCode {
-			if chunks, err := trie.ChunkifyCode(addr, obj.code); err == nil {
+			// Since the DB isn't updated with the code, don't update
+			// the offsets either.
+			if _, chunks, err := trie.ChunkifyCode(addr, obj.code); err == nil {
 				for i := range chunks {
 					s.trie.TryUpdate(trieUtils.GetTreeKeyCodeChunk(addr[:], uint256.NewInt(uint64(i))), chunks[i][:])
 				}
@@ -996,10 +998,12 @@ func (s *StateDB) Commit(deleteEmptyObjects bool) (common.Hash, error) {
 			// Write any contract code associated with the state object
 			if obj.code != nil && obj.dirtyCode {
 				if s.trie.IsVerkle() {
-					if chunks, err := trie.ChunkifyCode(addr, obj.code); err == nil {
+					if offsets, chunks, err := trie.ChunkifyCode(addr, obj.code); err == nil {
 						for i := range chunks {
 							s.trie.TryUpdate(trieUtils.GetTreeKeyCodeChunk(addr[:], uint256.NewInt(uint64(i))), chunks[i][:])
 						}
+
+						rawdb.WritePushDataOffsets(codeWriter, common.BytesToHash(obj.CodeHash()), offsets)
 					} else {
 						s.setError(err)
 					}
diff --git a/trie/verkle.go b/trie/verkle.go
index b2561b6f6b5e..7ce5dbe35090 100644
--- a/trie/verkle.go
+++ b/trie/verkle.go
@@ -240,7 +240,11 @@ const (
 	PUSH32 = 0x71
 )
 
-func ChunkifyCode(addr common.Address, code []byte) ([][32]byte, error) {
+// ChunkifyCode returns the list of code offsets for the chunks, as well
+// as the generated chunks. The chunks are to be inserted into the tree,
+// and the code offsets are to be saved into the database.
+func ChunkifyCode(addr common.Address, code []byte) ([]byte, [][32]byte, error) {
+	var offsets []byte
 	lastOffset := byte(0)
 	chunkCount := len(code) / 31
 	if len(code)%31 != 0 {
@@ -260,7 +264,8 @@ func ChunkifyCode(addr common.Address, code []byte) ([][32]byte, error) {
 			}
 		}
 		chunks[i][0] = lastOffset
+		offsets = append(offsets, lastOffset)
 	}
 
-	return chunks, nil
+	return offsets, chunks, nil
 }