// Copyright 2014 The go-ethereum Authors
// (original work)
// Copyright 2024 The Erigon Authors
// (modifications)
// This file is part of Erigon.
//
// Erigon 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.
//
// Erigon 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 Erigon. If not, see <http://www.gnu.org/licenses/>.

package trie

import (
	"bytes"
	"encoding/binary"
	"encoding/hex"
	"fmt"
	"math/big"
	"math/rand"
	"testing"

	"github.com/holiman/uint256"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"github.com/erigontech/erigon/common"
	"github.com/erigontech/erigon/common/crypto"
	"github.com/erigontech/erigon/execution/commitment/nibbles"
	"github.com/erigontech/erigon/execution/rlp"
	"github.com/erigontech/erigon/execution/types/accounts"
)

func newEmpty() *Trie {
	return New(common.Hash{})
}

// testAccount builds an account with the given nonce and balance, defaulting Root to
// EmptyRoot and CodeHash to EmptyCodeHash; opts override the remaining fields.
func testAccount(nonce uint64, balance uint256.Int, opts ...func(*accounts.Account)) *accounts.Account {
	acc := accounts.NewAccount()
	acc.Nonce = nonce
	acc.Balance = balance
	for _, o := range opts {
		o(&acc)
	}
	return &acc
}

func withIncarnation(inc uint64) func(*accounts.Account) {
	return func(a *accounts.Account) { a.Incarnation = inc }
}

func withRoot(root common.Hash) func(*accounts.Account) {
	return func(a *accounts.Account) { a.Root = root }
}

func withCodeHash(codeHash common.Hash) func(*accounts.Account) {
	return func(a *accounts.Account) { a.CodeHash = accounts.InternCodeHash(codeHash) }
}

func TestEmptyTrie(t *testing.T) {
	var trie Trie
	res := trie.Hash()
	exp := EmptyRoot
	if res != exp {
		t.Errorf("expected %x got %x", exp, res)
	}
}

func TestNull(t *testing.T) {
	var trie Trie
	key := make([]byte, 32)
	value := []byte("test")
	trie.Update(key, value)
	v, _ := trie.Get(key)
	if !bytes.Equal(v, value) {
		t.Fatal("wrong value")
	}
}

func TestUpdateEmptyValueDeletes(t *testing.T) {
	trie := newEmpty()
	key := []byte("test-key")
	value := []byte("test-value")

	// 1. Insert a value
	trie.Update(key, value)
	v, ok := trie.Get(key)
	require.True(t, ok)
	require.Equal(t, value, v)

	// 2. Update with empty value (should delete)
	trie.Update(key, []byte{})

	// 3. Verify deletion
	v, ok = trie.Get(key)
	require.True(t, ok) // true because we resolved it fully without hitting a hash node
	require.Nil(t, v, "key should be deleted when updated with empty value")

	// 4. Insert again
	trie.Update(key, value)
	v, ok = trie.Get(key)
	require.True(t, ok)
	require.Equal(t, value, v)

	// 5. Update with nil (should delete)
	trie.Update(key, nil)

	// 6. Verify deletion
	v, ok = trie.Get(key)
	require.True(t, ok)
	require.Nil(t, v, "key should be deleted when updated with nil value")
}

func TestLargeValue(t *testing.T) {
	trie := newEmpty()
	trie.Update([]byte("key1"), []byte{99, 99, 99, 99})
	trie.Update([]byte("key2"), bytes.Repeat([]byte{1}, 32))
	trie.Hash()
}

