node_test.go

package jamtrie_test

import (
	"testing"

	"github.com/shunsukew/jamtrie"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestBranchEncodeDecodeRoundtrip(t *testing.T) {
	left := [32]byte{255, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32}
	right := [32]byte{32, 31, 30, 29, 28, 27, 26, 25, 24, 23, 22, 21, 20, 19, 18, 17, 16, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1}

	node := jamtrie.NewBranch(left, right)
	encoded := node.Encode()

	assert.Equal(t, jamtrie.NodeSize, len(encoded))

	decoded, err := jamtrie.DecodeNode(encoded)
	require.NoError(t, err)

	branch, ok := decoded.(*jamtrie.Branch)
	require.True(t, ok, "decoded node should be a branch")

	// When decodes once encoded Branch node data, the left child's hash first bit is cleared.
	expectedLeftHash := left
	expectedLeftHash[0] &= 0b01111111
	assert.Equal(t, expectedLeftHash, branch.LeftHash())
	assert.Equal(t, right, branch.RightHash())
}

func TestEmbeddedValueLeafEncodeDecodeRoundtrip(t *testing.T) {
	tests := []struct {
		name      string
		stateKey  [32]byte
		value     []byte
		expectErr bool
	}{
		{
			name:     "max size value",
			stateKey: [32]byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32},
			value:    make([]byte, jamtrie.MaxEmbeddedValueSize),
		},
		{
			name:     "mid size value",
			stateKey: [32]byte{16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47},
			value:    []byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20},
		},
		{
			name:     "small value",
			stateKey: [32]byte{32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63},
			value:    []byte{1, 2, 3, 4, 5},
		},
		{
			name:     "empty value",
			stateKey: [32]byte{63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94},
			value:    []byte{},
		},
		{
			name:      "value too large",
			stateKey:  [32]byte{},
			value:     make([]byte, jamtrie.MaxEmbeddedValueSize+1),
			expectErr: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			node, err := jamtrie.NewEmbeddedValueLeaf(tt.stateKey, tt.value)
			if tt.expectErr {
				assert.Error(t, err)
				return
			}
			require.NoError(t, err)

			encoded := node.Encode()
			assert.Equal(t, jamtrie.NodeSize, len(encoded))

			decoded, err := jamtrie.DecodeNode(encoded)
			require.NoError(t, err)

			leaf, ok := decoded.(*jamtrie.EmbeddedValueLeaf)
			require.True(t, ok, "decoded node should be an embedded value leaf")

			// when encoding, last 1 byte of stateKey is lost, so we compare only first 31 bytes
			expected := [32]byte{}
			copy(expected[:31], tt.stateKey[:31])
			assert.Equal(t, expected, leaf.StateKey())
			assert.Equal(t, tt.value, leaf.Value())
		})
	}
}

func TestRegularLeafEncodeDecodeRoundtrip(t *testing.T) {
	tests := []struct {
		name      string
		stateKey  [32]byte
		valueHash [32]byte
	}{
		{
			name:     "state key is 31 bytes",
			stateKey: [32]byte{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31},
			valueHash: [32]byte{
				100, 101, 102, 103, 104, 105, 106, 107,
				108, 109, 110, 111, 112, 113, 114, 115,
				116, 117, 118, 119, 120, 121, 122, 123, 124,
				125, 126, 127, 128, 129, 130, 131,
			},
		},
		{
			name: "state key is more than 31 bytes",
			stateKey: [32]byte{
				0, 1, 2, 3, 4, 5, 6, 7, 8, 9,
				10, 11, 12, 13, 14, 15, 16, 17, 18, 19,
				20, 21, 22, 23, 24, 25, 26, 27, 28, 29,
				30, 31,
			},
			valueHash: [32]byte{
				200, 201, 202, 203, 204, 205, 206, 207,
				208, 209, 210, 211, 212, 213, 214, 215,
				216, 217, 218, 219, 220, 221, 222, 223, 224,
				225, 226, 227, 228, 229, 230, 231,
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			node := jamtrie.NewRegularLeaf(tt.stateKey, tt.valueHash)
			encoded := node.Encode()

			assert.Equal(t, jamtrie.NodeSize, len(encoded))

			decoded, err := jamtrie.DecodeNode(encoded)
			require.NoError(t, err)

			leaf, ok := decoded.(*jamtrie.RegularLeaf)
			require.True(t, ok, "decoded node should be a regular leaf")

			// when encoding, last 1 byte of stateKey is lost, so we compare only first 31 bytes
			expected := [32]byte{}
			copy(expected[:31], tt.stateKey[:31])
			assert.Equal(t, expected, leaf.StateKey())
			assert.Equal(t, tt.valueHash[:], leaf.Value())
		})
	}
}

func TestDecodeNodeInvalidSize(t *testing.T) {
	invalidData := make([]byte, jamtrie.NodeSize-1)
	_, err := jamtrie.DecodeNode(invalidData)
	assert.ErrorIs(t, err, jamtrie.ErrInvalidNodeSize)

	invalidData = make([]byte, jamtrie.NodeSize+1)
	_, err = jamtrie.DecodeNode(invalidData)
	assert.ErrorIs(t, err, jamtrie.ErrInvalidNodeSize)
}