peoplecoin

  1// Package trie implements the binary Merkle trie of the specification.
  2//
  3// Keys are 32 bytes and navigated bit by bit, MSB first (bit 0 is the most
  4// significant bit of byte 0); left child 0, right child 1. Branches exist
  5// only where keys diverge (a crit-bit trie), so the shape of the trie — and
  6// therefore its root hash — is a pure function of the key/value set,
  7// independent of insertion order.
  8//
  9//	Leaf hash:    SHA256(key ‖ value)
 10//	Branch hash:  SHA256(bit_position ‖ aggregate ‖ left_hash ‖ right_hash)
 11//	Empty:        [32]byte{0}
 12//
 13// bit_position is a uint8, the key bit at which the branch diverges. Storage
 14// tries (reg, rng, election) carry the number of entries below each branch as
 15// a big-endian uint64 aggregate, which gives every entry an index — its rank
 16// in key order, leftmost 0 — via Rank and Select. Plain tries (tx, coin,
 17// token) omit the aggregate from the hash.
 18//
 19// The trie is persistent: every mutation returns a new *Trie sharing
 20// unmodified subtrees with the old one, which is left untouched. Values are
 21// treated as immutable byte strings; Insert and Update copy them.
 22package trie
 23
 24import (
 25	"crypto/sha256"
 26	"encoding/binary"
 27	"errors"
 28)
 29
 30// Key is a 32-byte trie key: a public key, or a UTXO or transaction ID.
 31type Key [32]byte
 32
 33// Hash is a 32-byte SHA-256 digest.
 34type Hash [32]byte
 35
 36var (
 37	ErrExists   = errors.New("trie: key already present")
 38	ErrNotFound = errors.New("trie: key not found")
 39)
 40
 41// Trie is a persistent binary Merkle trie.
 42type Trie struct {
 43	root    node
 44	counted bool
 45}
 46
 47// NewStorage returns an empty storage trie (branches carry entry counts).
 48func NewStorage() *Trie { return &Trie{counted: true} }
 49
 50// NewPlain returns an empty plain trie (no aggregate in the branch hash).
 51func NewPlain() *Trie { return &Trie{counted: false} }
 52
 53// Counted reports whether this trie hashes entry counts into its branches.
 54func (t *Trie) Counted() bool { return t.counted }
 55
 56type node interface {
 57	hash() Hash
 58	count() uint64
 59}
 60
 61type leaf struct {
 62	key   Key
 63	value []byte
 64	h     Hash
 65}
 66
 67type branch struct {
 68	pos   uint8 // bit at which the children diverge
 69	n     uint64
 70	left  node
 71	right node
 72	h     Hash
 73}
 74
 75func (l *leaf) hash() Hash      { return l.h }
 76func (l *leaf) count() uint64   { return 1 }
 77func (b *branch) hash() Hash    { return b.h }
 78func (b *branch) count() uint64 { return b.n }
 79
 80// LeafHash is SHA256(key ‖ value).
 81func LeafHash(key Key, value []byte) Hash {
 82	h := sha256.New()
 83	h.Write(key[:])
 84	h.Write(value)
 85	var out Hash
 86	copy(out[:], h.Sum(nil))
 87	return out
 88}
 89
 90// BranchHash is SHA256(bit_position ‖ aggregate ‖ left ‖ right); the
 91// aggregate (big-endian uint64 count) is present only when counted is true.
 92func BranchHash(counted bool, pos uint8, count uint64, left, right Hash) Hash {
 93	h := sha256.New()
 94	h.Write([]byte{pos})
 95	if counted {
 96		var c [8]byte
 97		binary.BigEndian.PutUint64(c[:], count)
 98		h.Write(c[:])
 99	}
100	h.Write(left[:])
101	h.Write(right[:])
102	var out Hash
103	copy(out[:], h.Sum(nil))
104	return out
105}
106
107func newLeaf(key Key, value []byte) *leaf {
108	v := make([]byte, len(value))
109	copy(v, value)
110	return &leaf{key: key, value: v, h: LeafHash(key, v)}
111}
112
113func (t *Trie) newBranch(pos uint8, left, right node) *branch {
114	n := left.count() + right.count()
115	return &branch{
116		pos:   pos,
117		n:     n,
118		left:  left,
119		right: right,
120		h:     BranchHash(t.counted, pos, n, left.hash(), right.hash()),
121	}
122}
123
124// bit returns bit pos of key, MSB first.
125func bit(k *Key, pos uint8) uint8 {
126	return (k[pos>>3] >> (7 - pos&7)) & 1
127}
128
129// firstDiff returns the first bit position at which a and b differ.
130func firstDiff(a, b *Key) (uint8, bool) {
131	for i := 0; i < 32; i++ {
132		if x := a[i] ^ b[i]; x != 0 {
133			p := i * 8
134			for x&0x80 == 0 {
135				x <<= 1
136				p++
137			}
138			return uint8(p), true
139		}
140	}
141	return 0, false
142}
143
144func (t *Trie) with(root node) *Trie { return &Trie{root: root, counted: t.counted} }
145
146// Root returns the root hash; the zero hash for an empty trie.
147func (t *Trie) Root() Hash {
148	if t.root == nil {
149		return Hash{}
150	}
151	return t.root.hash()
152}
153
154// Count returns the number of entries.
155func (t *Trie) Count() uint64 {
156	if t.root == nil {
157		return 0
158	}
159	return t.root.count()
160}
161
162// findLeaf follows key's bits at each branch and returns the leaf reached,
163// which may or may not hold key.
164func (t *Trie) findLeaf(key *Key) *leaf {
165	n := t.root
166	for {
167		switch x := n.(type) {
168		case *leaf:
169			return x
170		case *branch:
171			if bit(key, x.pos) == 0 {
172				n = x.left
173			} else {
174				n = x.right
175			}
176		default:
177			return nil
178		}
179	}
180}
181
182// Get returns the value stored at key.
183func (t *Trie) Get(key Key) ([]byte, bool) {
184	l := t.findLeaf(&key)
185	if l == nil || l.key != key {
186		return nil, false
187	}
188	return l.value, true
189}
190
191// Has reports whether key is present.
192func (t *Trie) Has(key Key) bool {
193	_, ok := t.Get(key)
194	return ok
195}
196
197// Insert adds a new entry. It fails with ErrExists if key is present.
198func (t *Trie) Insert(key Key, value []byte) (*Trie, error) {
199	if t.root == nil {
200		return t.with(newLeaf(key, value)), nil
201	}
202	l := t.findLeaf(&key)
203	d, ok := firstDiff(&key, &l.key)
204	if !ok {
205		return nil, ErrExists
206	}
207	return t.with(t.insertAt(t.root, key, value, d)), nil
208}
209
210// insertAt splices a new leaf into the subtree n at the first node whose
211// branch position exceeds d, the bit at which key diverges from every key in
212// that subtree.
213func (t *Trie) insertAt(n node, key Key, value []byte, d uint8) node {
214	if b, ok := n.(*branch); ok && b.pos < d {
215		if bit(&key, b.pos) == 0 {
216			return t.newBranch(b.pos, t.insertAt(b.left, key, value, d), b.right)
217		}
218		return t.newBranch(b.pos, b.left, t.insertAt(b.right, key, value, d))
219	}
220	nl := newLeaf(key, value)
221	if bit(&key, d) == 0 {
222		return t.newBranch(d, nl, n)
223	}
224	return t.newBranch(d, n, nl)
225}
226
227// modify rebuilds the path to key, replacing its leaf with f(leaf); a nil
228// result deletes the leaf, collapsing the parent into the sibling.
229func (t *Trie) modify(n node, key *Key, f func(*leaf) node) (node, error) {
230	switch x := n.(type) {
231	case *leaf:
232		if x.key != *key {
233			return nil, ErrNotFound
234		}
235		return f(x), nil
236	case *branch:
237		if bit(key, x.pos) == 0 {
238			l, err := t.modify(x.left, key, f)
239			if err != nil {
240				return nil, err
241			}
242			if l == nil {
243				return x.right, nil
244			}
245			return t.newBranch(x.pos, l, x.right), nil
246		}
247		r, err := t.modify(x.right, key, f)
248		if err != nil {
249			return nil, err
250		}
251		if r == nil {
252			return x.left, nil
253		}
254		return t.newBranch(x.pos, x.left, r), nil
255	default:
256		return nil, ErrNotFound
257	}
258}
259
260// Update replaces the value of an existing entry.
261func (t *Trie) Update(key Key, value []byte) (*Trie, error) {
262	root, err := t.modify(t.root, &key, func(*leaf) node { return newLeaf(key, value) })
263	if err != nil {
264		return nil, err
265	}
266	return t.with(root), nil
267}
268
269// Delete removes an existing entry.
270func (t *Trie) Delete(key Key) (*Trie, error) {
271	root, err := t.modify(t.root, &key, func(*leaf) node { return nil })
272	if err != nil {
273		return nil, err
274	}
275	return t.with(root), nil
276}
277
278// Rank returns the index of key in key order (leftmost leaf 0).
279func (t *Trie) Rank(key Key) (uint64, bool) {
280	var idx uint64
281	n := t.root
282	for {
283		switch x := n.(type) {
284		case *leaf:
285			if x.key == key {
286				return idx, true
287			}
288			return 0, false
289		case *branch:
290			if bit(&key, x.pos) == 0 {
291				n = x.left
292			} else {
293				idx += x.left.count()
294				n = x.right
295			}
296		default:
297			return 0, false
298		}
299	}
300}
301
302// Select returns the entry at index i in key order.
303func (t *Trie) Select(i uint64) (Key, []byte, bool) {
304	if t.root == nil || i >= t.root.count() {
305		return Key{}, nil, false
306	}
307	n := t.root
308	for {
309		switch x := n.(type) {
310		case *leaf:
311			return x.key, x.value, true
312		case *branch:
313			if lc := x.left.count(); i < lc {
314				n = x.left
315			} else {
316				i -= lc
317				n = x.right
318			}
319		}
320	}
321}
322
323// Walk visits every entry in key order until fn returns false.
324func (t *Trie) Walk(fn func(Key, []byte) bool) {
325	var walk func(node) bool
326	walk = func(n node) bool {
327		switch x := n.(type) {
328		case *leaf:
329			return fn(x.key, x.value)
330		case *branch:
331			return walk(x.left) && walk(x.right)
332		}
333		return true
334	}
335	if t.root != nil {
336		walk(t.root)
337	}
338}