// TestRandomCases tests som cases that were found via random fuzzing
func TestRandomCases(t *testing.T) {
	var rt = []randTestStep{
		{op: 6, key: common.Hex2Bytes(""), value: common.Hex2Bytes("")},                                                                                         // step 0
		{op: 6, key: common.Hex2Bytes(""), value: common.Hex2Bytes("")},                                                                                         // step 1
		{op: 0, key: common.Hex2Bytes("d51b182b95d677e5f1c82508c0228de96b73092d78ce78b2230cd948674f66fd1483bd"), value: common.Hex2Bytes("0000000000000002")},   // step 2
		{op: 2, key: common.Hex2Bytes("c2a38512b83107d665c65235b0250002882ac2022eb00711552354832c5f1d030d0e408e"), value: common.Hex2Bytes("")},                 // step 3
		{op: 3, key: common.Hex2Bytes(""), value: common.Hex2Bytes("")},                                                                                         // step 4
		{op: 3, key: common.Hex2Bytes(""), value: common.Hex2Bytes("")},                                                                                         // step 5
		{op: 6, key: common.Hex2Bytes(""), value: common.Hex2Bytes("")},                                                                                         // step 6
		{op: 3, key: common.Hex2Bytes(""), value: common.Hex2Bytes("")},                                                                                         // step 7
		{op: 0, key: common.Hex2Bytes("c2a38512b83107d665c65235b0250002882ac2022eb00711552354832c5f1d030d0e408e"), value: common.Hex2Bytes("0000000000000008")}, // step 8
		{op: 0, key: common.Hex2Bytes("d51b182b95d677e5f1c82508c0228de96b73092d78ce78b2230cd948674f66fd1483bd"), value: common.Hex2Bytes("0000000000000009")},   // step 9
		{op: 2, key: common.Hex2Bytes("fd"), value: common.Hex2Bytes("")},                                                                                       // step 10
		{op: 6, key: common.Hex2Bytes(""), value: common.Hex2Bytes("")},                                                                                         // step 11
		{op: 6, key: common.Hex2Bytes(""), value: common.Hex2Bytes("")},                                                                                         // step 12
		{op: 0, key: common.Hex2Bytes("fd"), value: common.Hex2Bytes("000000000000000d")},                                                                       // step 13
		{op: 6, key: common.Hex2Bytes(""), value: common.Hex2Bytes("")},                                                                                         // step 14
		{op: 1, key: common.Hex2Bytes("c2a38512b83107d665c65235b0250002882ac2022eb00711552354832c5f1d030d0e408e"), value: common.Hex2Bytes("")},                 // step 15
		{op: 3, key: common.Hex2Bytes(""), value: common.Hex2Bytes("")},                                                                                         // step 16
		{op: 0, key: common.Hex2Bytes("c2a38512b83107d665c65235b0250002882ac2022eb00711552354832c5f1d030d0e408e"), value: common.Hex2Bytes("0000000000000011")}, // step 17
		{op: 5, key: common.Hex2Bytes(""), value: common.Hex2Bytes("")},                                                                                         // step 18
		{op: 3, key: common.Hex2Bytes(""), value: common.Hex2Bytes("")},                                                                                         // step 19
		// FIXME: fix these testcases for Erigon
		//{op: 0, key: common.Hex2Bytes("d51b182b95d677e5f1c82508c0228de96b73092d78ce78b2230cd948674f66fd1483bd"), value: common.Hex2Bytes("0000000000000014")},           // step 20
		//{op: 0, key: common.Hex2Bytes("d51b182b95d677e5f1c82508c0228de96b73092d78ce78b2230cd948674f66fd1483bd"), value: common.Hex2Bytes("0000000000000015")},           // step 21
		//{op: 0, key: common.Hex2Bytes("c2a38512b83107d665c65235b0250002882ac2022eb00711552354832c5f1d030d0e408e"), value: common.Hex2Bytes("0000000000000016")},         // step 22
		{op: 5, key: common.Hex2Bytes(""), value: common.Hex2Bytes("")}, // step 23
		//{op: 1, key: common.Hex2Bytes("980c393656413a15c8da01978ed9f89feb80b502f58f2d640e3a2f5f7a99a7018f1b573befd92053ac6f78fca4a87268"), value: common.Hex2Bytes("")}, // step 24
		//{op: 1, key: common.Hex2Bytes("fd"), value: common.Hex2Bytes("")}, // step 25
	}
	runRandTest(rt)

}

// randTest is a sequence of random trie operations replayed by runRandTest.
type randTest []randTestStep

