niko_trust/internal/smt/smt.go

// Package smt implements a compressed sparse Merkle trie over fixed-size
// keys, committed to by a single root hash.
//
// Purpose. Two relays that store the same set of objects must compute the
// same root regardless of insertion order: the root is a function of the
// set, so comparing one 32-byte string detects divergence between mirrors
// instantly (docs/CHECKPOINT.md). The trie is grow-only — objects are never
// removed from the log even when revoked — which keeps every operation
// monotone and the proofs simple.
//
// Hash domain (SHA-256 throughout):
//
//	leaf   = H(0x00 || key)
//	branch = H(0x01 || uvarint(prefixLen) || prefixBits || left || right)
//	empty  = H(0x02)
//
// The prefix is inside the branch hash so the trie's shape is committed to:
// without it, proofs could not be checked for canonicality. A branch exists
// only where both children exist and splits exactly at its keys' first
// differing bit, which makes the shape a function of the set alone.
//
// Proofs. An inclusion proof walks from the key's leaf to the root through
// sibling hashes. An absence proof witnesses the exact spot where the key
// would have lived: either a neighbouring leaf whose path agrees with the
// queried key up to that leaf's depth (variant A), or the branch whose
// prefix the queried key leaves mid-way (variant B). Both recombine with
// the queried key's own bits, so a proof for one key verifies for no other.
package smt

import (
	"crypto/sha256"
	"encoding/binary"
	"errors"
)

// KeySize is the exact key size: an object content ID.
const KeySize = 32

// EmptyRoot is the root of a trie containing nothing.
var EmptyRoot = sha256.Sum256([]byte{0x02})

// Errors reported by proof generation and verification.
var (
	ErrPresent  = errors.New("smt: key is present")
	ErrAbsent   = errors.New("smt: key is absent")
	ErrBadProof = errors.New("smt: bad proof")
)

type node struct {
	leaf bool

	// Leaf fields.
	key [KeySize]byte

	// Branch fields: prefix holds the bits between the parent's split and
	// this node's own split, packed MSB-first, prefixLen valid bits. A
	// branch always has both children.
	prefix    []byte
	prefixLen int
	left      *node
	right     *node
}

// Trie is a grow-only set commitment. Not safe for concurrent use; callers
// serialise access.
type Trie struct {
	root *node
	n    int
}

// New returns an empty trie.
func New() *Trie { return &Trie{} }

// Len returns the number of distinct keys inserted.
func (t *Trie) Len() int { return t.n }

// ---------------------------------------------------------------- primitives

func bit(k *[KeySize]byte, i int) byte {
	return k[i>>3] >> (7 - uint(i&7)) & 1
}

func prefixBit(p []byte, i int) byte {
	return p[i>>3] >> (7 - uint(i&7)) & 1
}

func commonPrefixLen(a, b *[KeySize]byte) int {
	for i := 0; i < KeySize*8; i++ {
		if bit(a, i) != bit(b, i) {
			return i
		}
	}
	return KeySize * 8
}

func packBits(k *[KeySize]byte, from, n int) []byte {
	out := make([]byte, (n+7)/8)
	for i := 0; i < n; i++ {
		if bit(k, from+i) == 1 {
			out[i>>3] |= 0x80 >> uint(i&7)
		}
	}
	return out
}

func appendUvarint(b []byte, v int) []byte {
	var tmp [10]byte
	n := binary.PutUvarint(tmp[:], uint64(v))
	return append(b, tmp[:n]...)
}

func leafHash(key *[KeySize]byte) [32]byte {
	h := sha256.New()
	h.Write([]byte{0x00})
	h.Write(key[:])
	var out [32]byte
	copy(out[:], h.Sum(nil))
	return out
}

func branchHash(prefix []byte, prefixLen int, left, right *[32]byte) [32]byte {
	h := sha256.New()
	h.Write([]byte{0x01})
	var tmp [10]byte
	n := binary.PutUvarint(tmp[:], uint64(prefixLen))
	h.Write(tmp[:n])
	h.Write(prefix)
	h.Write(left[:])
	h.Write(right[:])
	var out [32]byte
	copy(out[:], h.Sum(nil))
	return out
}

