From 0ac1eea9be091b6c4b2838dc5309b9ec77c5bdf1 Mon Sep 17 00:00:00 2001
From: tidwall <joshbaker77@gmail.com>
Date: Thu, 11 Oct 2018 03:58:43 -0700
Subject: [PATCH] Fixed panic on set after delete

Fix btree reset single item
---
 pkg/ds/btree.go      |  1 +
 pkg/ds/btree_test.go | 46 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 47 insertions(+)

diff --git a/pkg/ds/btree.go b/pkg/ds/btree.go
index 7b2d011e..9a2386ea 100644
--- a/pkg/ds/btree.go
+++ b/pkg/ds/btree.go
@@ -195,6 +195,7 @@ func (tr *BTree) Delete(key string) (prev interface{}, deleted bool) {
 	tr.length--
 	if tr.length == 0 {
 		tr.root = nil
+		tr.height = 0
 	}
 	return
 }
diff --git a/pkg/ds/btree_test.go b/pkg/ds/btree_test.go
index cb6a04ee..5f4ad370 100644
--- a/pkg/ds/btree_test.go
+++ b/pkg/ds/btree_test.go
@@ -428,3 +428,49 @@ func BenchmarkTidwallGet(b *testing.B) {
 		tr.Get(keys[i])
 	}
 }
+
+func TestBTreeOne(t *testing.T) {
+	var tr BTree
+	tr.Set("1", "1")
+	tr.Delete("1")
+	tr.Set("1", "1")
+	tr.Delete("1")
+	tr.Set("1", "1")
+	tr.Delete("1")
+}
+
+func TestBTree256(t *testing.T) {
+	var tr BTree
+	var n int
+	for j := 0; j < 2; j++ {
+		for _, i := range rand.Perm(256) {
+			tr.Set(fmt.Sprintf("%d", i), i)
+			n++
+			if tr.Len() != n {
+				t.Fatalf("expected 256, got %d", n)
+			}
+		}
+		for _, i := range rand.Perm(256) {
+			v, ok := tr.Get(fmt.Sprintf("%d", i))
+			if !ok {
+				t.Fatal("expected true")
+			}
+			if v.(int) != i {
+				t.Fatalf("expected %d, got %d", i, v.(int))
+			}
+		}
+		for _, i := range rand.Perm(256) {
+			tr.Delete(fmt.Sprintf("%d", i))
+			n--
+			if tr.Len() != n {
+				t.Fatalf("expected 256, got %d", n)
+			}
+		}
+		for _, i := range rand.Perm(256) {
+			_, ok := tr.Get(fmt.Sprintf("%d", i))
+			if ok {
+				t.Fatal("expected false")
+			}
+		}
+	}
+}