type randTestStep struct {
	op    int
	key   []byte // for opUpdate, opDelete, opGet
	value []byte // for opUpdate
	err   error  // for debugging
}

const (
	opUpdate = iota
	opDelete
	opGet
	opCommit
	opHash
	opReset
	opItercheckhash
	opCheckCacheInvariant
	opMax // boundary value, not an actual op
)

func runRandTest(rt randTest) bool {
	tr := newEmpty()
	values := make(map[string]string) // tracks content of the trie

	for i, step := range rt {
		switch step.op {
		case opUpdate:
			tr.Update(step.key, step.value)
			values[string(step.key)] = string(step.value)
		case opDelete:
			tr.Delete(step.key)
			delete(values, string(step.key))
		case opGet:
			v, _ := tr.Get(step.key)
			want := values[string(step.key)]
			if string(v) != want {
				rt[i].err = fmt.Errorf("mismatch for key 0x%x, got 0x%x want 0x%x", step.key, v, want)
			}
		case opCommit:
		case opHash:
			tr.Hash()
		case opReset:
			hash := tr.Hash()
			newtr := New(hash)
			tr = newtr
		case opItercheckhash:
			// FIXME: restore for Erigon
			/*
				checktr := newEmpty()
				it := NewIterator(tr.NodeIterator(nil))
				for it.Next() {
					checktr.Update(it.Key, it.Value)
				}
				if tr.Hash() != checktr.Hash() {
					rt[i].err = fmt.Errorf("hash mismatch in opItercheckhash")
				}
			*/
		case opCheckCacheInvariant:
		}
		// Abort the test on error.
		if rt[i].err != nil {
			return false
		}
	}
	return true
}

// Benchmarks the trie hashing. Since the trie caches the result of any operation,
// we cannot use b.N as the number of hashing rouns, since all rounds apart from
// the first one will be NOOP. As such, we'll use b.N as the number of account to
// insert into the trie before measuring the hashing.
func BenchmarkHash(b *testing.B) {
	// Make the random benchmark deterministic
	random := rand.New(rand.NewSource(0))

	// Create a realistic account trie to hash
	addresses := make([][20]byte, b.N)
	for i := range addresses {
		for j := range len(addresses[i]) {
			addresses[i][j] = byte(random.Intn(256))
		}
	}
	accounts := make([][]byte, len(addresses))
	for i := range accounts {
		var (
			nonce   = uint64(random.Int63())
			balance = new(big.Int).Rand(random, new(big.Int).Exp(common.Big2, common.Big256, nil))
			root    = EmptyRoot
			code    = crypto.Keccak256(nil)
		)
		accounts[i], _ = rlp.EncodeToBytes([]any{nonce, balance, root, code})
	}
	// Insert the accounts into the trie and hash it
	trie := newEmpty()
	for i := range addresses {
		trie.Update(crypto.Keccak256(addresses[i][:]), accounts[i])
	}
	b.ResetTimer()
	b.ReportAllocs()
	trie.Hash()
}

func TestDeepHash(t *testing.T) {
	acc := accounts.NewAccount()
	prefix := "prefix"
	var testdata = [][]struct {
		key   string
		value string
	}{
		{{"key1", "value1"}},
		{{"key1", "value1"}, {"key2", "value2"}},
		{{"key1", "value1"}, {"key2", "value2"}, {"key3", "value3"}},
		{{"key1", "value1"}, {"key2", "value2"}, {"\xffek3", "value3"}},
	}
	for i, keyVals := range testdata {
		trie := newEmpty()
		for _, keyVal := range keyVals {
			trie.Update([]byte(keyVal.key), []byte(keyVal.value))
		}
		hash1 := trie.Hash()

		prefixTrie := newEmpty()
		prefixTrie.UpdateAccount([]byte(prefix), &acc)
		for _, keyVal := range keyVals {
			// Add a prefix to every key
			prefixTrie.Update([]byte(prefix+keyVal.key), []byte(keyVal.value))
		}

		got2, hash2 := prefixTrie.DeepHash([]byte(prefix))
		if !got2 {
			t.Errorf("Expected DeepHash returning true, got false, testcase %d", i)
		}
		if hash1 != hash2 {
			t.Errorf("DeepHash mistmatch: %x, expected %x, testcase %d", hash2, hash1, i)
		}
	}
}

