Various fixes and test additions.

key/key.go

@@ -11,7 +11,7 @@ import (
 type Key []byte
 
 func (k Key) BitAt(offset int) byte {
-	if k[offset/8]&(1<<(offset%8)) == 0 {
+	if k[offset/8]&(byte(1)<<(offset%8)) == 0 {
 		return 0
 	} else {
 		return 1

key/key_test.go

@@ -4,5 +4,23 @@ import "testing"
 
 func TestKeyString(t *testing.T) {
 	key := Key{0x05, 0xf0}
-	println(key.String())
+	if key.BitString() != "1111000000000101" {
+		t.Errorf("unexpected bit string: %s", key.BitString())
+	}
+}
+
+func TestBitAt(t *testing.T) {
+	key := Key{0x21, 0x84}
+	switch {
+	case key.BitAt(0) != 1:
+		t.Errorf("bit 0 flipped")
+	case key.BitAt(4+1) != 1:
+		t.Errorf("bit 5 flipped")
+	case key.BitAt(8) != 0:
+		t.Errorf("bit 8 flipped")
+	case key.BitAt(8+2) != 1:
+		t.Errorf("bit 10 flipped")
+	case key.BitAt(8+4+3) != 1:
+		t.Errorf("bit 15 flipped")
+	}
 }

trie/add.go

@@ -22,7 +22,7 @@ func (trie *Trie) AddAtDepth(depth int, q key.Key) (insertedDepth int, insertedO
 			trie.Key = nil
 			// both branches are nil
 			trie.Branch[0], trie.Branch[1] = &Trie{}, &Trie{}
-			trie.Branch[p.BitAt(depth)].AddAtDepth(depth+1, p)
+			trie.Branch[p.BitAt(depth)].Key = p
 			return trie.Branch[q.BitAt(depth)].AddAtDepth(depth+1, q)
 		}
 	default:
@@ -44,17 +44,7 @@ func AddAtDepth(depth int, trie *Trie, q key.Key) *Trie {
 		if key.Equal(trie.Key, q) {
 			return trie
 		} else {
-			dir := q.BitAt(depth)
-			s := &Trie{}
-			if q.BitAt(depth) == trie.Key.BitAt(depth) {
-				s.Branch[dir] = AddAtDepth(depth+1, &Trie{Key: trie.Key}, q)
-				s.Branch[1-dir] = &Trie{}
-				return s
-			} else {
-				s.Branch[dir] = AddAtDepth(depth+1, &Trie{Key: trie.Key}, q)
-				s.Branch[1-dir] = &Trie{}
-			}
-			return s
+			return trieForTwo(depth, trie.Key, q)
 		}
 	default:
 		dir := q.BitAt(depth)
@@ -64,3 +54,18 @@ func AddAtDepth(depth int, trie *Trie, q key.Key) *Trie {
 		return s
 	}
 }
+
+func trieForTwo(depth int, p, q key.Key) *Trie {
+	pDir, qDir := p.BitAt(depth), q.BitAt(depth)
+	if qDir == pDir {
+		s := &Trie{}
+		s.Branch[pDir] = trieForTwo(depth+1, p, q)
+		s.Branch[1-pDir] = &Trie{}
+		return s
+	} else {
+		s := &Trie{}
+		s.Branch[pDir] = &Trie{Key: p}
+		s.Branch[qDir] = &Trie{Key: q}
+		return s
+	}
+}

trie/add_test.go

@@ -9,7 +9,7 @@ import (
 
 // Verify mutable and immutable add do the same thing.
 func TestMutableAndImmutableAddSame(t *testing.T) {
-	for _, s := range testAddSameSamples {
+	for _, s := range append(testAddSamples, randomTestAddSamples(100)...) {
 		mut := New()
 		immut := New()
 		for _, k := range s.Keys {
@@ -25,10 +25,11 @@ func TestMutableAndImmutableAddSame(t *testing.T) {
 }
 
 func TestAddIsOrderIndependent(t *testing.T) {
-	for _, s := range testAddSameSamples {
+	for _, s := range append(testAddSamples, randomTestAddSamples(100)...) {
 		base := New()
 		for _, k := range s.Keys {
 			base.Add(k)
+			base.CheckInvariant()
 		}
 		base.CheckInvariant()
 		for j := 0; j < 100; j++ {
@@ -45,11 +46,31 @@ func TestAddIsOrderIndependent(t *testing.T) {
 	}
 }
 
-type testAddSameSample struct {
+type testAddSample struct {
 	Keys []key.Key
 }
 
-var testAddSameSamples = []*testAddSameSample{
-	{Keys: []key.Key{{1, 3, 5, 7, 11, 13}}},
-	{Keys: []key.Key{{11, 22, 23, 25, 27, 28, 31, 32, 33}}},
+var testAddSamples = []*testAddSample{
+	{Keys: []key.Key{{1}, {3}, {5}, {7}, {11}, {13}}},
+	{Keys: []key.Key{{11}, {22}, {23}, {25}, {27}, {28}, {31}, {32}, {33}}},
+}
+
+func randomTestAddSamples(count int) []*testAddSample {
+	s := make([]*testAddSample, count)
+	for i := range s {
+		s[i] = randomTestAddSample(31, 2)
+	}
+	return s
+}
+
+func randomTestAddSample(setSize, keySizeByte int) *testAddSample {
+	keySet := make([]key.Key, setSize)
+	for i := range keySet {
+		k := make(key.Key, keySizeByte)
+		rand.Read(k)
+		keySet[i] = k
+	}
+	return &testAddSample{
+		Keys: keySet,
+	}
 }

trie/check.go

 package trie
 
+import (
+	"github.com/libp2p/go-libp2p-xor/key"
+)
+
 // CheckInvariant panics of the trie does not meet its invariant.
 func (trie *Trie) CheckInvariant() {
+	trie.checkInvariant(0, nil)
+}
+
+func (trie *Trie) checkInvariant(depth int, pathSoFar *triePath) {
 	switch {
-	case trie.IsLeaf():
-		return
+	case trie.IsEmptyLeaf(): // ok
+	case trie.IsNonEmptyLeaf():
+		if !pathSoFar.matchesKey(trie.Key) {
+			panic("key found at invalid location in trie")
+		}
 	default:
 		if trie.IsEmpty() {
 			b0, b1 := trie.Branch[0], trie.Branch[1]
-			b0.CheckInvariant()
-			b1.CheckInvariant()
+			b0.checkInvariant(depth+1, pathSoFar.Push(0))
+			b1.checkInvariant(depth+1, pathSoFar.Push(1))
 			switch {
 			case b0.IsEmptyLeaf() && b1.IsEmptyLeaf():
 				panic("intermediate node with two empty leaves")
@@ -23,3 +34,60 @@ func (trie *Trie) CheckInvariant() {
 		}
 	}
 }
+
+type triePath struct {
+	parent *triePath
+	bit    byte
+}
+
+func (p *triePath) Push(bit byte) *triePath {
+	return &triePath{parent: p, bit: bit}
+}
+
+func (p *triePath) RootPath() []byte {
+	if p == nil {
+		return nil
+	} else {
+		return append(p.parent.RootPath(), p.bit)
+	}
+}
+
+func (p *triePath) matchesKey(k key.Key) bool {
+	// Slower, but more explicit:
+	for i, b := range p.RootPath() {
+		if k.BitAt(i) != b {
+			return false
+		}
+	}
+	return true
+	// ok, _ := p.matchesKeyWalk(k, 0)
+	// return ok
+}
+
+func (p *triePath) matchesKeyWalk(k key.Key, depthToLeaf int) (ok bool, depthToRoot int) {
+	if p == nil {
+		return true, 0
+	} else {
+		parOk, parDepthToRoot := p.parent.matchesKeyWalk(k, depthToLeaf+1)
+		return k.BitAt(parDepthToRoot+1) == p.bit && parOk, parDepthToRoot + 1
+	}
+}
+
+func (p *triePath) String() string {
+	return p.string(0)
+}
+
+func (p *triePath) string(depthToLeaf int) string {
+	if p == nil {
+		return ""
+	} else {
+		switch {
+		case p.bit == 0:
+			return p.parent.string(depthToLeaf+1) + "0"
+		case p.bit == 1:
+			return p.parent.string(depthToLeaf+1) + "1"
+		default:
+			panic("bit digit > 1")
+		}
+	}
+}