Refs #4 - Remove original trie when no key removed

trie/remove.go

@@ -33,12 +33,20 @@ func Remove(trie *Trie, q key.Key) (*Trie) {
 
 func RemoveAtDepth(depth int, trie *Trie, q key.Key) (*Trie) {
 	switch {
-	case trie.IsEmptyLeaf() || trie.IsNonEmptyLeaf():
+	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] = RemoveAtDepth(depth + 1, trie.Branch[dir], q)
+		copy.Branch[dir] = afterDelete
 		copy.Branch[1 - dir] = trie.Branch[1 - dir]
 		copy.shrink()
 		return copy

trie/remove_test.go

@@ -3,6 +3,7 @@ package trie
 import (
 	"testing"
 	"math/rand"
+	"github.com/libp2p/go-libp2p-xor/key"
 )
 
 func TestImmutableRemoveIsImmutable(t *testing.T) {
@@ -60,3 +61,12 @@ func TestRemoveIsOrderIndependent(t *testing.T) {
 		}
 	}
 }
+
+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")
+	}
+}