func genRandomByteArrayOfLen(length uint) []byte {
	array := make([]byte, length)
	for i := range length {
		array[i] = byte(rand.Intn(256))
	}
	return array
}

func getAddressForIndex(index int) [20]byte {
	var address [20]byte
	binary.BigEndian.PutUint32(address[:], uint32(index))
	return address
}

// randomAccountWithCode returns an account with random nonce and balance whose code hash
// commits to codeValue (the empty code hash when codeValue is nil).
func randomAccountWithCode(random *rand.Rand, codeValue []byte) accounts.Account {
	acc := accounts.NewAccount()
	acc.Nonce = uint64(random.Int63())
	acc.Balance.SetFromBig(new(big.Int).Rand(random, new(big.Int).Exp(common.Big2, common.Big256, nil)))
	acc.Root = EmptyRoot
	if codeValue == nil {
		acc.CodeHash = accounts.EmptyCodeHash
	} else {
		acc.CodeHash = accounts.InternCodeHash(crypto.Keccak256Hash(codeValue))
	}
	return acc
}

// insertAccountWithCode inserts an account at keccak(address) along with its code.
func insertAccountWithCode(t *testing.T, trie *Trie, random *rand.Rand, address [20]byte, codeValue []byte) accounts.Account {
	t.Helper()
	acc := randomAccountWithCode(random, codeValue)
	trie.UpdateAccount(crypto.Keccak256(address[:]), &acc)
	err := trie.UpdateAccountCode(crypto.Keccak256(address[:]), codeValue)
	require.NoError(t, err, "should successfully insert code")
	return acc
}

// newCodeTrie returns the common prologue for the TestCodeNode* tests: an empty trie, a
// zero-seeded RNG, and the address at index 0.
func newCodeTrie(t *testing.T) (*Trie, *rand.Rand, [20]byte) {
	t.Helper()
	return newEmpty(), rand.New(rand.NewSource(0)), getAddressForIndex(0)
}

func TestCodeNodeValid(t *testing.T) {
	trie, random, _ := newCodeTrie(t)

	numberOfAccounts := 20

	addresses := make([][20]byte, numberOfAccounts)
	for i := range addresses {
		addresses[i] = getAddressForIndex(i)
	}
	codeValues := make([][]byte, len(addresses))
	for i := range addresses {
		codeValues[i] = genRandomByteArrayOfLen(128)
		insertAccountWithCode(t, trie, random, addresses[i], codeValues[i])
	}

	for i := range addresses {
		value, gotValue := trie.GetAccountCode(crypto.Keccak256(addresses[i][:]))
		assert.True(t, gotValue, "should receive code value")
		assert.True(t, bytes.Equal(value, codeValues[i]), "should receive the right code")
	}
}

func TestCodeNodeUpdateNotExisting(t *testing.T) {
	trie, random, address := newCodeTrie(t)
	insertAccountWithCode(t, trie, random, address, genRandomByteArrayOfLen(128))

	nonExistingAddress := getAddressForIndex(9999)
	codeValue2 := genRandomByteArrayOfLen(128)

	err := trie.UpdateAccountCode(crypto.Keccak256(nonExistingAddress[:]), codeValue2)
	assert.Error(t, err, "should return an error for non existing acc")
}

func TestCodeNodeGetNotExistingAccount(t *testing.T) {
	trie, random, address := newCodeTrie(t)
	insertAccountWithCode(t, trie, random, address, genRandomByteArrayOfLen(128))

	nonExistingAddress := getAddressForIndex(9999)

	value, gotValue := trie.GetAccountCode(crypto.Keccak256(nonExistingAddress[:]))
	assert.True(t, gotValue, "should indicate that account doesn't exist at all (not just hashed)")
	assert.Nil(t, value, "the value should be nil")
}

