From 6c7445772cbe9dd1a6386b3a409c5504f50f5363 Mon Sep 17 00:00:00 2001
From: wdvxdr <wdvxdr1123@gmail.com>
Date: Sun, 30 Jan 2022 15:35:14 +0800
Subject: [PATCH] fix: fix btree iterate

---
 internal/btree/btree.go      |  6 ++++++
 internal/btree/btree_test.go | 25 +++++++++++++++++++++++++
 2 files changed, 31 insertions(+)

diff --git a/internal/btree/btree.go b/internal/btree/btree.go
index 24b5f99..95068f6 100644
--- a/internal/btree/btree.go
+++ b/internal/btree/btree.go
@@ -526,6 +526,7 @@ func (d *DB) Delete(hash *byte) error {
 	return nil
 }
 
+// Foreach iterates over all items in the database file.
 func (d *DB) Foreach(iter func(key [16]byte, value []byte)) {
 	top := d.get(d.top)
 	d.iterate(top, iter)
@@ -542,4 +543,9 @@ func (d *DB) iterate(table *table, iter func(key [16]byte, value []byte)) {
 			d.iterate(child, iter)
 		}
 	}
+	item := table.items[table.size]
+	if item.child != 0 {
+		child := d.get(item.child)
+		d.iterate(child, iter)
+	}
 }
diff --git a/internal/btree/btree_test.go b/internal/btree/btree_test.go
index accec46..a080fbd 100644
--- a/internal/btree/btree_test.go
+++ b/internal/btree/btree_test.go
@@ -5,6 +5,7 @@ import (
 	"os"
 	"testing"
 
+	"github.com/Mrs4s/MiraiGo/utils"
 	assert2 "github.com/stretchr/testify/assert"
 )
 
@@ -65,3 +66,27 @@ func TestBtree(t *testing.T) {
 
 	assert2.NoError(t, bt.Close())
 }
+
+func TestDB_Foreach(t *testing.T) {
+	const elemSize = 100
+	set := make([]string, elemSize)
+	for i := 0; i < elemSize; i++ {
+		set[i] = utils.RandomString(20)
+	}
+	f := tempfile(t)
+	defer os.Remove(f)
+	bt, err := Create(f)
+	defer bt.Close()
+	assert2.NoError(t, err)
+	for _, v := range set {
+		hash := sha1.New()
+		hash.Write([]byte(v))
+		bt.Insert(&hash.Sum(nil)[0], []byte(v))
+	}
+
+	var ss []string
+	bt.Foreach(func(key [16]byte, value []byte) {
+		ss = append(ss, string(value))
+	})
+	assert2.ElementsMatch(t, set, ss)
+}