Add IterateKV (#265)

z/btree.go

@@ -308,6 +308,32 @@ func (t *Tree) Iterate(fn func(node)) {
 	t.iterate(root, fn)
 }
 
+// IterateKV iterates through all keys and values in the tree.
+// If newVal is non-zero, it will be set in the tree.
+func (t *Tree) IterateKV(f func(key, val uint64) (newVal uint64)) {
+	t.Iterate(func(n node) {
+		// Only leaf nodes contain keys.
+		if !n.isLeaf() {
+			return
+		}
+
+		for i := 0; i < n.numKeys(); i++ {
+			key := n.key(i)
+			val := n.val(i)
+
+			// A zero value here means that this is a bogus entry.
+			if val == 0 {
+				continue
+			}
+
+			newVal := f(key, val)
+			if newVal != 0 {
+				n.setAt(valOffset(i), newVal)
+			}
+		}
+	})
+}
+
 func (t *Tree) print(n node, parentID uint64) {
 	n.print(parentID)
 	if n.isLeaf() {

z/btree_test.go

@@ -134,6 +134,35 @@ func TestTreeCycle(t *testing.T) {
 	require.GreaterOrEqual(t, stats.NumPagesFree, int(float64(stats.NumPages)*0.95))
 }
 
+func TestTreeIterateKV(t *testing.T) {
+	bt := NewTree()
+
+	// Set entries: (i, i*10)
+	const n = uint64(1 << 20)
+	for i := uint64(1); i <= n; i++ {
+		bt.Set(i, i*10)
+	}
+
+	// Validate entries: (i, i*10)
+	// Set entries: (i, i*20)
+	count := uint64(0)
+	bt.IterateKV(func(k, v uint64) uint64 {
+		require.Equal(t, k*10, v)
+		count++
+		return k * 20
+	})
+	require.Equal(t, n, count)
+
+	// Validate entries: (i, i*20)
+	count = uint64(0)
+	bt.IterateKV(func(k, v uint64) uint64 {
+		require.Equal(t, k*20, v)
+		count++
+		return 0
+	})
+	require.Equal(t, n, count)
+}
+
 func TestOccupancyRatio(t *testing.T) {
 	// atmax 4 keys per node
 	setPageSize(16 * 5)

z/flags_test.go

@@ -43,7 +43,6 @@ func TestFlagDefault(t *testing.T) {
 }
 
 func TestGetPath(t *testing.T) {
-
 	usr, err := user.Current()
 	require.NoError(t, err)
 	homeDir := usr.HomeDir