func TestCodeNodeGetHashedAccount(t *testing.T) {
	trie, _, address := newCodeTrie(t)

	fakeAccount := genRandomByteArrayOfLen(50)
	fakeAccountHash := crypto.Keccak256Hash(fakeAccount)

	hex := nibbles.KeybytesToHex(crypto.Keccak256(address[:]))

	_, trie.RootNode = trie.insert(trie.RootNode, hex, &HashNode{hash: fakeAccountHash[:]})

	value, gotValue := trie.GetAccountCode(crypto.Keccak256(address[:]))
	assert.False(t, gotValue, "should indicate that account exists but hashed")
	assert.Nil(t, value, "the value should be nil")
}

func TestCodeNodeGetExistingAccountNoCodeNotEmpty(t *testing.T) {
	trie, random, address := newCodeTrie(t)

	acc := randomAccountWithCode(random, genRandomByteArrayOfLen(128))
	trie.UpdateAccount(crypto.Keccak256(address[:]), &acc)

	value, gotValue := trie.GetAccountCode(crypto.Keccak256(address[:]))
	assert.False(t, gotValue, "should indicate that account exists with code but the code isn't in cache")
	assert.Nil(t, value, "the value should be nil")
}

func TestCodeNodeGetExistingAccountEmptyCode(t *testing.T) {
	trie, random, address := newCodeTrie(t)

	acc := randomAccountWithCode(random, nil)
	trie.UpdateAccount(crypto.Keccak256(address[:]), &acc)

	value, gotValue := trie.GetAccountCode(crypto.Keccak256(address[:]))
	assert.True(t, gotValue, "should indicate that account exists with empty code")
	assert.Nil(t, value, "the value should be nil")
}

func TestCodeNodeWrongHash(t *testing.T) {
	trie, random, address := newCodeTrie(t)

	acc := randomAccountWithCode(random, genRandomByteArrayOfLen(128))
	trie.UpdateAccount(crypto.Keccak256(address[:]), &acc)

	codeValue2 := genRandomByteArrayOfLen(128)
	err := trie.UpdateAccountCode(crypto.Keccak256(address[:]), codeValue2)
	assert.Error(t, err, "should NOT be able to insert code with wrong hash")
}

func TestCodeNodeUpdateAccountAndCodeValidHash(t *testing.T) {
	trie, random, address := newCodeTrie(t)

	acc := insertAccountWithCode(t, trie, random, address, genRandomByteArrayOfLen(128))

	codeValue2 := genRandomByteArrayOfLen(128)
	acc.CodeHash = accounts.InternCodeHash(crypto.Keccak256Hash(codeValue2))

	trie.UpdateAccount(crypto.Keccak256(address[:]), &acc)
	err := trie.UpdateAccountCode(crypto.Keccak256(address[:]), codeValue2)
	require.NoError(t, err, "should successfully insert code")
}

func TestCodeNodeUpdateAccountAndCodeInvalidHash(t *testing.T) {
	trie, random, address := newCodeTrie(t)

	acc := insertAccountWithCode(t, trie, random, address, genRandomByteArrayOfLen(128))

	codeValue2 := genRandomByteArrayOfLen(128)
	codeValue3 := genRandomByteArrayOfLen(128)
	acc.CodeHash = accounts.InternCodeHash(crypto.Keccak256Hash(codeValue2))

	trie.UpdateAccount(crypto.Keccak256(address[:]), &acc)
	err := trie.UpdateAccountCode(crypto.Keccak256(address[:]), codeValue3)
	require.Error(t, err, "should NOT be able to insert code with wrong hash")
}

func TestCodeNodeUpdateAccountChangeCodeHash(t *testing.T) {
	trie, random, address := newCodeTrie(t)

	acc := insertAccountWithCode(t, trie, random, address, genRandomByteArrayOfLen(128))

	codeValue2 := genRandomByteArrayOfLen(128)
	acc.CodeHash = accounts.InternCodeHash(crypto.Keccak256Hash(codeValue2))

	trie.UpdateAccount(crypto.Keccak256(address[:]), &acc)
	value, gotValue := trie.GetAccountCode(crypto.Keccak256(address[:]))
	assert.Nil(t, value, "the value should reset after the code change happen")
	assert.False(t, gotValue, "should indicate that the code isn't in the cache")
}

