Merge pull request #187 from benbjohnson/fix-seek

Fix in-between page seek

2 files changed, 87 insertions, 22 deletions

diff --git a/cursor.go b/cursor.go
index 1c14d05..dbc2833 100644
--- a/cursor.go
+++ b/cursor.go
@@ -55,26 +55,7 @@ func (c *Cursor) Last() (key []byte, value []byte) {
 // If the cursor is at the end of the bucket then a nil key and value are returned.
 func (c *Cursor) Next() (key []byte, value []byte) {
 	_assert(c.bucket.tx.db != nil, "tx closed")
-
-	// Attempt to move over one element until we're successful.
-	// Move up the stack as we hit the end of each page in our stack.
-	for i := len(c.stack) - 1; i >= 0; i-- {
-		elem := &c.stack[i]
-		if elem.index < elem.count()-1 {
-			elem.index++
-			break
-		}
-		c.stack = c.stack[:i]
-	}
-
-	// If we've hit the end then return nil.
-	if len(c.stack) == 0 {
-		return nil, nil
-	}
-
-	// Move down the stack to find the first element of the first leaf under this branch.
-	c.first()
-	k, v, flags := c.keyValue()
+	k, v, flags := c.next()
 	if (flags & uint32(bucketLeafFlag)) != 0 {
 		return k, nil
 	}
@@ -116,6 +97,12 @@ func (c *Cursor) Prev() (key []byte, value []byte) {
 // follow, a nil key is returned.
 func (c *Cursor) Seek(seek []byte) (key []byte, value []byte) {
 	k, v, flags := c.seek(seek)
+
+	// If we ended up after the last element of a page then move to the next one.
+	if ref := &c.stack[len(c.stack)-1]; ref.index >= ref.count() {
+		k, v, flags = c.next()
+	}
+
 	if k == nil {
 		return nil, nil
 	} else if (flags & uint32(bucketLeafFlag)) != 0 {
@@ -125,8 +112,7 @@ func (c *Cursor) Seek(seek []byte) (key []byte, value []byte) {
 }
 
 // seek moves the cursor to a given key and returns it.
-// If the key does not exist then the next key is used. If no keys
-// follow, a nil value is returned.
+// If the key does not exist then the next key is used.
 func (c *Cursor) seek(seek []byte) (key []byte, value []byte, flags uint32) {
 	_assert(c.bucket.tx.db != nil, "tx closed")
 
@@ -189,6 +175,33 @@ func (c *Cursor) last() {
 	}
 }
 
+// next moves to the next leaf element and returns the key and value.
+// If the cursor is at the last leaf element then it stays there and returns nil.
+func (c *Cursor) next() (key []byte, value []byte, flags uint32) {
+	// Attempt to move over one element until we're successful.
+	// Move up the stack as we hit the end of each page in our stack.
+	var i int
+	for i = len(c.stack) - 1; i >= 0; i-- {
+		elem := &c.stack[i]
+		if elem.index < elem.count()-1 {
+			elem.index++
+			break
+		}
+	}
+
+	// If we've hit the root page then stop and return. This will leave the
+	// cursor on the last element of the last page.
+	if i == -1 {
+		return nil, nil, 0
+	}
+
+	// Otherwise start from where we left off in the stack and find the
+	// first element of the first leaf page.
+	c.stack = c.stack[:i+1]
+	c.first()
+	return c.keyValue()
+}
+
 // search recursively performs a binary search against a given page/node until it finds a given key.
 func (c *Cursor) search(key []byte, pgid pgid) {
 	p, n := c.bucket.pageNode(pgid)
diff --git a/cursor_test.go b/cursor_test.go
index 23b3a8e..b44bf53 100644
--- a/cursor_test.go
+++ b/cursor_test.go
@@ -1,6 +1,7 @@
 package bolt
 
 import (
+	"encoding/binary"
 	"sort"
 	"testing"
 	"testing/quick"
@@ -66,6 +67,57 @@ func TestCursor_Seek(t *testing.T) {
 	})
 }
 
+// Ensure that a Tx cursor can seek to the appropriate keys when there are a
+// large number of keys. This test also checks that seek will always move
+// forward to the next key.
+//
+// Related: https://github.com/boltdb/bolt/pull/187
+func TestCursor_Seek_Large(t *testing.T) {
+	withOpenDB(func(db *DB, path string) {
+		var count = 10000
+
+		// Insert every other key between 0 and $count.
+		db.Update(func(tx *Tx) error {
+			b, _ := tx.CreateBucket([]byte("widgets"))
+			for i := 0; i < count; i += 100 {
+				for j := i; j < i+100; j += 2 {
+					k := make([]byte, 8)
+					binary.BigEndian.PutUint64(k, uint64(j))
+					b.Put(k, make([]byte, 100))
+				}
+			}
+			return nil
+		})
+
+		db.View(func(tx *Tx) error {
+			c := tx.Bucket([]byte("widgets")).Cursor()
+			for i := 0; i < count; i++ {
+				seek := make([]byte, 8)
+				binary.BigEndian.PutUint64(seek, uint64(i))
+
+				k, _ := c.Seek(seek)
+
+				// The last seek is beyond the end of the the range so
+				// it should return nil.
+				if i == count-1 {
+					assert.Nil(t, k)
+					continue
+				}
+
+				// Otherwise we should seek to the exact key or the next key.
+				num := binary.BigEndian.Uint64(k)
+				if i%2 == 0 {
+					assert.Equal(t, uint64(i), num)
+				} else {
+					assert.Equal(t, uint64(i+1), num)
+				}
+			}
+
+			return nil
+		})
+	})
+}
+
 // Ensure that a cursor can iterate over an empty bucket without error.
 func TestCursor_EmptyBucket(t *testing.T) {
 	withOpenDB(func(db *DB, path string) {