go-ethereum/trie/bintrie/stem_node.go

// Copyright 2025 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 (
	"math/bits"

	"github.com/ethereum/go-ethereum/common"
)

// StemNode holds up to 256 values sharing a 31-byte stem, packed via bitmap.
//
// Invariant: dirty=false implies mustRecompute=false. Every mutation that
// invalidates the cached hash MUST also mark the blob for re-flush.
type StemNode struct {
	Stem      [StemSize]byte
	bitmap    [StemBitmapSize]byte
	valueData []byte
	count     uint16
	depth     uint8
	shared    bool // true if valueData is shared with serialized input

	mustRecompute bool        // hash is stale (cleared by Hash)
	dirty         bool        // on-disk blob is stale (cleared by CollectNodes)
	hash          common.Hash // cached hash when mustRecompute == false
}

// posInData returns the offset within valueData, or -1 if absent.
func (sn *StemNode) posInData(suffix byte) int {
	idx := int(suffix)
	if sn.bitmap[idx/8]>>(7-(idx%8))&1 == 0 {
		return -1
	}
	// Count the bits set before this position to determine the offset
	pos := 0
	byteIdx := idx / 8
	for i := 0; i < byteIdx; i++ {
		pos += bits.OnesCount8(sn.bitmap[i])
	}
	// Count bits in the partial byte
	mask := byte(0xFF) << (8 - (idx % 8))
	pos += bits.OnesCount8(sn.bitmap[byteIdx] & mask)
	return pos
}

func (sn *StemNode) getValue(suffix byte) []byte {
	pos := sn.posInData(suffix)
	if pos < 0 {
		return nil
	}
	start := pos * HashSize
	return sn.valueData[start : start+HashSize]
}

func (sn *StemNode) hasValue(suffix byte) bool {
	idx := int(suffix)
	return sn.bitmap[idx/8]>>(7-(idx%8))&1 == 1
}

// allValues returns all 256 values (nil for absent positions).
func (sn *StemNode) allValues() [][]byte {
	values := make([][]byte, StemNodeWidth)
	dataIdx := 0
	for i := range StemNodeWidth {
		if sn.bitmap[i/8]>>(7-(i%8))&1 == 1 {
			values[i] = sn.valueData[dataIdx*HashSize : (dataIdx+1)*HashSize]
			dataIdx++
		}
	}
	return values
}

// ensureWritable copies valueData if shared (copy-on-write).
func (sn *StemNode) ensureWritable() {
	if sn.shared || cap(sn.valueData)-len(sn.valueData) < HashSize {
		newData := make([]byte, len(sn.valueData), len(sn.valueData)+HashSize*4)
		copy(newData, sn.valueData)
		sn.valueData = newData
		sn.shared = false
	}
}

func (sn *StemNode) setValue(suffix byte, value []byte) {
	sn.ensureWritable()
	idx := int(suffix)
	pos := sn.posInData(suffix)
	if pos >= 0 {
		copy(sn.valueData[pos*HashSize:], value[:HashSize])
		return
	}
	sn.bitmap[idx/8] |= 1 << (7 - (idx % 8))
	sn.count++

	insertPos := 0
	byteIdx := idx / 8
	for i := 0; i < byteIdx; i++ {
		insertPos += bits.OnesCount8(sn.bitmap[i])
	}
	mask := byte(0xFF) << (8 - (idx % 8))
	insertPos += bits.OnesCount8(sn.bitmap[byteIdx] & mask)

	insertOffset := insertPos * HashSize
	sn.valueData = append(sn.valueData, make([]byte, HashSize)...)
	copy(sn.valueData[insertOffset+HashSize:], sn.valueData[insertOffset:len(sn.valueData)-HashSize])
	copy(sn.valueData[insertOffset:], value[:HashSize])
}

func (sn *StemNode) Hash() common.Hash {
	if !sn.mustRecompute {
		return sn.hash
	}

	var data [StemNodeWidth]common.Hash
	h := newSha256()
	defer returnSha256(h)

	// Hash each present value
	dataIdx := 0
	for i := range StemNodeWidth {
		if sn.bitmap[i/8]>>(7-(i%8))&1 == 1 {
			v := sn.valueData[dataIdx*HashSize : (dataIdx+1)*HashSize]
			h.Reset()
			h.Write(v)
			h.Sum(data[i][:0])
			dataIdx++
		}
	}
	h.Reset()

	for level := 1; level <= 8; level++ {
		for i := range StemNodeWidth / (1 << level) {
			h.Reset()

			if data[i*2] == (common.Hash{}) && data[i*2+1] == (common.Hash{}) {
				data[i] = common.Hash{}
				continue
			}

			h.Write(data[i*2][:])
			h.Write(data[i*2+1][:])
			data[i] = common.Hash(h.Sum(nil))
		}
	}

	h.Reset()
	h.Write(sn.Stem[:])
	h.Write([]byte{0})
	h.Write(data[0][:])
	sn.hash = common.BytesToHash(h.Sum(nil))
	sn.mustRecompute = false
	return sn.hash
}

func (sn *StemNode) Key(i int) []byte {
	var ret [HashSize]byte
	copy(ret[:], sn.Stem[:])
	ret[StemSize] = byte(i)
	return ret[:]
}