func (t *Trie) hash(n *node) [32]byte {
	if n.leaf {
		return leafHash(&n.key)
	}
	l := t.hash(n.left)
	r := t.hash(n.right)
	return branchHash(n.prefix, n.prefixLen, &l, &r)
}

// Root returns the set commitment.
func (t *Trie) Root() [32]byte {
	if t.root == nil {
		return EmptyRoot
	}
	return t.hash(t.root)
}

// ------------------------------------------------------------------- insert

// Insert adds key, reporting whether it was new.
func (t *Trie) Insert(key [KeySize]byte) bool {
	root, added := insertNode(t.root, &key, 0)
	t.root = root
	if added {
		t.n++
	}
	return added
}

// Contains reports whether key has been inserted.
func (t *Trie) Contains(key [KeySize]byte) bool {
	n := t.root
	d := 0
	for n != nil && !n.leaf {
		i := 0
		for i < n.prefixLen && bit(&key, d+i) == prefixBit(n.prefix, i) {
			i++
		}
		if i < n.prefixLen {
			return false // diverged inside the prefix
		}
		d += n.prefixLen
		if bit(&key, d) == 0 {
			n = n.left
		} else {
			n = n.right
		}
		d++
	}
	return n != nil && n.key == key
}

// insertNode inserts key into the subtree rooted at n, whose split decision
// begins at depth d. It returns the possibly-replaced subtree.
func insertNode(n *node, key *[KeySize]byte, d int) (*node, bool) {
	if n == nil {
		return &node{leaf: true, key: *key}, true
	}
	if n.leaf {
		if n.key == *key {
			return n, false
		}
		k := commonPrefixLen(&n.key, key)
		// Shared bits d..k-1 become the new branch's prefix; bit k splits.
		newBranch := &node{
			prefix:    packBits(key, d, k-d),
			prefixLen: k - d,
		}
		if bit(key, k) == 0 {
			newBranch.left = &node{leaf: true, key: *key}
			newBranch.right = n
		} else {
			newBranch.left = n
			newBranch.right = &node{leaf: true, key: *key}
		}
		return newBranch, true
	}

	// Branch: follow the prefix while it matches.
	i := 0
	for i < n.prefixLen && bit(key, d+i) == prefixBit(n.prefix, i) {
		i++
	}
	if i < n.prefixLen {
		// The key diverges inside the prefix at absolute bit d+i. Split the
		// branch there: the upper half becomes the new branch's prefix, the
		// old branch keeps the remainder as its own prefix.
		j := d + i
		upperPrefix := make([]byte, (i+7)/8)
		for b := 0; b < i; b++ {
			if prefixBit(n.prefix, b) == 1 {
				upperPrefix[b>>3] |= 0x80 >> uint(b&7)
			}
		}
		upper := &node{
			prefix:    upperPrefix,
			prefixLen: i,
		}
		tail := copyTailBits(n.prefix, i, n.prefixLen)
		lower := &node{
			prefix:    tail,
			prefixLen: n.prefixLen - i - 1,
			left:      n.left,
			right:     n.right,
		}
		leaf := &node{leaf: true, key: *key}
		if bit(key, j) == 0 {
			upper.left = leaf
			upper.right = lower
		} else {
			upper.left = lower
			upper.right = leaf
		}
		return upper, true
	}

	sd := d + n.prefixLen
	var added bool
	if bit(key, sd) == 0 {
		n.left, added = insertNode(n.left, key, sd+1)
	} else {
		n.right, added = insertNode(n.right, key, sd+1)
	}
	return n, added
}

// copyTailBits extracts bits [from+1, end) of p, MSB-packed.
func copyTailBits(p []byte, from, end int) []byte {
	n := end - from - 1
	out := make([]byte, (n+7)/8)
	for i := 0; i < n; i++ {
		if prefixBit(p, from+1+i) == 1 {
			out[i>>3] |= 0x80 >> uint(i&7)
		}
	}
	return out
}