index: convert AssociatedSet to go iterators

internal/repository/index/associated_data.go

@@ -1,6 +1,8 @@
 package index
 
 import (
+	"iter"
+	"slices"
 	"sort"
 
 	"github.com/restic/restic/internal/restic"
@@ -108,42 +110,48 @@ func (a *AssociatedSet[T]) Delete(bh restic.BlobHandle) {
 
 func (a *AssociatedSet[T]) Len() int {
 	count := 0
-	a.For(func(_ restic.BlobHandle, _ T) {
+	for range a.All() {
 		count++
-	})
+	}
 	return count
 }
 
-func (a *AssociatedSet[T]) For(cb func(bh restic.BlobHandle, val T)) {
-	for k, v := range a.overflow {
-		cb(k, v)
-	}
-
-	for pb := range a.idx.Values() {
-		if _, ok := a.overflow[pb.BlobHandle]; ok {
-			// already reported via overflow set
-			continue
+func (a *AssociatedSet[T]) All() iter.Seq2[restic.BlobHandle, T] {
+	return func(yield func(restic.BlobHandle, T) bool) {
+		for k, v := range a.overflow {
+			if !yield(k, v) {
+				return
+			}
 		}
 
-		val, known := a.Get(pb.BlobHandle)
-		if known {
-			cb(pb.BlobHandle, val)
+		for pb := range a.idx.Values() {
+			if _, ok := a.overflow[pb.BlobHandle]; ok {
+				// already reported via overflow set
+				continue
+			}
+
+			val, known := a.Get(pb.BlobHandle)
+			if known {
+				if !yield(pb.BlobHandle, val) {
+					return
+				}
+			}
 		}
 	}
 }
 
-// List returns a sorted slice of all BlobHandle in the set.
-func (a *AssociatedSet[T]) List() restic.BlobHandles {
-	list := make(restic.BlobHandles, 0)
-	a.For(func(bh restic.BlobHandle, _ T) {
-		list = append(list, bh)
-	})
-
-	return list
+func (a *AssociatedSet[T]) Keys() iter.Seq[restic.BlobHandle] {
+	return func(yield func(restic.BlobHandle) bool) {
+		for bh := range a.All() {
+			if !yield(bh) {
+				return
+			}
+		}
+	}
 }
 
 func (a *AssociatedSet[T]) String() string {
-	list := a.List()
+	list := restic.BlobHandles(slices.Collect(a.Keys()))
 	sort.Sort(list)
 
 	str := list.String()

internal/repository/index/associated_data_test.go

@@ -2,6 +2,7 @@ package index
 
 import (
 	"context"
+	"slices"
 	"testing"
 
 	"github.com/restic/restic/internal/crypto"
@@ -31,6 +32,10 @@ func makeFakePackedBlob() (restic.BlobHandle, restic.PackedBlob) {
 	return bh, blob
 }
 
+func list(bs *AssociatedSet[uint8]) restic.BlobHandles {
+	return restic.BlobHandles(slices.Collect(bs.Keys()))
+}
+
 func TestAssociatedSet(t *testing.T) {
 	bh, blob := makeFakePackedBlob()
 
@@ -40,7 +45,7 @@ func TestAssociatedSet(t *testing.T) {
 
 	bs := NewAssociatedSet[uint8](mi)
 	test.Equals(t, bs.Len(), 0)
-	test.Equals(t, bs.List(), restic.BlobHandles{})
+	test.Equals(t, list(bs), restic.BlobHandles(nil))
 
 	// check non existent
 	test.Equals(t, bs.Has(bh), false)
@@ -51,7 +56,7 @@ func TestAssociatedSet(t *testing.T) {
 	bs.Insert(bh)
 	test.Equals(t, bs.Has(bh), true)
 	test.Equals(t, bs.Len(), 1)
-	test.Equals(t, bs.List(), restic.BlobHandles{bh})
+	test.Equals(t, list(bs), restic.BlobHandles{bh})
 	test.Equals(t, 0, len(bs.overflow))
 
 	// test set
@@ -69,7 +74,7 @@ func TestAssociatedSet(t *testing.T) {
 	bs.Delete(bh)
 	test.Equals(t, bs.Len(), 0)
 	test.Equals(t, bs.Has(bh), false)
-	test.Equals(t, bs.List(), restic.BlobHandles{})
+	test.Equals(t, list(bs), restic.BlobHandles(nil))
 
 	test.Equals(t, "{}", bs.String())
 
@@ -99,7 +104,7 @@ func TestAssociatedSet(t *testing.T) {
 	val, ok = bs.Get(of)
 	test.Equals(t, true, ok)
 	test.Equals(t, uint8(7), val)
-	test.Equals(t, bs.List(), restic.BlobHandles{of, bh})
+	test.Equals(t, list(bs), restic.BlobHandles{of, bh})
 	// update
 	bs.Set(of, 8)
 	val, ok = bs.Get(of)
@@ -110,7 +115,7 @@ func TestAssociatedSet(t *testing.T) {
 	bs.Delete(of)
 	test.Equals(t, bs.Len(), 1)
 	test.Equals(t, bs.Has(of), false)
-	test.Equals(t, bs.List(), restic.BlobHandles{bh})
+	test.Equals(t, list(bs), restic.BlobHandles{bh})
 	test.Equals(t, 0, len(bs.overflow))
 }
 
@@ -138,7 +143,7 @@ func TestAssociatedSetWithExtendedIndex(t *testing.T) {
 	val, ok := bs.Get(of)
 	test.Equals(t, true, ok)
 	test.Equals(t, uint8(5), val)
-	test.Equals(t, bs.List(), restic.BlobHandles{of})
+	test.Equals(t, list(bs), restic.BlobHandles{of})
 	// update
 	bs.Set(of, 8)
 	val, ok = bs.Get(of)
@@ -149,6 +154,6 @@ func TestAssociatedSetWithExtendedIndex(t *testing.T) {
 	bs.Delete(of)
 	test.Equals(t, bs.Len(), 0)
 	test.Equals(t, bs.Has(of), false)
-	test.Equals(t, bs.List(), restic.BlobHandles{})
+	test.Equals(t, list(bs), restic.BlobHandles(nil))
 	test.Equals(t, 0, len(bs.overflow))
 }

internal/repository/prune.go

@@ -176,12 +176,12 @@ func packInfoFromIndex(ctx context.Context, idx restic.ListBlobser, usedBlobs *i
 
 	// Check if all used blobs have been found in index
 	missingBlobs := restic.NewBlobSet()
-	usedBlobs.For(func(bh restic.BlobHandle, count uint8) {
+	for bh, count := range usedBlobs.All() {
 		if count == 0 {
 			// blob does not exist in any pack files
 			missingBlobs.Insert(bh)
 		}
-	})
+	}
 
 	if len(missingBlobs) != 0 {
 		printer.E("%v not found in the index\n\n"+
@@ -311,11 +311,11 @@ func packInfoFromIndex(ctx context.Context, idx restic.ListBlobser, usedBlobs *i
 
 	// Sanity check. If no duplicates exist, all blobs have value 1. After handling
 	// duplicates, this also applies to duplicates.
-	usedBlobs.For(func(_ restic.BlobHandle, count uint8) {
+	for _, count := range usedBlobs.All() {
 		if count != 1 {
 			panic("internal error during blob selection")
 		}
-	})
+	}
 
 	return usedBlobs, indexPack, nil
 }