From b01a6a60e24b19caa37cef4e5450df2e08272e16 Mon Sep 17 00:00:00 2001
From: Pierre Mdawar <pierre@mdawar.dev>
Date: Fri, 23 Aug 2024 18:57:57 +0300
Subject: [PATCH] refactor!: use an iterator function for Entries

This change removes the need for a context to cancel a partial iteration.
---
 README.md        |  4 ++--
 examples_test.go | 21 ++-------------------
 go.mod           |  2 +-
 map.go           | 36 ++++++++----------------------------
 map_test.go      | 23 +++++++++--------------
 5 files changed, 22 insertions(+), 64 deletions(-)

diff --git a/README.md b/README.md
index e600e5a..fdaf722 100644
--- a/README.md
+++ b/README.md
@@ -65,8 +65,8 @@ func main() {
 	removed := m.RemoveExpired() // Returns the number of removed keys.
 
 	// Iterate over the map entries.
-	for entry := range m.Entries(context.TODO()) {
-		fmt.Println("Key:", entry.Key, "-", "Value:", entry.Value)
+	for key, value := range m.Entries() {
+		fmt.Println("Key:", key, "-", "Value:", value)
 	}
 }
 ```
diff --git a/examples_test.go b/examples_test.go
index ac445f7..446189e 100644
--- a/examples_test.go
+++ b/examples_test.go
@@ -1,7 +1,6 @@
 package xmap_test
 
 import (
-	"context"
 	"fmt"
 	"time"
 
@@ -69,23 +68,7 @@ func ExampleMap_Entries() {
 	m := xmap.New[string, int]()
 	defer m.Stop()
 
-	for entry := range m.Entries(context.TODO()) {
-		fmt.Println("Key:", entry.Key, "-", "Value:", entry.Value)
-	}
-}
-
-func ExampleMap_Entries_partial_iteration() {
-	m := xmap.New[string, int]()
-	defer m.Stop()
-
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
-
-	for entry := range m.Entries(ctx) {
-		fmt.Println("Key:", entry.Key, "-", "Value:", entry.Value)
-		// With a partial iteration, the context must be canceled
-		// to prevent a deadlock (A read lock is held during the iteration).
-		cancel()
-		break
+	for k, v := range m.Entries() {
+		fmt.Println("Key:", k, "-", "Value:", v)
 	}
 }
diff --git a/go.mod b/go.mod
index 381ce5c..0b45ada 100644
--- a/go.mod
+++ b/go.mod
@@ -1,5 +1,5 @@
 module github.com/mdawar/xmap
 
-go 1.21
+go 1.23
 
 require go.uber.org/goleak v1.3.0
diff --git a/map.go b/map.go
index 6d25faf..43dc9b2 100644
--- a/map.go
+++ b/map.go
@@ -2,7 +2,7 @@
 package xmap
 
 import (
-	"context"
+	"iter"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -169,44 +169,24 @@ func (m *Map[K, V]) GetWithExpiration(key K) (V, time.Time, bool) {
 	return zero, time.Time{}, false
 }
 
-// Entry represents a key/value pair in the Map.
-type Entry[K comparable, V any] struct {
-	Key   K
-	Value V
-}
-
-// Entries returns a read-only channel of Entry elements representing the
-// current entries in the Map.
-//
-// This channel can be used in a for range loop to iterate over the current
-// map entries. Only the elements that have not expired are sent on this
-// channel. The channel is closed after all the entries have been sent.
+// Entries returns an iterator over key-value pairs of the Map entries.
 //
-// Like the map type, the iteration order is not guaranteed.
+// Only the entries that have not expired are produced during the iteration.
 //
-// A read lock is held during the iteration, so it's important to consume
-// all of the elements sent on the channel or cancel the passed context
-// if a full iteration is not needed.
-func (m *Map[K, V]) Entries(ctx context.Context) <-chan Entry[K, V] {
-	ch := make(chan Entry[K, V])
-
-	go func() {
+// Similar to the map type, the iteration order is not guaranteed.
+func (m *Map[K, V]) Entries() iter.Seq2[K, V] {
+	return func(yield func(K, V) bool) {
 		m.mu.RLock()
 		defer m.mu.RUnlock()
-		defer close(ch)
 
 		for key, entry := range m.kv {
 			if !m.expired(entry) {
-				select {
-				case <-ctx.Done():
+				if !yield(key, entry.value) {
 					return
-				case ch <- Entry[K, V]{key, entry.value}:
 				}
 			}
 		}
-	}()
-
-	return ch
+	}
 }
 
 // Delete removes a key from the map.
diff --git a/map_test.go b/map_test.go
index a30e7fb..c650470 100644
--- a/map_test.go
+++ b/map_test.go
@@ -1,7 +1,6 @@
 package xmap_test
 
 import (
-	"context"
 	"maps"
 	"testing"
 	"time"
@@ -653,8 +652,8 @@ func TestMapIterateOverMapEntries(t *testing.T) {
 		t.Helper()
 
 		gotEntries := make(map[string]int)
-		for entry := range m.Entries(context.Background()) {
-			gotEntries[entry.Key] = entry.Value
+		for k, v := range m.Entries() {
+			gotEntries[k] = v
 		}
 
 		if !maps.Equal(wantEntries, gotEntries) {
@@ -708,21 +707,17 @@ func TestMapPartialIterationOverEntries(t *testing.T) {
 		m.Set(entry.key, entry.value, entry.ttl)
 	}
 
-	// Number of entries consumed.
-	var gotCount int
+	// Consumed entries.
+	consumed := make(map[string]int)
 
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
-
-	for range m.Entries(ctx) {
-		gotCount++
-		cancel() // Must cancel the context to release the lock.
-		break    // Stop after consuming 1 entry.
+	for k, v := range m.Entries() {
+		consumed[k] = v
+		break // Stop after consuming 1 entry.
 	}
 
 	// Make sure we consume at least 1 entry.
-	if gotCount != 1 {
-		t.Errorf("want to consume 1 entry, got %d", gotCount)
+	if len(consumed) != 1 {
+		t.Errorf("want to consume 1 entry, got %d", len(consumed))
 	}
 
 	// Channel used to wait for stopping the map.