node.go

package jamtrie

import (
	"errors"

	"golang.org/x/crypto/blake2b"
)

const (
	NodeSize = 64 // fixed size in bytes

	// Maximum size for embedded values (small values stored directly in node)
	MaxEmbeddedValueSize = 32
)

var (
	ErrInvalidValueSize   = errors.New("value too large to embed in leaf node")
	ErrInvalidNodeSize    = errors.New("invalid node size")
	ErrInvalidNodeType    = errors.New("invalid node type")
	ErrDecodeNotSupported = errors.New("decoding not supported for this node type")
)

type Node interface {
	Hash() []byte
	Encode() []byte
	Decode(data []byte) error
}

type Leaf interface {
	Node
	StateKey() [32]byte
	Value() []byte
}

func DecodeNode(data []byte) (Node, error) {
	if len(data) != NodeSize {
		return nil, ErrInvalidNodeSize
	}

	// Determine node type by prefix.
	// First bit 0 indicates Branch node.
	if (data[0] & 0b10000000) == 0 {
		node := &Branch{}
		err := node.Decode(data)
		if err != nil {
			return nil, err
		}
		return node, nil
	}

	switch (data[0] >> 6) & 0b11 {
	case 0b10:
		// Embedded-value leaf node (first two bits 10)
		node := &EmbeddedValueLeaf{}
		err := node.Decode(data)
		if err != nil {
			return nil, err
		}
		return node, nil
	case 0b11:
		// Regular leaf node (first two bits 11)
		node := &RegularLeaf{}
		err := node.Decode(data)
		if err != nil {
			return nil, err
		}
		return node, nil
	}

	return nil, ErrInvalidNodeType
}

// Branch node represents an internal node with left and right children
type Branch struct {
	left  [32]byte // Hash of left child
	right [32]byte // Hash of right child
}

func NewBranch(left, right [32]byte) Node {
	return &Branch{
		left:  left,
		right: right,
	}
}

func (b *Branch) Hash() []byte {
	hasher, _ := blake2b.New256(nil)
	hasher.Write(b.Encode())
	hash := hasher.Sum(nil)
	return hash
}

func (b *Branch) Encode() []byte {
	encoded := make([]byte, NodeSize)
	encoded[0] = b.left[0] & 0b01111111 // Clear most significant bit (MSB)
	copy(encoded[1:32], b.left[1:32])
	copy(encoded[32:64], b.right[0:32])
	return encoded
}

// Decode decodes a Branch node from given byte slice.
// Note that the left child's hash first bit was cleared during encoding,
// Thus, Decode cannot ensure the exact the same Branch recovery, but this has no problem for trie node operations.
func (b *Branch) Decode(data []byte) error {
	if len(data) != NodeSize {
		return ErrInvalidNodeSize
	}

	// Branch must have MSB=0
	if data[0]&0b10000000 != 0 {
		return ErrInvalidNodeType
	}

	// Reconstruct the left hash **with MSB cleared** (as per encoding)
	// MSB of left[0] is already ensured to be 0 by the check above
	copy(b.left[:], data[0:32])
	copy(b.right[:], data[32:64])

	return nil
}

func (b *Branch) LeftHash() [32]byte {
	return b.left
}

func (b *Branch) RightHash() [32]byte {
	return b.right
}

// Embedded-value leaf node stores small values directly in the node
type EmbeddedValueLeaf struct {
	stateKey [32]byte // State key 32 bytes as defined in JAM spec v0.4.0
	value    []byte   // Small value stored directly (max 32 bytes)
}

func NewEmbeddedValueLeaf(stateKey [32]byte, value []byte) (Leaf, error) {
	if len(value) > MaxEmbeddedValueSize {
		return nil, ErrInvalidValueSize
	}
	return &EmbeddedValueLeaf{
		stateKey: stateKey,
		value:    append([]byte(nil), value...), // Copy value to avoid external mutation
	}, nil
}

func (e *EmbeddedValueLeaf) Hash() []byte {
	hasher, _ := blake2b.New256(nil)
	hasher.Write(e.Encode())
	hash := hasher.Sum(nil)
	return hash
}

func (e *EmbeddedValueLeaf) Encode() []byte {
	valueSize := len(e.value)
	encodedValueSize := EncodeFixedLength(uint32(valueSize), 1)

	encoded := make([]byte, NodeSize)
	// Set MSB=1 and value size in byte. Because value size is <=32, first 2 bits of encodedValueSize[0] is always 0.
	// Thus, first 2 bits is ensured to be 0b10.
	encoded[0] = 0b10000000 | encodedValueSize[0]
	// Copy state key (31 bytes), trim the first 31 bytes as per JAM spec v0.4.0
	// Lost 1 byte information, this information loss is fixed at least after JAM spec v0.7.0
	copy(encoded[1:32], e.stateKey[:])
	// Copy value (up to 32 bytes)
	copy(encoded[32:32+valueSize], e.value)
	// Remaining bytes are already zero-initialized (padding)

	return encoded
}

func (e *EmbeddedValueLeaf) Decode(data []byte) error {
	if len(data) != NodeSize {
		return ErrInvalidNodeSize
	}

	// Embedded leaf must have MSBs == 0b10
	if (data[0] & 0b11000000) != 0b10000000 {
		return ErrInvalidNodeType
	}

	// Extract value length from the first byte, lower 7 bits holds value size data
	valueSize := uint(data[0] & 0b01111111)
	if valueSize > 32 {
		return ErrInvalidValueSize
	}

	copy(e.stateKey[:], data[1:32])

	e.value = make([]byte, valueSize)
	copy(e.value[:], data[32:32+valueSize])

	return nil
}

func (e *EmbeddedValueLeaf) StateKey() [32]byte {
	return e.stateKey
}

func (e *EmbeddedValueLeaf) Value() []byte {
	return e.value
}

// Regular leaf node stores a key-value pair where value is stored externally
type RegularLeaf struct {
	stateKey  [32]byte // State key (first 31 bytes as trimed per JAM spec v0.4.0)
	valueHash [32]byte // Hash of the value stored externally
}

func NewRegularLeaf(stateKey [32]byte, valueHash [32]byte) Leaf {
	return &RegularLeaf{
		stateKey:  stateKey,
		valueHash: valueHash,
	}
}

func (r *RegularLeaf) Hash() []byte {
	hasher, _ := blake2b.New256(nil)
	hasher.Write(r.Encode())
	hash := hasher.Sum(nil)
	return hash
}

func (r *RegularLeaf) Encode() []byte {
	encoded := make([]byte, NodeSize)

	// The first 1 byte is always 0b11000000 (2 bits prefix + 6 bits padding) for RegularLeaf
	encoded[0] = 0b11000000
	// Copy state key (31 bytes), trim the first 31 bytes as per JAM spec v0.4.0
	// Lost 1 byte information, this information loss is fixed at least after JAM spec v0.7.0
	copy(encoded[1:32], r.stateKey[:])
	// Copy value (up to 32 bytes)
	copy(encoded[32:64], r.valueHash[:])

	return encoded
}

func (r *RegularLeaf) Decode(data []byte) error {
	if len(data) != NodeSize {
		return ErrInvalidNodeSize
	}

	// Regular leaf must have 0x11000000 as the first byte
	if data[0] != 0b11000000 {
		return ErrInvalidNodeType
	}

	copy(r.stateKey[:], data[1:32])
	copy(r.valueHash[:], data[32:64])

	return nil
}

func (r *RegularLeaf) StateKey() [32]byte {
	return r.stateKey
}

func (r *RegularLeaf) Value() []byte {
	return r.valueHash[:]
}