From 334493d73655d30e8cd8d3bd6f051185e6e5132c Mon Sep 17 00:00:00 2001
From: Ondrej Prazak <o.prazak5@seznam.cz>
Date: Sun, 4 Jul 2021 12:05:23 +0200
Subject: [PATCH] Refs #4 - Remove original trie when no key removed

---
 trie/remove.go      | 12 ++++++++++--
 trie/remove_test.go | 10 ++++++++++
 2 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/trie/remove.go b/trie/remove.go
index 52986de..571ae3c 100644
--- a/trie/remove.go
+++ b/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
diff --git a/trie/remove_test.go b/trie/remove_test.go
index 24597ba..9036ee5 100644
--- a/trie/remove_test.go
+++ b/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")
+	}
+}