func TestCodeNodeUpdateAccountNoChangeCodeHash(t *testing.T) {
	trie, random, address := newCodeTrie(t)

	codeValue1 := genRandomByteArrayOfLen(128)
	acc := insertAccountWithCode(t, trie, random, address, codeValue1)

	acc.Nonce = uint64(random.Int63())
	acc.Balance.SetFromBig(new(big.Int).Rand(random, new(big.Int).Exp(common.Big2, common.Big256, nil)))

	trie.UpdateAccount(crypto.Keccak256(address[:]), &acc)
	value, gotValue := trie.GetAccountCode(crypto.Keccak256(address[:]))
	assert.Equal(t, codeValue1, value, "the value should NOT reset after account's non codehash had changed")
	assert.True(t, gotValue, "should indicate that the code is still in the cache")
}

func TestShortNode(t *testing.T) {
	extensionKeyStr := "0b00000d020809020b080f0d0a090a0a0705070f050d0002090f0d0d0e07080e0402070e010c08030d0e0409060e040b0e0406060b0c080b0503010005"
	extensionKeyNibbles := make([]byte, len(extensionKeyStr)/2+1)
	for i := 0; i < len(extensionKeyStr)/2; i++ {
		c, err := hex.DecodeString(extensionKeyStr[2*i : 2*i+2])
		if err != nil {
			t.Errorf("parsing error")
		}
		extensionKeyNibbles[i] = c[0]
	}
	extensionKeyNibbles[len(extensionKeyNibbles)-1] = 16

	accRlpHex := "f84b08873f54314ab16fd0a056e81f171bcc55a6ff8345e692c0f86e5b48e01b996cadc001622fb5e363b421a0c5d2460186f7233c927e7db2dcc703c0e500b653ca82273b7bfad8045d85a470"
	accRlp, err := hex.DecodeString(accRlpHex)
	if err != nil {
		t.Fatalf("failed to decode rlp from hex")
	}

	valNode := ValueNode(accRlp)

	shortNode := &ShortNode{Key: extensionKeyNibbles, Val: valNode}
	trie := NewInMemoryTrieRLPEncoded(shortNode)
	trie.valueNodesRLPEncoded = true
	rootHash := trie.Root()

	expectedHashStr := "fa7297cfa8b2d445e356e9f752e23b018ae8717fdf117cbe4fcef1b16217d3c2"
	expectedHashBytes, err := hex.DecodeString(expectedHashStr)
	if err != nil {
		t.Fatalf("unable to decode expected hash %v", err)
	}
	if !bytes.Equal(rootHash, expectedHashBytes) {
		t.Fatalf("rootHash(%x) != expectedHash(%x)", rootHash, expectedHashBytes)
	}
}

