peoplecoin

  1package trie
  2
  3import (
  4	"bytes"
  5	"crypto/sha256"
  6	"encoding/binary"
  7	"math/rand"
  8	"sort"
  9	"testing"
 10)
 11
 12func k(b ...byte) Key {
 13	var key Key
 14	copy(key[:], b)
 15	return key
 16}
 17
 18func mustInsert(t *testing.T, tr *Trie, key Key, v []byte) *Trie {
 19	t.Helper()
 20	n, err := tr.Insert(key, v)
 21	if err != nil {
 22		t.Fatalf("insert %x: %v", key[:2], err)
 23	}
 24	return n
 25}
 26
 27func TestEmptyRoot(t *testing.T) {
 28	if NewStorage().Root() != (Hash{}) || NewPlain().Root() != (Hash{}) {
 29		t.Fatal("empty root must be the zero hash")
 30	}
 31	if NewStorage().Count() != 0 {
 32		t.Fatal("empty count")
 33	}
 34}
 35
 36func TestSingleLeafRootIsLeafHash(t *testing.T) {
 37	key, val := k(0xAB), []byte("value")
 38	tr := mustInsert(t, NewStorage(), key, val)
 39	want := sha256.Sum256(append(append([]byte{}, key[:]...), val...))
 40	if tr.Root() != Hash(want) {
 41		t.Fatal("single leaf root must equal SHA256(key ‖ value)")
 42	}
 43}
 44
 45// Two keys diverging at bit 0: the root is a branch at position 0 with the
 46// smaller key on the left. Checked by hand against the hash formula, with
 47// and without the count aggregate.
 48func TestTwoLeafBranchHash(t *testing.T) {
 49	a, b := k(0x00, 0x01), k(0x80, 0x02)
 50	va, vb := []byte("a"), []byte("b")
 51	la, lb := LeafHash(a, va), LeafHash(b, vb)
 52
 53	for _, counted := range []bool{true, false} {
 54		tr := NewPlain()
 55		if counted {
 56			tr = NewStorage()
 57		}
 58		tr = mustInsert(t, mustInsert(t, tr, b, vb), a, va)
 59
 60		h := sha256.New()
 61		h.Write([]byte{0}) // bit_position
 62		if counted {
 63			var c [8]byte
 64			binary.BigEndian.PutUint64(c[:], 2)
 65			h.Write(c[:])
 66		}
 67		h.Write(la[:])
 68		h.Write(lb[:])
 69		var want Hash
 70		copy(want[:], h.Sum(nil))
 71		if tr.Root() != want {
 72			t.Fatalf("counted=%v: root mismatch", counted)
 73		}
 74	}
 75	// The two flavours must not collide.
 76	if mustInsert(t, mustInsert(t, NewStorage(), a, va), b, vb).Root() ==
 77		mustInsert(t, mustInsert(t, NewPlain(), a, va), b, vb).Root() {
 78		t.Fatal("storage and plain roots should differ")
 79	}
 80}
 81
 82// Divergence deeper than bit 0: keys sharing the first 13 bits.
 83func TestDeepDivergencePosition(t *testing.T) {
 84	a := k(0xFF, 0xF8, 0x00) // 1111 1111 1111 1000
 85	b := k(0xFF, 0xFC, 0x00) // 1111 1111 1111 1100 — differs at bit 13
 86	tr := mustInsert(t, mustInsert(t, NewStorage(), a, []byte{1}), b, []byte{2})
 87	br, ok := tr.root.(*branch)
 88	if !ok || br.pos != 13 {
 89		t.Fatalf("expected branch at bit 13, got %+v", tr.root)
 90	}
 91	if br.left.(*leaf).key != a || br.right.(*leaf).key != b {
 92		t.Fatal("bit 0 goes left, bit 1 goes right")
 93	}
 94}
 95
 96func TestErrors(t *testing.T) {
 97	tr := mustInsert(t, NewStorage(), k(1), []byte{1})
 98	if _, err := tr.Insert(k(1), []byte{9}); err != ErrExists {
 99		t.Fatalf("duplicate insert: %v", err)
100	}
101	if _, err := tr.Update(k(2), nil); err != ErrNotFound {
102		t.Fatalf("update missing: %v", err)
103	}
104	if _, err := tr.Delete(k(2)); err != ErrNotFound {
105		t.Fatalf("delete missing: %v", err)
106	}
107	if _, err := NewStorage().Delete(k(2)); err != ErrNotFound {
108		t.Fatalf("delete from empty: %v", err)
109	}
110}
111
112func TestPersistence(t *testing.T) {
113	t1 := mustInsert(t, NewStorage(), k(1), []byte{1})
114	r1 := t1.Root()
115	t2 := mustInsert(t, t1, k(2), []byte{2})
116	t3, err := t2.Update(k(1), []byte{7})
117	if err != nil {
118		t.Fatal(err)
119	}
120	t4, err := t3.Delete(k(2))
121	if err != nil {
122		t.Fatal(err)
123	}
124	if t1.Root() != r1 || t1.Count() != 1 || t1.Has(k(2)) {
125		t.Fatal("older version was mutated")
126	}
127	if v, _ := t2.Get(k(1)); !bytes.Equal(v, []byte{1}) {
128		t.Fatal("update leaked into older version")
129	}
130	if v, _ := t4.Get(k(1)); !bytes.Equal(v, []byte{7}) || t4.Count() != 1 {
131		t.Fatal("wrong content after update+delete")
132	}
133	// Same content as a fresh trie holding {1: 7}.
134	if t4.Root() != mustInsert(t, NewStorage(), k(1), []byte{7}).Root() {
135		t.Fatal("root must depend only on content")
136	}
137}
138
139func randomKeys(r *rand.Rand, n int) []Key {
140	seen := map[Key]bool{}
141	keys := make([]Key, 0, n)
142	for len(keys) < n {
143		var key Key
144		r.Read(key[:])
145		if r.Intn(4) == 0 && len(keys) > 0 {
146			// Share a long prefix with an existing key to exercise deep splices.
147			base := keys[r.Intn(len(keys))]
148			copy(key[:], base[:r.Intn(31)+1])
149		}
150		if !seen[key] {
151			seen[key] = true
152			keys = append(keys, key)
153		}
154	}
155	return keys
156}
157
158func TestOrderIndependence(t *testing.T) {
159	r := rand.New(rand.NewSource(1))
160	keys := randomKeys(r, 300)
161	build := func(order []int) *Trie {
162		tr := NewStorage()
163		for _, i := range order {
164			tr = mustInsert(t, tr, keys[i], keys[i][:4])
165		}
166		return tr
167	}
168	order := r.Perm(len(keys))
169	a := build(order)
170	r.Shuffle(len(order), func(i, j int) { order[i], order[j] = order[j], order[i] })
171	b := build(order)
172	if a.Root() != b.Root() {
173		t.Fatal("root depends on insertion order")
174	}
175}
176
177// checkInvariants verifies branch positions strictly increase along every
178// path, every branch has two children, counts add up, and keys under a
179// branch agree on all bits before its position and split at it.
180func checkInvariants(t *testing.T, tr *Trie) {
181	t.Helper()
182	var walk func(n node, prefix *Key, upto int) uint64
183	walk = func(n node, prefix *Key, upto int) uint64 {
184		switch x := n.(type) {
185		case *leaf:
186			for p := 0; p < upto; p++ {
187				if bit(&x.key, uint8(p)) != bit(prefix, uint8(p)) {
188					t.Fatalf("leaf disagrees with prefix at bit %d", p)
189				}
190			}
191			return 1
192		case *branch:
193			if x.left == nil || x.right == nil {
194				t.Fatal("branch with a missing child")
195			}
196			if int(x.pos) < upto {
197				t.Fatalf("branch position %d not increasing (parent %d)", x.pos, upto-1)
198			}
199			lk := leftmost(x.left)
200			rk := leftmost(x.right)
201			if bit(&lk, x.pos) != 0 || bit(&rk, x.pos) != 1 {
202				t.Fatal("children on the wrong side")
203			}
204			c := walk(x.left, &lk, int(x.pos)+1) + walk(x.right, &rk, int(x.pos)+1)
205			if c != x.n {
206				t.Fatalf("count %d, want %d", x.n, c)
207			}
208			if x.h != BranchHash(tr.counted, x.pos, x.n, x.left.hash(), x.right.hash()) {
209				t.Fatal("stale branch hash")
210			}
211			return c
212		}
213		return 0
214	}
215	if tr.root != nil {
216		lk := leftmost(tr.root)
217		walk(tr.root, &lk, 0)
218	}
219}
220
221func leftmost(n node) Key {
222	for {
223		switch x := n.(type) {
224		case *leaf:
225			return x.key
226		case *branch:
227			n = x.left
228		}
229	}
230}
231
232func TestRandomizedAgainstReference(t *testing.T) {
233	r := rand.New(rand.NewSource(42))
234	keys := randomKeys(r, 2000)
235	ref := map[Key][]byte{}
236	tr := NewStorage()
237	for _, key := range keys {
238		v := make([]byte, r.Intn(40))
239		r.Read(v)
240		ref[key] = v
241		tr = mustInsert(t, tr, key, v)
242	}
243	checkInvariants(t, tr)
244	compare(t, tr, ref)
245
246	// Update a third, delete a third.
247	for i, key := range keys {
248		switch i % 3 {
249		case 0:
250			v := []byte{byte(i)}
251			ref[key] = v
252			n, err := tr.Update(key, v)
253			if err != nil {
254				t.Fatal(err)
255			}
256			tr = n
257		case 1:
258			delete(ref, key)
259			n, err := tr.Delete(key)
260			if err != nil {
261				t.Fatal(err)
262			}
263			tr = n
264		}
265	}
266	checkInvariants(t, tr)
267	compare(t, tr, ref)
268
269	// Rebuild from scratch; must reach the same root.
270	fresh := NewStorage()
271	for key, v := range ref {
272		fresh = mustInsert(t, fresh, key, v)
273	}
274	if fresh.Root() != tr.Root() {
275		t.Fatal("mutated trie differs from freshly built equivalent")
276	}
277
278	// Delete everything.
279	for key := range ref {
280		n, err := tr.Delete(key)
281		if err != nil {
282			t.Fatal(err)
283		}
284		tr = n
285	}
286	if tr.Root() != (Hash{}) || tr.Count() != 0 {
287		t.Fatal("trie not empty after deleting all")
288	}
289}
290
291func compare(t *testing.T, tr *Trie, ref map[Key][]byte) {
292	t.Helper()
293	sorted := make([]Key, 0, len(ref))
294	for key := range ref {
295		sorted = append(sorted, key)
296	}
297	sort.Slice(sorted, func(i, j int) bool { return bytes.Compare(sorted[i][:], sorted[j][:]) < 0 })
298
299	if tr.Count() != uint64(len(sorted)) {
300		t.Fatalf("count %d, want %d", tr.Count(), len(sorted))
301	}
302	for i, key := range sorted {
303		v, ok := tr.Get(key)
304		if !ok || !bytes.Equal(v, ref[key]) {
305			t.Fatalf("get %x: ok=%v", key[:4], ok)
306		}
307		idx, ok := tr.Rank(key)
308		if !ok || idx != uint64(i) {
309			t.Fatalf("rank %x: got %d ok=%v, want %d", key[:4], idx, ok, i)
310		}
311		sk, sv, ok := tr.Select(uint64(i))
312		if !ok || sk != key || !bytes.Equal(sv, ref[key]) {
313			t.Fatalf("select %d: wrong entry", i)
314		}
315	}
316	if _, _, ok := tr.Select(uint64(len(sorted))); ok {
317		t.Fatal("select past the end")
318	}
319	var absent Key
320	absent[0] = 0x55
321	if _, ok := tr.Rank(absent); ok && ref[absent] == nil {
322		t.Fatal("rank of absent key")
323	}
324	i := 0
325	tr.Walk(func(key Key, v []byte) bool {
326		if key != sorted[i] {
327			t.Fatalf("walk order broken at %d", i)
328		}
329		i++
330		return true
331	})
332	if i != len(sorted) {
333		t.Fatal("walk visited wrong number of entries")
334	}
335}
336
337func TestSelectDistribution(t *testing.T) {
338	// Adjacent keys packed into one byte: index must equal the key value.
339	tr := NewStorage()
340	for i := 0; i < 256; i++ {
341		tr = mustInsert(t, tr, k(byte(i)), nil)
342	}
343	for i := 0; i < 256; i++ {
344		key, _, ok := tr.Select(uint64(i))
345		if !ok || key[0] != byte(i) {
346			t.Fatalf("select %d gave %d", i, key[0])
347		}
348	}
349}
350
351func BenchmarkInsert(b *testing.B) {
352	r := rand.New(rand.NewSource(7))
353	keys := randomKeys(r, 100000)
354	tr := NewStorage()
355	b.ResetTimer()
356	for i := 0; i < b.N; i++ {
357		tr, _ = tr.Insert(keys[i%len(keys)], nil)
358		if i%len(keys) == len(keys)-1 {
359			tr = NewStorage()
360		}
361	}
362}