From e3d53acb29ba08bb5a4a3fdc20bf0e2e7d64cf41 Mon Sep 17 00:00:00 2001
From: Tom Godkin <tomgodkin@pm.me>
Date: Mon, 9 Sep 2024 21:56:21 +0100
Subject: [PATCH] Add MustCollect consumer

---
 README.md       | 20 +++++++++++++++++---
 it/iter.go      | 17 +++++++++++++++++
 it/iter_test.go | 26 ++++++++++++++++++++++++++
 3 files changed, 60 insertions(+), 3 deletions(-)

diff --git a/README.md b/README.md
index 45c1209..f4ef7ed 100644
--- a/README.md
+++ b/README.md
@@ -76,7 +76,7 @@ keys, values := it.Collect2(maps.All(map[string]int{"one": 1, "two": 2}))
 keys, values := itx.FromMap(map[string]int{"one": 1, "two": 2}).Collect()
 ```
 
-### TryCollect
+<h3 id="trycollect">TryCollect & MustCollect</h3>
 
 Dealing with iterators that return `T, error` can involve the boilerplate of checking that the
 returned slice of errors only contains `nil`. `TryCollect` solves this by collecting all values into
@@ -90,6 +90,19 @@ if lines, err := it.TryCollect(it.LinesString(text)); err != nil {
 }
 ```
 
+MustCollect is similar except that if an error is encountered then a panic will occur.
+
+```go
+text := strings.NewReader("one\ntwo\nthree\n")
+
+lines := it.MustCollect(it.LinesString(text))
+```
+
+<!-- prettier-ignore -->
+> [!TIP]
+> Use `MustCollect` when you can guarantee that no error will occur (such as with
+> [strings.Reader](https://pkg.go.dev/strings#Reader)).
+
 <!-- prettier-ignore -->
 > [!NOTE]
 > If an error is encountered, collection stops. This means the iterator being collected may not be
@@ -97,7 +110,8 @@ if lines, err := it.TryCollect(it.LinesString(text)); err != nil {
 
 <!-- prettier-ignore -->
 > [!NOTE]
-> The `itx` package does not contain `TryCollect` due to limitations with Go's type system.
+> The `itx` package does not contain `TryCollect` or `MustCollect` due to limitations with Go's type
+> system.
 
 ### ForEach
 
@@ -316,7 +330,7 @@ itx.From2(it.Map2(slices.All([]int{1, 2, 3}), printValue2)).Drain()
 
 <!-- prettier-ignore -->
 > [!TIP]
-> Use Drain to consume an iterator to invoke any side effects when you don't need to collect the
+> Use `Drain` to consume an iterator to invoke any side effects when you don't need to collect the
 > values.
 
 ## Iterators
diff --git a/it/iter.go b/it/iter.go
index 74ae142..df94e72 100644
--- a/it/iter.go
+++ b/it/iter.go
@@ -2,6 +2,7 @@ package it
 
 import (
 	"cmp"
+	"fmt"
 	"iter"
 )
 
@@ -142,6 +143,22 @@ func TryCollect[V any](iterator func(func(V, error) bool)) ([]V, error) {
 	return values, nil
 }
 
+// MustCollect consumes an [iter.Seq2] where the right side yields errors and
+// returns a slice of values. If an error is encountered this function will
+// panic.
+func MustCollect[V any](iterator func(func(V, error) bool)) []V {
+	var values []V
+
+	for v, err := range iterator {
+		if err != nil {
+			panic(fmt.Sprintf("it: MustCollect: error yielded by iterator: %s", err.Error()))
+		}
+		values = append(values, v)
+	}
+
+	return values
+}
+
 // Collect2 consumes an [iter.Seq2] and returns two slices of values.
 func Collect2[V, W any](iterator func(func(V, W) bool)) ([]V, []W) {
 	var (
diff --git a/it/iter_test.go b/it/iter_test.go
index 79c54b5..b10109b 100644
--- a/it/iter_test.go
+++ b/it/iter_test.go
@@ -225,3 +225,29 @@ func ExampleDrain2() {
 	// 2
 	// 3
 }
+
+func ExampleMustCollect() {
+	buffer := strings.NewReader("one\ntwo")
+	lines := it.MustCollect(it.LinesString(buffer))
+
+	fmt.Println(lines)
+	// Output: [one two]
+}
+
+func TestMustCollectPanic(t *testing.T) {
+	t.Parallel()
+
+	defer func() {
+		r := recover()
+
+		if r == nil {
+			t.Errorf("expected panic")
+		}
+
+		if fmt.Sprint(r) != "it: MustCollect: error yielded by iterator: read error" {
+			t.Errorf("wrong panic message")
+		}
+	}()
+
+	it.MustCollect(it.LinesString(new(failSecondTime)))
+}