Add index iterator interface with B-tree implementation and tests

Author: XiXi-2024

index/btree.go

@@ -147,3 +147,129 @@ func (mt *MemoryBTree) DescendLessOrEqual(key []byte, handleFn func(key []byte,
 		return cont
 	})
 }
+
+func (mt *MemoryBTree) Iterator(reverse bool) Iterator {
+	if mt.tree == nil {
+		return nil
+	}
+	mt.lock.RLock()
+	defer mt.lock.RUnlock()
+
+	return newMemoryBTreeIterator(mt.tree, reverse)
+}
+
+// memoryBTreeIterator represents a B-tree index iterator implementation
+type memoryBTreeIterator struct {
+	tree    *btree.BTree // underlying B-tree implementation
+	reverse bool         // indicates whether to traverse in descending order
+	current *item        // current element being traversed
+	valid   bool         // indicates if the iterator is valid
+}
+
+func newMemoryBTreeIterator(tree *btree.BTree, reverse bool) *memoryBTreeIterator {
+	var current *item
+	var valid bool
+	if tree.Len() > 0 {
+		if reverse {
+			current = tree.Max().(*item)
+		} else {
+			current = tree.Min().(*item)
+		}
+		valid = true
+	}
+	return &memoryBTreeIterator{
+		tree:    tree.Clone(),
+		reverse: reverse,
+		current: current,
+		valid:   valid,
+	}
+}
+
+func (it *memoryBTreeIterator) Rewind() {
+	if !it.valid {
+		return
+	}
+	if it.reverse {
+		if mx, ok := it.tree.Max().(*item); ok {
+			it.current = mx
+		}
+	} else {
+		if mi, ok := it.tree.Min().(*item); ok {
+			it.current = mi
+		}
+	}
+}
+
+func (it *memoryBTreeIterator) Seek(key []byte) {
+	if !it.valid {
+		return
+	}
+	seekItem := &item{key: key}
+	it.valid = false
+	if it.reverse {
+		it.tree.DescendLessOrEqual(seekItem, func(i btree.Item) bool {
+			it.current = i.(*item)
+			it.valid = true
+			return false
+		})
+	} else {
+		it.tree.AscendGreaterOrEqual(seekItem, func(i btree.Item) bool {
+			it.current = i.(*item)
+			it.valid = true
+			return false
+		})
+	}
+}
+
+func (it *memoryBTreeIterator) Next() {
+	if !it.valid {
+		return
+	}
+	it.valid = false
+	if it.reverse {
+		it.tree.DescendLessOrEqual(it.current, func(i btree.Item) bool {
+			if !i.(*item).Less(it.current) {
+				return true
+			}
+			it.current = i.(*item)
+			it.valid = true
+			return false
+		})
+	} else {
+		it.tree.AscendGreaterOrEqual(it.current, func(i btree.Item) bool {
+			if !it.current.Less(i.(*item)) {
+				return true
+			}
+			it.current = i.(*item)
+			it.valid = true
+			return false
+		})
+	}
+	if !it.valid {
+		it.current = nil
+	}
+}
+
+func (it *memoryBTreeIterator) Valid() bool {
+	return it.valid
+}
+
+func (it *memoryBTreeIterator) Key() []byte {
+	if !it.valid {
+		return nil
+	}
+	return it.current.key
+}
+
+func (it *memoryBTreeIterator) Value() *wal.ChunkPosition {
+	if !it.valid {
+		return nil
+	}
+	return it.current.pos
+}
+
+func (it *memoryBTreeIterator) Close() {
+	it.tree = nil
+	it.current = nil
+	it.valid = false
+}

index/btree_test.go

@@ -3,6 +3,7 @@ package index
 import (
 	"bytes"
 	"fmt"
+	"github.com/stretchr/testify/assert"
 	"testing"
 
 	"github.com/rosedblabs/wal"
@@ -193,3 +194,93 @@ func TestMemoryBTree_AscendGreaterOrEqual_DescendLessOrEqual(t *testing.T) {
 		return true, nil
 	})
 }
