add deep tree tests, some cleanup

2 files changed, 106 insertions, 32 deletions

diff --git a/c/cursor.go b/c/cursor.go
index 49fdcd8..53756cd 100644
--- a/c/cursor.go
+++ b/c/cursor.go
@@ -88,8 +88,6 @@ elem_ref *cursor_current(bolt_cursor *c);
 
 elem_ref *cursor_pop(bolt_cursor *c);
 
-void cursor_first_leaf(bolt_cursor *c);
-
 void cursor_key_value(bolt_cursor *c, bolt_val *key, bolt_val *value, uint32_t *flags);
 
 void cursor_search(bolt_cursor *c, bolt_val key, pgid id);
@@ -116,7 +114,6 @@ void bolt_cursor_first(bolt_cursor *c, bolt_val *key, bolt_val *value, uint32_t
 	elem_ref *ref = cursor_push(c, c->root);
 
 	// Find first leaf and return key/value.
-	cursor_first_leaf(c);
 	cursor_key_value(c, key, value, flags);
 }
 
@@ -133,16 +130,7 @@ void bolt_cursor_next(bolt_cursor *c, bolt_val *key, bolt_val *value, uint32_t *
 		cursor_pop(c);
 	};
 
-	// If we are at the top of the stack then return a blank key/value pair.
-	if (ref == NULL) {
-		key->size = value->size = 0;
-		key->data = value->data = NULL;
-		*flags = 0;
-		return;
-	};
-
 	// Find first leaf and return key/value.
-	cursor_first_leaf(c);
 	cursor_key_value(c, key, value, flags);
 }
 
@@ -153,18 +141,8 @@ void bolt_cursor_seek(bolt_cursor *c, bolt_val seek, bolt_val *key, bolt_val *va
 	// Start from root page/node and traverse to correct page.
 	cursor_push(c, c->root);
 	if (seek.size > 0) cursor_search(c, seek, c->root);
-	elem_ref *ref = cursor_current(c);
-
-	// If the cursor is pointing to the end of page then return nil.
-	if (ref == NULL) {
-		key->size = value->size = 0;
-		key->data = value->data = NULL;
-		*flags = 0;
-		return;
-	};
 
 	// Find first leaf and return key/value.
-	cursor_first_leaf(c);
 	cursor_key_value(c, key, value, flags);
 }
 
