provide loop break mechanism in map.ForEach

alphadose · Sep 6, 2022 · f4af9b5 · f4af9b5
1 parent 0caaec1
commit f4af9b5
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 8 deletions.
diff --git a/README.md b/README.md
@@ -43,8 +43,9 @@ func main() {
 	mep.Set(3, "three")
 
 	// ForEach loop to iterate over all key-value pairs and execute the given lambda
-	mep.ForEach(func(key int, value string) {
+	mep.ForEach(func(key int, value string) bool {
 		fmt.Printf("Key -> %d | Value -> %s\n", key, value)
+		return true // return `true` to continue iteration and `false` to break iteration
 	})
 
 	// delete values

diff --git a/examples/simple.go b/examples/simple.go
@@ -22,8 +22,9 @@ func main() {
 	mep.Set(3, "three")
 
 	// ForEach loop to iterate over all key-value pairs and execute the given lambda
-	mep.ForEach(func(key int, value string) {
+	mep.ForEach(func(key int, value string) bool {
 		fmt.Printf("Key -> %d | Value -> %s\n", key, value)
+		return true
 	})
 
 	// delete values

diff --git a/map.go b/map.go
@@ -199,10 +199,14 @@ func (m *HashMap[K, V]) fillIndexItems(mapData *hashMapData[K, V]) {
 }
 
 // ForEach iterates over key-value pairs and executes the lambda provided for each such pair
-func (m *HashMap[K, V]) ForEach(lambda func(K, V)) {
-	for item := m.listHead; item != nil; item = item.next() {
-		if item.keyHash != marked {
-			lambda(item.key, *item.value.Load())
+// lambda must return `true` to continue iteration and `false` to break iteration
+func (m *HashMap[K, V]) ForEach(lambda func(K, V) bool) {
+	for item := m.listHead.nextPtr.Load(); item != nil; item = item.next() {
+		if item.keyHash == marked {
+			continue
+		}
+		if !lambda(item.key, *item.value.Load()) {
+			return
 		}
 	}
 }

diff --git a/tests/e2e_test.go b/tests/e2e_test.go
@@ -134,8 +134,9 @@ func TestDelete(t *testing.T) {
 func TestIterator(t *testing.T) {
 	m := haxmap.New[int, *Animal]()
 
-	m.ForEach(func(i int, a *Animal) {
+	m.ForEach(func(i int, a *Animal) bool {
 		t.Errorf("map should be empty but got key -> %d and value -> %#v.", i, a)
+		return true
 	})
 
 	itemCount := 16
@@ -144,11 +145,12 @@ func TestIterator(t *testing.T) {
 	}
 
 	counter := 0
-	m.ForEach(func(i int, a *Animal) {
+	m.ForEach(func(i int, a *Animal) bool {
 		if a == nil {
 			t.Error("Expecting an object.")
 		}
 		counter++
+		return true
 	})
 
 	if counter != itemCount {