From ef017f0404990ba21c8c5c06af624b47fe58d729 Mon Sep 17 00:00:00 2001
From: Ben Johnson <benbjohnson@yahoo.com>
Date: Tue, 28 Jan 2014 22:50:09 -0500
Subject: [PATCH] Add branch.put().

---
 branch.go      | 30 ++++++++++++++----------------
 branch_test.go | 48 ++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 62 insertions(+), 16 deletions(-)
 create mode 100644 branch_test.go

diff --git a/branch.go b/branch.go
index 77cc1ca..c9fc7ca 100644
--- a/branch.go
+++ b/branch.go
@@ -2,7 +2,6 @@ package bolt
 
 import (
 	"bytes"
-	"sort"
 	"unsafe"
 )
 
@@ -12,24 +11,23 @@ type branch struct {
 	items branchItems
 }
 
-// insert inserts a new item after a given pgid.
-func (b *branch) insert(key []byte, previd pgid, id pgid) {
-	// Find previous insertion index.
-	index := sort.Search(len(b.items), func(i int) bool { return b.items[i].pgid >= previd })
-
-	// If there is no existing key then add a new item.
-	b.items = append(b.items, branchItem{})
-	if index < len(b.items) {
-		copy(b.items[index+1:], b.items[index:])
+// put adds a new node or replaces an existing node.
+func (b *branch) put(id pgid, newid pgid, key []byte, replace bool) {
+	var index int
+	for ; index < len(b.items); index++ {
+		if b.items[index].pgid == id {
+			break
+		}
 	}
 
-	b.items[index].pgid = id
-	b.items[index].key = key
-}
+	if !replace {
+		index++
+		b.items = append(b.items, branchItem{})
+		if index < len(b.items) {
+			copy(b.items[index+1:], b.items[index:])
+		}
+	}
 
-// replace swaps out an existing node id for a new one id.
-func (b *branch) replace(oldid pgid, newid pgid, key []byte) {
-	index := sort.Search(len(b.items), func(i int) bool { return b.items[i].pgid >= oldid })
 	b.items[index].pgid = newid
 	b.items[index].key = key
 }
diff --git a/branch_test.go b/branch_test.go
new file mode 100644
index 0000000..0687b20
--- /dev/null
+++ b/branch_test.go
@@ -0,0 +1,48 @@
+package bolt
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+// Ensure that a branch can replace a key.
+func TestBranchPutReplace(t *testing.T) {
+	b := &branch{
+		items: branchItems{
+			branchItem{pgid: 1, key: []byte("bar")},
+			branchItem{pgid: 2, key: []byte("baz")},
+			branchItem{pgid: 3, key: []byte("foo")},
+		},
+	}
+	b.put(1, 4, []byte("bar"), true)
+	b.put(2, 5, []byte("boo"), true)
+	assert.Equal(t, len(b.items), 3)
+	assert.Equal(t, b.items[0].pgid, pgid(4))
+	assert.Equal(t, string(b.items[0].key), "bar")
+	assert.Equal(t, b.items[1].pgid, pgid(5))
+	assert.Equal(t, string(b.items[1].key), "boo")
+	assert.Equal(t, b.items[2].pgid, pgid(3))
+	assert.Equal(t, string(b.items[2].key), "foo")
+}
+
+// Ensure that a branch can insert a key.
+func TestBranchPutInsert(t *testing.T) {
+	b := &branch{
+		items: branchItems{
+			branchItem{pgid: 1, key: []byte("bar")},
+			branchItem{pgid: 2, key: []byte("foo")},
+		},
+	}
+	b.put(1, 4, []byte("baz"), false)
+	b.put(2, 5, []byte("zzz"), false)
+	assert.Equal(t, len(b.items), 4)
+	assert.Equal(t, b.items[0].pgid, pgid(1))
+	assert.Equal(t, string(b.items[0].key), "bar")
+	assert.Equal(t, b.items[1].pgid, pgid(4))
+	assert.Equal(t, string(b.items[1].key), "baz")
+	assert.Equal(t, b.items[2].pgid, pgid(2))
+	assert.Equal(t, string(b.items[2].key), "foo")
+	assert.Equal(t, b.items[3].pgid, pgid(5))
+	assert.Equal(t, string(b.items[3].key), "zzz")
+}