+
+func TestMemoryBTree_Iterator(t *testing.T) {
+	mt := newBTree()
+	// Test iterator for empty tree
+	it1 := mt.Iterator(false)
+	assert.Equal(t, false, it1.Valid())
+
+	// Build test data
+	testData := map[string]*wal.ChunkPosition{
+		"acee": {SegmentId: 1, BlockNumber: 2, ChunkOffset: 3, ChunkSize: 100},
+		"bbcd": {SegmentId: 2, BlockNumber: 3, ChunkOffset: 4, ChunkSize: 200},
+		"code": {SegmentId: 3, BlockNumber: 4, ChunkOffset: 5, ChunkSize: 300},
+		"eede": {SegmentId: 4, BlockNumber: 5, ChunkOffset: 6, ChunkSize: 400},
+	}
+
+	// Insert test data
+	for k, v := range testData {
+		mt.Put([]byte(k), v)
+	}
+
+	// Test ascending iteration
+	iter := mt.Iterator(false)
+	var prevKey string
+	count := 0
+	for iter.Rewind(); iter.Valid(); iter.Next() {
+		currKey := string(iter.Key())
+		pos := iter.Value()
+
+		// Verify key order
+		if prevKey != "" {
+			assert.True(t, currKey > prevKey)
+		}
+
+		// Verify value correctness
+		expectedPos := testData[currKey]
+		assert.Equal(t, expectedPos, pos)
+
+		prevKey = currKey
+		count++
+	}
+	assert.Equal(t, len(testData), count)
+
+	// Test descending iteration
+	iter = mt.Iterator(true)
+	prevKey = ""
+	count = 0
+	for iter.Rewind(); iter.Valid(); iter.Next() {
+		currKey := string(iter.Key())
+		pos := iter.Value()
+
+		// Verify key order
+		if prevKey != "" {
+			assert.True(t, currKey < prevKey)
+		}
+
+		// Verify value correctness
+		expectedPos := testData[currKey]
+		assert.Equal(t, expectedPos, pos)
+
+		prevKey = currKey
+		count++
+	}
+	assert.Equal(t, len(testData), count)
+
+	// Test Seek operation
+	testCases := []struct {
+		seekKey    string
+		expectKey  string
+		shouldFind bool
+	}{
+		{"b", "bbcd", true},   // Should find bbcd
+		{"cc", "code", true},  // Should find code
+		{"d", "eede", true},   // Should find eede
+		{"f", "", false},      // Should not find any element
+		{"aaa", "acee", true}, // Should find acee
+	}
+
+	for _, tc := range testCases {
+		iter = mt.Iterator(false)
+		iter.Seek([]byte(tc.seekKey))
+
+		if tc.shouldFind {
+			assert.True(t, iter.Valid())
+			assert.Equal(t, tc.expectKey, string(iter.Key()))
+			assert.Equal(t, testData[tc.expectKey], iter.Value())
+		} else {
+			assert.False(t, iter.Valid())
+		}
+	}
+}

index/index.go

@@ -42,6 +42,9 @@ type Indexer interface {
 	// DescendLessOrEqual iterates in descending order, starting from key <= given key,
 	// invoking handleFn. Stops if handleFn returns false.
 	DescendLessOrEqual(key []byte, handleFn func(key []byte, position *wal.ChunkPosition) (bool, error))
+
+	// Iterator returns an index iterator.
+	Iterator(reverse bool) Iterator
 }
 
 type IndexerType = byte
@@ -61,3 +64,27 @@ func NewIndexer() Indexer {
 		panic("unexpected index type")
 	}
 }
+
+// Iterator represents a generic index iterator interface.
+type Iterator interface {
+	// Rewind resets the iterator to its initial position.
+	Rewind()
+
+	// Seek positions the cursor to the element with the specified key.
+	Seek(key []byte)
+
+	// Next moves the cursor to the next element.
+	Next()
+
+	// Valid checks if the iterator is still valid for reading.
+	Valid() bool
+
+	// Key returns the key of the current element.
+	Key() []byte
+
+	// Value returns the value (chunk position) of the current element.
+	Value() *wal.ChunkPosition
+
+	// Close releases the resources associated with the iterator.
+	Close()
+}