store the offsets aside from the code

core/rawdb/accessors_state.go

@@ -61,13 +61,25 @@ 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 {
 		log.Crit("Failed to store contract code", "err", err)
 	}
 }
 
+// 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 {

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) {

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.

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)
 					}

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
 }