Bug fixes in add, remove and find.

key/key.go

 package key
 
-import "bytes"
+import (
+	"bytes"
+	"fmt"
+	"strings"
+)
 
 // Key is a vector of bits backed by a Go byte slice in big endian byte order and big-endian bit order.
 type Key []byte
 
-func (bs Key) BitAt(offset int) byte {
-	if bs[offset/8]&(1<<(offset%8)) == 0 {
+func (k Key) BitAt(offset int) byte {
+	if k[offset/8]&(1<<(offset%8)) == 0 {
 		return 0
 	} else {
 		return 1
 	}
 }
 
-func (bs Key) BitLen() int {
-	return 8 * len(bs)
+func (k Key) BitLen() int {
+	return 8 * len(k)
+}
+
+func (k Key) String() string {
+	s := make([]string, len(k))
+	for i, b := range k {
+		s[len(k)-i-1] = fmt.Sprintf("%08b", b)
+	}
+	return strings.Join(s, "")
 }
 
 func Equal(x, y Key) bool {

key/key_test.go

+package key
+
+import "testing"
+
+func TestKeyString(t *testing.T) {
+	key := Key{0x05, 0xf0}
+	println(key.String())
+}

trie/add.go

@@ -11,24 +11,23 @@ func (trie *Trie) Add(q key.Key) (insertedDepth int, insertedOK bool) {
 }
 
 func (trie *Trie) AddAtDepth(depth int, q key.Key) (insertedDepth int, insertedOK bool) {
-	if qb := trie.Branch[q.BitAt(depth)]; qb != nil {
-		return qb.AddAtDepth(depth+1, q)
-	} else {
-		if trie.Key == nil {
-			trie.Key = q
-			return depth, true
+	switch {
+	case trie.IsEmptyLeaf():
+		trie.Key = q
+		return depth, true
+	case trie.IsNonEmptyLeaf():
+		if key.Equal(trie.Key, q) {
+			return depth, false
 		} else {
-			if key.Equal(trie.Key, q) {
-				return depth, false
-			} else {
-				p := trie.Key
-				trie.Key = nil
-				// both Branches are nil
-				trie.Branch[0], trie.Branch[1] = &Trie{}, &Trie{}
-				trie.Branch[p.BitAt(depth)].AddAtDepth(depth+1, p)
-				return trie.Branch[q.BitAt(depth)].AddAtDepth(depth+1, q)
-			}
+			p := trie.Key
+			trie.Key = nil
+			// both branches are nil
+			trie.Branch[0], trie.Branch[1] = &Trie{}, &Trie{}
+			trie.Branch[p.BitAt(depth)].AddAtDepth(depth+1, p)
+			return trie.Branch[q.BitAt(depth)].AddAtDepth(depth+1, q)
 		}
+	default:
+		return trie.Branch[q.BitAt(depth)].AddAtDepth(depth+1, q)
 	}
 }
 
@@ -39,30 +38,30 @@ func Add(trie *Trie, q key.Key) *Trie {
 }
 
 func AddAtDepth(depth int, trie *Trie, q key.Key) *Trie {
-	dir := q.BitAt(depth)
-	if !trie.IsLeaf() {
-		s := &Trie{}
-		s.Branch[dir] = AddAtDepth(depth+1, trie.Branch[dir], q)
-		s.Branch[1-dir] = trie.Branch[1-dir]
-		return s
-	} else {
-		if trie.Key == nil {
-			return &Trie{Key: q}
+	switch {
+	case trie.IsEmptyLeaf():
+		return &Trie{Key: q}
+	case trie.IsNonEmptyLeaf():
+		if key.Equal(trie.Key, q) {
+			return trie
 		} else {
-			if key.Equal(trie.Key, q) {
-				return trie
-			} else {
-				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{}
-				}
+			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
 		}
+	default:
+		dir := q.BitAt(depth)
+		s := &Trie{}
+		s.Branch[dir] = AddAtDepth(depth+1, trie.Branch[dir], q)
+		s.Branch[1-dir] = trie.Branch[1-dir]
+		return s
 	}
 }

trie/add_test.go

@@ -27,4 +27,5 @@ type testAddSameSample struct {
 
 var testAddSameSamples = []*testAddSameSample{
 	{Keys: []key.Key{{1, 3, 5, 7, 11, 13}}},
+	{Keys: []key.Key{{11, 22, 23, 25, 27, 28, 31, 32, 33}}},
 }

trie/find.go

@@ -12,13 +12,12 @@ func (trie *Trie) Find(q key.Key) (reachedDepth int, found bool) {
 }
 
 func (trie *Trie) FindAtDepth(depth int, q key.Key) (reachedDepth int, found bool) {
-	if qb := trie.Branch[q.BitAt(depth)]; qb != nil {
-		return qb.FindAtDepth(depth+1, q)
-	} else {
-		if trie.Key == nil {
-			return depth, false
-		} else {
-			return depth, key.Equal(trie.Key, q)
-		}
+	switch {
+	case trie.IsEmptyLeaf():
+		return depth, false
+	case trie.IsNonEmptyLeaf():
+		return depth, key.Equal(trie.Key, q)
+	default:
+		return trie.Branch[q.BitAt(depth)].FindAtDepth(depth+1, q)
 	}
 }

trie/intersect_test.go

 package trie
 
 import (
+	"math/rand"
 	"testing"
 
 	"github.com/libp2p/go-libp2p-xor/key"
 )
 
 func TestIntersectRandom(t *testing.T) {
+	for i := 0; i < 100; i++ {
+		testIntersect(t, randomTestIntersectSample(10, 10, 5))
+	}
+}
+
+func TestIntersect(t *testing.T) {
 	for _, s := range testIntersectSamples {
 		testIntersect(t, s)
 	}
@@ -42,6 +49,17 @@ func setIntersect(left, right []key.Key) []key.Key {
 	return intersection
 }
 
+func randomTestIntersectSample(leftSize, rightSize, intersectSize int) *testIntersectSample {
+	keys := make([]key.Key, leftSize+rightSize-intersectSize)
+	for i := range keys {
+		keys[i] = key.Key{byte(rand.Intn(256))}
+	}
+	return &testIntersectSample{
+		LeftKeys:  keys[:leftSize],
+		RightKeys: keys[leftSize-intersectSize:],
+	}
+}
+
 type testIntersectSample struct {
 	LeftKeys  []key.Key
 	RightKeys []key.Key

trie/remove.go

@@ -5,25 +5,24 @@ import (
 )
 
 // Remove removes the key q from the trie. Remove mutates the trie.
-// TODO: Also implement an immutable version of Add.
+// TODO: Also implement an immutable version of Remove.
 func (trie *Trie) Remove(q key.Key) (removedDepth int, removed bool) {
 	return trie.RemoveAtDepth(0, q)
 }
 
 func (trie *Trie) RemoveAtDepth(depth int, q key.Key) (reachedDepth int, removed bool) {
-	if qb := trie.Branch[q.BitAt(depth)]; qb != nil {
-		if d, ok := qb.RemoveAtDepth(depth+1, q); ok {
+	switch {
+	case trie.IsEmptyLeaf():
+		return depth, false
+	case trie.IsNonEmptyLeaf():
+		trie.Key = nil
+		return depth, true
+	default:
+		if d, removed := trie.Branch[q.BitAt(depth)].RemoveAtDepth(depth+1, q); removed {
 			trie.shrink()
 			return d, true
 		} else {
 			return d, false
 		}
-	} else {
-		if trie.Key != nil && key.Equal(q, trie.Key) {
-			trie.Key = nil
-			return depth, true
-		} else {
-			return depth, false
-		}
 	}
 }