func TestRLPEncodeDecodeWithAccountsAndStorage(t *testing.T) {
	stateTrie := newEmpty()

	// test addressed
	addresses := []common.Address{
		common.HexToAddress("0x1111111111111111111111111111111111111111"), // EOA
		common.HexToAddress("0x2222222222222222222222222222222222222222"), // Contract with storage
		common.HexToAddress("0x3333333333333333333333333333333333333333"), // EOA
		common.HexToAddress("0xdeadbeefdeadbeefdeadbeefdeadbeefdeadbeef"), // Contract with storage
	}

	// Hash the addresses for trie keys
	addrHashes := make([]common.Hash, len(addresses))
	for i, addr := range addresses {
		addrHashes[i] = crypto.Keccak256Hash(addr[:])
	}

	// Create accounts
	testAccounts := []*accounts.Account{
		testAccount(1, *uint256.NewInt(1000000000000000000)), // 1 ETH
		testAccount(42, *uint256.NewInt(10000000000000000000), // 10 ETH
			withCodeHash(crypto.Keccak256Hash([]byte{0x60, 0x80, 0x60, 0x40}))),
		testAccount(0, *uint256.NewInt(0)),
		testAccount(999, *uint256.NewInt(0xffffffffffffffff),
			withCodeHash(crypto.Keccak256Hash([]byte("contract code")))),
	}

	// Insert accounts using UpdateAccount (creates AccountNode entries)
	for i, addr := range addresses {
		key := crypto.Keccak256(addr[:])
		stateTrie.UpdateAccount(key, testAccounts[i])
	}

	// Define storage for contract at index 1
	contract1AddrHash := addrHashes[1]
	v1 := common.HexToHash("0x1")
	v2 := common.HexToHash("0xdeadbeef")
	v3 := common.HexToHash("0x1234567890abcdef")
	v4 := common.HexToHash("0xff")
	storageSlots1 := []struct {
		slot  common.Hash
		value []byte
	}{
		{common.HexToHash("0x0"), v1[:]},
		{common.HexToHash("0x1"), v2[:]},
		{common.HexToHash("0x2"), v3[:]},
		{common.HexToHash("0x100"), v4[:]},
	}

	// Insert storage using composite keys: addressHash + keccak256(slot)
	// This inserts into AccountNode.Storage
	for _, slot := range storageSlots1 {
		compositeKey := make([]byte, 64)
		copy(compositeKey[:32], contract1AddrHash[:])
		copy(compositeKey[32:], crypto.Keccak256(slot.slot[:]))
		stateTrie.Update(compositeKey, slot.value)
	}

	// Define storage for contract at index 3
	contract2AddrHash := addrHashes[3]
	v5 := common.HexToHash("0xabcd")
	v6 := common.HexToHash("0x9999")
	storageSlots2 := []struct {
		slot  common.Hash
		value []byte
	}{
		{common.HexToHash("0x0"), v5[:]},
		{common.HexToHash("0x5"), v6[:]},
	}

	for _, slot := range storageSlots2 {
		compositeKey := make([]byte, 64)
		copy(compositeKey[:32], contract2AddrHash[:])
		copy(compositeKey[32:], crypto.Keccak256(slot.slot[:]))
		stateTrie.Update(compositeKey, slot.value)
	}

	// Get the storage root hashes via DeepHash
	_, storageRoot1 := stateTrie.DeepHash(contract1AddrHash[:])
	_, storageRoot2 := stateTrie.DeepHash(contract2AddrHash[:])

	// Update expected accounts with computed storage roots
	// (storage was added via Update, so the trie's AccountNode.Root is updated)
	testAccounts[1].Root = storageRoot1
	testAccounts[3].Root = storageRoot2

	// Compute original state root BEFORE encoding
	originalStateRoot := stateTrie.Hash()

	// Encode the unified trie (includes accounts AND their storage subtries)
	encoded, err := stateTrie.RLPEncode()
	require.NoError(t, err)
	require.NotEmpty(t, encoded)

	// Decode the unified trie back
	decodedStateTrie, err := RLPDecode(encoded)
	require.NoError(t, err)

	// Verify the decoded trie hash matches the original
	decodedStateRoot := decodedStateTrie.Hash()
	require.Equal(t, originalStateRoot, decodedStateRoot, "decoded state trie hash should match original")

	// Verify that encoding captured all expected nodes:
	// - Account trie nodes
	// - Storage subtrie nodes for both contracts
	require.GreaterOrEqual(t, len(encoded), 10, "should have multiple nodes for accounts + storage")

	for i, addr := range addresses {
		key := crypto.Keccak256(addr[:])
		acc, ok := decodedStateTrie.GetAccount(key)
		require.True(t, ok)
		require.EqualValues(t, testAccounts[i], acc)
	}

}