@@ -216,6 +194,19 @@ leaf_element *page_leaf_element(page *p, uint16_t index) {
 // Returns the key/value pair for the current position of the cursor.
 void cursor_key_value(bolt_cursor *c, bolt_val *key, bolt_val *value, uint32_t *flags) {
 	elem_ref *ref = cursor_current(c);
+	if (ref == NULL) {
+		key->size = value->size = 0;
+		key->data = value->data = NULL;
+		*flags = 0;
+		return;
+	};
+
+	// Descend to the current leaf page if we're on branch page
+	while (ref->page->flags & PAGE_BRANCH) {
+		branch_element *elem = page_branch_element(ref->page,ref->index);
+		ref = cursor_push(c, elem->pgid);
+	};
+
 	leaf_element *elem = page_leaf_element(ref->page,ref->index);
 
 	// Assign key pointer.
@@ -230,15 +221,6 @@ void cursor_key_value(bolt_cursor *c, bolt_val *key, bolt_val *value, uint32_t *
 	*flags = elem->flags;
 }
 
-// Traverses from the current stack position down to the first leaf element.
-void cursor_first_leaf(bolt_cursor *c) {
-	elem_ref *ref = cursor_current(c);
-	while (ref->page->flags & PAGE_BRANCH) {
-		branch_element *elem = page_branch_element(ref->page,ref->index);
-		ref = cursor_push(c, elem->pgid);
-	};
-}
-
 // Recursively performs a binary search against a given page/node until it finds a given key.
 void cursor_search(bolt_cursor *c, bolt_val key, pgid id) {
 	// Push page onto the cursor stack.
@@ -267,7 +249,7 @@ void cursor_search_leaf(bolt_cursor *c, bolt_val key) {
 
 		// printf("? %.*s | %.*s\n", key.size, key.data, elem->ksize, ((void*)elem) + elem->pos);
 		// printf("rc=%d; key.size(%d) >= elem->ksize(%d)\n", rc, key.size, elem->ksize);
-		if ((rc == 0 && key.size >= elem->ksize) || rc < 0) {
+		if ((rc == 0 && key.size <= elem->ksize) || rc < 0) {
 			ref->index = i;
 			return;
 		}
@@ -341,6 +323,9 @@ func (c *Cursor) First() (key, value []byte) {
 	var k, v C.bolt_val
 	var flags C.uint32_t
 	C.bolt_cursor_first(c.C, &k, &v, &flags)
+	if k.data == nil {
+		return nil, nil
+	}
 	return C.GoBytes(k.data, C.int(k.size)), C.GoBytes(v.data, C.int(v.size))
 }
 
@@ -350,6 +335,9 @@ func (c *Cursor) Next() (key, value []byte) {
 	var k, v C.bolt_val
 	var flags C.uint32_t
 	C.bolt_cursor_next(c.C, &k, &v, &flags)
+	if k.data == nil {
+		return nil, nil
+	}
 	return C.GoBytes(k.data, C.int(k.size)), C.GoBytes(v.data, C.int(v.size))
 }
 
@@ -369,6 +357,7 @@ func (c *Cursor) Seek(seek []byte) (key, value []byte, flags int) {
 	if k.data == nil {
 		return nil, nil, 0
 	}
+
 	return C.GoBytes(k.data, C.int(k.size)), C.GoBytes(v.data, C.int(v.size)), int(_flags)
 }
 
diff --git a/c/cursor_test.go b/c/cursor_test.go
index 9e7cb1b..46d2a58 100644
--- a/c/cursor_test.go
+++ b/c/cursor_test.go
@@ -140,6 +140,40 @@ func TestCursor_Iterate_Large(t *testing.T) {
 	})
 }
 
+// Ensure that a C cursor can iterate over branches and leafs.
+func TestCursor_Iterate_Deep(t *testing.T) {
+	withDB(func(db *bolt.DB) {
+		pgsz := db.Info().PageSize / 10
+		assert.True(t, pgsz > 100)
+		db.Update(func(tx *bolt.Tx) error {
+			b, _ := tx.CreateBucket([]byte("widgets"))
+			for i := 0; i < 1000; i++ {
+				kv := []byte(fmt.Sprintf("%0*d", pgsz, i))
+				b.Put(kv, kv)
+			}
+			return nil
+		})
+		db.View(func(tx *bolt.Tx) error {
+			b := tx.Bucket([]byte("widgets"))
+			s := b.Stats()
+			assert.True(t, s.Depth > 3)
+
+			var index int
+			c := NewCursor(b)
+			for k, v := c.First(); len(k) > 0; k, v = c.Next() {
+				kv := fmt.Sprintf("%0*d", pgsz, index)
+				assert.Equal(t, kv, string(k))
+				assert.Equal(t, kv, string(v))
+				index++
+			}
+			assert.Equal(t, 1000, index)
+			k, _ := c.Next()
+			assert.Equal(t, nil, k)
+			return nil
+		})
+	})
+}
+
 // Ensure that a C cursor can seek over branches and leafs.
 func TestCursor_Seek_Large(t *testing.T) {
 	withDB(func(db *bolt.DB) {
@@ -178,6 +212,57 @@ func TestCursor_Seek_Large(t *testing.T) {
 	})
 }
 
+// Ensure that a C cursor can seek over branches and leafs.
+func TestCursor_Seek_Deep(t *testing.T) {
+	withDB(func(db *bolt.DB) {
+		pgsz := db.Info().PageSize / 10
+		assert.True(t, pgsz > 100)
+		db.Update(func(tx *bolt.Tx) error {
+			b, _ := tx.CreateBucket([]byte("widgets"))
+			for i := 1; i < 1000; i++ {
+				kv := []byte(fmt.Sprintf("%0*d", pgsz, i*10))
+				b.Put(kv, kv)
+			}
+			return nil
+		})
+		db.View(func(tx *bolt.Tx) error {
+			b := tx.Bucket([]byte("widgets"))
+			s := b.Stats()
+			assert.True(t, s.Depth > 3)
+
+			c := NewCursor(b)
+
+			// Exact match should go to the key.
+			seek := fmt.Sprintf("%0*d", pgsz, 5000)
+			k, v, _ := c.Seek([]byte(seek))
+			assert.Equal(t, seek, string(k))
+			assert.Equal(t, seek, string(v))
+
+			// Inexact match should go to the next key.
+			seek = fmt.Sprintf("%0*d", pgsz, 7495)
+			found := fmt.Sprintf("%0*d", pgsz, 7500)
+			k, v, _ = c.Seek([]byte(seek))
+			assert.Equal(t, found, string(k))
+			assert.Equal(t, found, string(v))
+
+			// Low key should go to the first key.
+			seek = fmt.Sprintf("%0*d", pgsz, 0)
+			found = fmt.Sprintf("%0*d", pgsz, 10)
+			k, v, _ = c.Seek([]byte(seek))
+			assert.Equal(t, found, string(k))
+			assert.Equal(t, found, string(v))
+
+			// High key should return no key.
+			seek = fmt.Sprintf("%0*d", pgsz, 40000)
+			k, v, _ = c.Seek([]byte(seek))
+			assert.Equal(t, "", string(k))
+			assert.Equal(t, "", string(v))
+
+			return nil
+		})
+	})
+}
+
 // tempfile returns a temporary path.
 func tempfile() string {
 	f, _ := ioutil.TempFile("", "bolt-c-")