trie/bintrie: fix code chunk key encoding

trie/bintrie/key_encoding.go

@@ -90,8 +90,8 @@ func GetBinaryTreeKeyStorageSlot(address common.Address, key []byte) []byte {
 }
 
 func GetBinaryTreeKeyCodeChunk(address common.Address, chunknr *uint256.Int) []byte {
-	chunkOffset := new(uint256.Int).Add(codeOffset, chunknr).Bytes()
-	return GetBinaryTreeKey(address, chunkOffset)
+	chunkOffset := new(uint256.Int).Add(codeOffset, chunknr).Bytes32()
+	return GetBinaryTreeKey(address, chunkOffset[:])
 }
 
 func StorageIndex(storageKey []byte) (*uint256.Int, byte) {

trie/bintrie/key_encoding_test.go

@@ -0,0 +1,59 @@
+// Copyright 2025 The go-ethereum Authors
+// This file is part of the go-ethereum library.
+//
+// The go-ethereum library is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Lesser General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// The go-ethereum library is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Lesser General Public License for more details.
+//
+// You should have received a copy of the GNU Lesser General Public License
+// along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.
+
+package bintrie
+
+import (
+	"bytes"
+	"encoding/binary"
+	"testing"
+
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/holiman/uint256"
+)
+
+func TestGetBinaryTreeKeyCodeChunkBoundaries(t *testing.T) {
+	addr := common.HexToAddress("0x1234567890abcdef1234567890abcdef12345678")
+
+	for _, chunknr := range []uint64{0, 1, 127, 128, 255, 256, 257, 1024} {
+		got := GetBinaryTreeKeyCodeChunk(addr, uint256.NewInt(chunknr))
+
+		var offset [HashSize]byte
+		binary.BigEndian.PutUint64(offset[24:], chunknr+128)
+		want := GetBinaryTreeKey(addr, offset[:])
+
+		if !bytes.Equal(got, want) {
+			t.Fatalf("wrong code chunk key for chunk=%d: got=%x want=%x", chunknr, got, want)
+		}
+	}
+}
+
+func TestGetBinaryTreeKeyCodeChunkLargeIndex(t *testing.T) {
+	addr := common.HexToAddress("0x1234567890abcdef1234567890abcdef12345678")
+
+	var chunknr uint256.Int
+	chunknr.SetBytes(common.FromHex("0x0102030405060708090a0b0c0d0e0f10"))
+	got := GetBinaryTreeKeyCodeChunk(addr, &chunknr)
+
+	var offset uint256.Int
+	offset.Add(codeOffset, &chunknr)
+	offsetBytes := offset.Bytes32()
+	want := GetBinaryTreeKey(addr, offsetBytes[:])
+
+	if !bytes.Equal(got, want) {
+		t.Fatalf("wrong code chunk key for large chunk index: got=%x want=%x", got, want)
+	}
+}

trie/bintrie/trie.go

@@ -383,9 +383,9 @@ func (t *BinaryTrie) UpdateContractCode(addr common.Address, codeHash common.Has
 		groupOffset := (chunknr + 128) % StemNodeWidth
 		if groupOffset == 0 /* start of new group */ || chunknr == 0 /* first chunk in header group */ {
 			values = make([][]byte, StemNodeWidth)
-			var offset [HashSize]byte
-			binary.BigEndian.PutUint64(offset[24:], chunknr+128)
-			key = GetBinaryTreeKey(addr, offset[:])
+			var chunknrInt uint256.Int
+			chunknrInt.SetUint64(chunknr)
+			key = GetBinaryTreeKeyCodeChunk(addr, &chunknrInt)
 		}
 		values[groupOffset] = chunks[i : i+HashSize]
 

trie/bintrie/trie_test.go

@@ -267,6 +267,59 @@ func TestStorageRoundTrip(t *testing.T) {
 	}
 }
 
+// TestContractCodeRoundTrip verifies that UpdateContractCode stores the first
+// code chunk under the key returned by GetBinaryTreeKeyCodeChunk.
+func TestContractCodeRoundTrip(t *testing.T) {
+	tracer := trie.NewPrevalueTracer()
+	tr := &BinaryTrie{
+		root:   NewBinaryNode(),
+		tracer: tracer,
+	}
+	addr := common.HexToAddress("0x1234567890abcdef1234567890abcdef12345678")
+
+	acc := &types.StateAccount{
+		Nonce:    1,
+		Balance:  uint256.NewInt(1000),
+		CodeHash: common.HexToHash("c5d2460186f7233c927e7db2dcc703c0e500b653ca82273b7bfad8045d85a470").Bytes(),
+	}
+	if err := tr.UpdateAccount(addr, acc, 0); err != nil {
+		t.Fatalf("UpdateAccount error: %v", err)
+	}
+
+	code := common.FromHex("0x6001600055")
+	if err := tr.UpdateContractCode(addr, common.BytesToHash(acc.CodeHash), code); err != nil {
+		t.Fatalf("UpdateContractCode error: %v", err)
+	}
+
+	key := GetBinaryTreeKeyCodeChunk(addr, uint256.NewInt(0))
+	var (
+		got []byte
+		err error
+	)
+	switch root := tr.root.(type) {
+	case *InternalNode:
+		got, err = root.Get(key, tr.nodeResolver)
+		if err != nil {
+			t.Fatalf("Get error: %v", err)
+		}
+	case *StemNode:
+		values, err := root.GetValuesAtStem(key[:StemSize], nil)
+		if err != nil {
+			t.Fatalf("GetValuesAtStem error: %v", err)
+		}
+		if values == nil {
+			t.Fatal("expected stem values for code chunk key")
+		}
+		got = values[key[StemSize]]
+	default:
+		t.Fatalf("unexpected root type %T", tr.root)
+	}
+	want := ChunkifyCode(code)[:HashSize]
+	if !bytes.Equal(got, want) {
+		t.Fatalf("wrong chunk for chunk #0 key: got=%x want=%x", got, want)
+	}
+}
+
 func TestBinaryTrieWitness(t *testing.T) {
 	tracer := trie.NewPrevalueTracer()