// The witness prune paths (WitnessNodesForKeys and WitnessNodesForKeysFromNodes, the
// serializer debug_executionWitness uses) must omit inlined (<32B) nodes and stay
// equivalent. A node whose RLP is <32 bytes is inlined into its parent and never
// referenced by hash, so it must not appear as a standalone witness node. This pins a
// storage trie whose branch inlines two 31-byte leaves (55 remaining nibbles + 1-byte
// value each): the prune must emit the branch but not the two leaves.
func TestWitnessPruneOmitsInlinedNodes(t *testing.T) {
	stateTrie := newEmpty()

	contractAddr := common.HexToAddress("0x2222222222222222222222222222222222222222")
	contract := &accounts.Account{
		Nonce:    7,
		Balance:  *uint256.NewInt(5),
		Root:     EmptyRoot,
		CodeHash: accounts.InternCodeHash(crypto.Keccak256Hash([]byte{0x60, 0x00})),
	}
	stateTrie.UpdateAccount(crypto.Keccak256(contractAddr[:]), contract)

	contractHash := crypto.Keccak256Hash(contractAddr[:])
	slotA := make([]byte, 32)
	slotB := make([]byte, 32)
	slotA[4] = 0x10
	slotB[4] = 0x20
	for i := 5; i < 32; i++ {
		slotA[i] = byte(i)
		slotB[i] = byte(0xff - i)
	}
	keyA := append(append([]byte{}, contractHash[:]...), slotA...)
	keyB := append(append([]byte{}, contractHash[:]...), slotB...)
	stateTrie.Update(keyA, []byte{0x01})
	stateTrie.Update(keyB, []byte{0x02})
	_ = stateTrie.Hash()

	toNibbles := func(b []byte) []byte {
		out := make([]byte, 0, len(b)*2)
		for _, x := range b {
			out = append(out, x>>4, x&0x0f)
		}
		return out
	}
	provedKeys := [][]byte{toNibbles(keyA), toNibbles(keyB)}

	full, err := stateTrie.RLPEncode()
	require.NoError(t, err)

	// Production always prunes a trie decoded from the captured node bytes (branches
	// are FullNodes there), so walk RLPDecode(full), not the Update-built trie.
	wt, err := RLPDecode(full)
	require.NoError(t, err)
	fromTrie, err := wt.WitnessNodesForKeys(provedKeys)
	require.NoError(t, err)
	require.NotEmpty(t, fromTrie)
	fromNodes, err := WitnessNodesForKeysFromNodes(full, provedKeys)
	require.NoError(t, err)

	for i, n := range fromTrie {
		require.GreaterOrEqualf(t, len(n), 32, "WitnessNodesForKeys node %d len %d (<32B)", i, len(n))
	}
	for i, n := range fromNodes {
		require.GreaterOrEqualf(t, len(n), 32, "WitnessNodesForKeysFromNodes node %d len %d (<32B)", i, len(n))
	}

	set := func(ns [][]byte) map[string]struct{} {
		m := make(map[string]struct{}, len(ns))
		for _, n := range ns {
			m[string(n)] = struct{}{}
		}
		return m
	}
	require.Equal(t, set(fromTrie), set(fromNodes), "byHash and RLPDecode prune must agree")
}

// The trie root is force-hashed (Trie.Hash) and referenced by that hash even when its
// own RLP is <32B, so both prune paths must keep it despite the inline-node filter.
func TestWitnessPruneKeepsSmallRoot(t *testing.T) {
	leaf := &ShortNode{Key: []byte{1, 2, 3, 4, 5, 16}, Val: ValueNode([]byte{0x2a})}

	hasher := newHasher(false)
	rootRLP, err := hasher.hashChildren(leaf, 0)
	returnHasherToPool(hasher)
	require.NoError(t, err)
	require.Less(t, len(rootRLP), 32, "test needs a <32B root")

	provedKeys := [][]byte{{1, 2, 3, 4, 5}}

	tr := New(common.Hash{})
	tr.RootNode = leaf
	fromTrie, err := tr.WitnessNodesForKeys(provedKeys)
	require.NoError(t, err)
	require.Len(t, fromTrie, 1)
	require.Equal(t, rootRLP, []byte(fromTrie[0]))

	fromNodes, err := WitnessNodesForKeysFromNodes([][]byte{rootRLP}, provedKeys)
	require.NoError(t, err)
	require.Len(t, fromNodes, 1)
	require.Equal(t, rootRLP, []byte(fromNodes[0]))
}
