add btree Ascend、Descend method and unitest.

index/btree.go

@@ -64,3 +64,23 @@ func (mt *MemoryBTree) Delete(key []byte) (*wal.ChunkPosition, bool) {
 func (mt *MemoryBTree) Size() int {
 	return mt.tree.Len()
 }
+
+// Ascend calls f for each item in the tree in ascending order. If f returns false, iteration stops.
+func (mt *MemoryBTree) Ascend(f func(k []byte, v *wal.ChunkPosition) bool) {
+	mt.lock.RLock()
+	defer mt.lock.RUnlock()
+
+	mt.tree.Ascend(func(i btree.Item) bool {
+		return f(i.(*item).key, i.(*item).pos)
+	})
+}
+
+// Descend calls f for each item in the tree in descending order. If f returns false, iteration stops.
+func (mt *MemoryBTree) Descend(f func(k []byte, v *wal.ChunkPosition) bool) {
+	mt.lock.RLock()
+	defer mt.lock.RUnlock()
+
+	mt.tree.Descend(func(i btree.Item) bool {
+		return f(i.(*item).key, i.(*item).pos)
+	})
+}

index/btree_test.go

@@ -0,0 +1,117 @@
+package index
+
+import (
+	"bytes"
+	"testing"
+
+	"github.com/rosedblabs/wal"
+)
+
+func TestMemoryBTree_Put_Get(t *testing.T) {
+	mt := newBTree()
+	w, _ := wal.Open(wal.DefaultOptions)
+
+	key := []byte("testKey")
+	chunkPosition, _ := w.Write([]byte("some data 1"))
+
+	// Test Put
+	oldPos := mt.Put(key, chunkPosition)
+	if oldPos != nil {
+		t.Fatalf("expected nil, got %+v", oldPos)
+	}
+
+	// Test Get
+	gotPos := mt.Get(key)
+	if chunkPosition.ChunkOffset != gotPos.ChunkOffset {
+		t.Fatalf("expected %+v, got %+v", chunkPosition, gotPos)
+	}
+}
+
+func TestMemoryBTree_Delete(t *testing.T) {
+	mt := newBTree()
+	w, _ := wal.Open(wal.DefaultOptions)
+
+	key := []byte("testKey")
+	chunkPosition, _ := w.Write([]byte("some data 2"))
+
+	mt.Put(key, chunkPosition)
+
+	// Test Delete
+	delPos, ok := mt.Delete(key)
+	if !ok {
+		t.Fatal("expected item to be deleted")
+	}
+	if chunkPosition.ChunkOffset != delPos.ChunkOffset {
+		t.Fatalf("expected %+v, got %+v", chunkPosition, delPos)
+	}
+
+	// Ensure the key is deleted
+	if mt.Get(key) != nil {
+		t.Fatal("expected nil, got value")
+	}
+}
+
+func TestMemoryBTree_Size(t *testing.T) {
+	mt := newBTree()
+
+	if mt.Size() != 0 {
+		t.Fatalf("expected size to be 0, got %d", mt.Size())
+	}
+
+	w, _ := wal.Open(wal.DefaultOptions)
+	key := []byte("testKey")
+	chunkPosition, _ := w.Write([]byte("some data 3"))
+
+	mt.Put(key, chunkPosition)
+
+	if mt.Size() != 1 {
+		t.Fatalf("expected size to be 1, got %d", mt.Size())
+	}
+}
+
+func TestMemoryBTree_Ascend_Descend(t *testing.T) {
+	mt := newBTree()
+	w, _ := wal.Open(wal.DefaultOptions)
+
+	data := map[string][]byte{
+		"apple":  []byte("some data 4"),
+		"banana": []byte("some data 5"),
+		"cherry": []byte("some data 6"),
+	}
+
+	positionMap := make(map[string]*wal.ChunkPosition)
+
+	for k, v := range data {
+		chunkPosition, _ := w.Write(v)
+		positionMap[k] = chunkPosition
+		mt.Put([]byte(k), chunkPosition)
+	}
+
+	// Test Ascend
+	prevKey := ""
+	mt.Ascend(func(key []byte, pos *wal.ChunkPosition) bool {
+		if prevKey != "" && bytes.Compare([]byte(prevKey), key) >= 0 {
+			t.Fatalf("items are not in ascending order")
+		}
+		expectedPos := positionMap[string(key)]
+		if expectedPos.ChunkOffset != pos.ChunkOffset {
+			t.Fatalf("expected position %+v, got %+v", expectedPos, pos)
+		}
+		prevKey = string(key)
+		return true
+	})
+
+	// Test Descend
+	prevKey = "zzzzzz"
+	mt.Descend(func(key []byte, pos *wal.ChunkPosition) bool {
+		if bytes.Compare([]byte(prevKey), key) <= 0 {
+			t.Fatalf("items are not in descending order")
+		}
+		expectedPos := positionMap[string(key)]
+		if expectedPos.ChunkOffset != pos.ChunkOffset {
+			t.Fatalf("expected position %+v, got %+v", expectedPos, pos)
+		}
+		prevKey = string(key)
+		return true
+	})
+}