add ForEach

tree.go

@@ -239,6 +239,24 @@ func (t *Tree[K, V]) Remove(key K) (*Tree[K, V], bool) {
 	return t, false
 }
 
+// ForEach will call f for each node in this tree, in dfs order. If f returns false, interation stops.
+func (t *Tree[K, V]) ForEach(f func(K, V) bool) {
+	if t == nil {
+		return
+	}
+	t.forEach(f)
+}
+
+func (t *Tree[K, V]) forEach(f func(K, V) bool) bool {
+	if t.left != nil && !t.left.forEach(f) {
+		return false
+	}
+	if !f(t.key, t.value) {
+		return false
+	}
+	return t.right == nil || t.right.forEach(f)
+}
+
 // balance checks if, after Insert or Remove, the tree has become unbalanced.
 // If it has, balance rotates (or double rotates) to fix it, updating childSize as needed.
 // rightHeavy should be true if a new node was inserted into right (or removed from left), else false.

tree_test.go

@@ -159,3 +159,27 @@ func TestTree2Remove(t *testing.T) {
 		t.Fatal("should have nothing left")
 	}
 }
+
+func TestTree2ForEach(t *testing.T) {
+	rand.Seed(time.Now().UnixNano())
+	var test *Tree[cmpUint, cmpUint]
+	for i := 0; i < cmpUintTestSize; i++ {
+		newVal := cmpUint(rand.Uint64())
+		test, _ = test.Insert(newVal, newVal)
+	}
+
+	keys := test.Keys()
+	var diy []cmpUint
+	test.ForEach(func(key, _ cmpUint) bool {
+		diy = append(diy, key)
+		return true
+	})
+	if len(diy) != len(keys) {
+		t.Fatalf("should be same length")
+	}
+	for i, key := range keys {
+		if diy[i] != key {
+			t.Fatalf("should be equal @%d %v %v", i, key, diy[i])
+		}
+	}
+}