Fixes #4 - Add immutable remove operation

trie/remove.go

@@ -26,3 +26,29 @@ func (trie *Trie) RemoveAtDepth(depth int, q key.Key) (reachedDepth int, removed
 		}
 	}
 }
+
+func Remove(trie *Trie, q key.Key) *Trie {
+	return RemoveAtDepth(0, trie, q)
+}
+
+func RemoveAtDepth(depth int, trie *Trie, q key.Key) *Trie {
+	switch {
+	case trie.IsEmptyLeaf():
+		return trie
+	case trie.IsNonEmptyLeaf() && !key.Equal(trie.Key, q):
+		return trie
+	case trie.IsNonEmptyLeaf() && key.Equal(trie.Key, q):
+		return &Trie{}
+	default:
+		dir := q.BitAt(depth)
+		afterDelete := RemoveAtDepth(depth+1, trie.Branch[dir], q)
+		if afterDelete == trie.Branch[dir] {
+			return trie
+		}
+		copy := &Trie{}
+		copy.Branch[dir] = afterDelete
+		copy.Branch[1-dir] = trie.Branch[1-dir]
+		copy.shrink()
+		return copy
+	}
+}

trie/remove_test.go

@@ -0,0 +1,72 @@
+package trie
+
+import (
+	"github.com/libp2p/go-libp2p-xor/key"
+	"math/rand"
+	"testing"
+)
+
+func TestImmutableRemoveIsImmutable(t *testing.T) {
+	for _, keySet := range testAddSamples {
+		trie := FromKeys(keySet.Keys)
+		for _, key := range keySet.Keys {
+			updated := Remove(trie, key)
+			if Equal(trie, updated) {
+				t.Fatalf("immuatble remove should not mutate trie, original: %v, updated: %v", trie, updated)
+			}
+			trie = updated
+		}
+	}
+}
+
+func TestMutableAndImmutableRemoveSame(t *testing.T) {
+	for _, keySet := range append(testAddSamples, randomTestAddSamples(100)...) {
+		mut := FromKeys(keySet.Keys)
+		immut := FromKeys(keySet.Keys)
+
+		for _, key := range keySet.Keys {
+			mut.Remove(key)
+			immut = Remove(immut, key)
+			if d := mut.CheckInvariant(); d != nil {
+				t.Fatalf("mutable trie invariant discrepancy: %v", d)
+			}
+			if d := immut.CheckInvariant(); d != nil {
+				t.Fatalf("immutable trie invariant discrepancy: %v", d)
+			}
+			if !Equal(mut, immut) {
+				t.Errorf("mutable trie %v differs from immutable trie %v", mut, immut)
+			}
+		}
+	}
+}
+
+func TestRemoveIsOrderIndependent(t *testing.T) {
+	for _, keySet := range append(testAddSamples, randomTestAddSamples(100)...) {
+		mut := FromKeys(keySet.Keys)
+		immut := FromKeys(keySet.Keys)
+
+		for j := 0; j < 100; j++ {
+			perm := rand.Perm(len(keySet.Keys))
+			for _, idx := range perm {
+				mut.Remove(keySet.Keys[idx])
+				immut = Remove(immut, keySet.Keys[idx])
+
+				if d := immut.CheckInvariant(); d != nil {
+					t.Fatalf("trie invariant discrepancy: %v", d)
+				}
+				if !Equal(mut, immut) {
+					t.Errorf("trie %v differs from trie %v", mut, immut)
+				}
+			}
+		}
+	}
+}
+
+func TestRemoveReturnsOriginalWhenNoKeyRemoved(t *testing.T) {
+	trie := FromKeys(testAddSamples[0].Keys)
+
+	result := Remove(trie, key.ByteKey(2))
+	if trie != result {
+		t.Fatalf("Remove should return original trie when no key was removed")
+	}
+}