From 93f447163daea106d3f934f48ea4a9f2b4730d5e Mon Sep 17 00:00:00 2001
From: Ophestra <cat@gensokyo.uk>
Date: Mon, 14 Jul 2025 19:00:11 +0900
Subject: [PATCH] instantiated: implement decoding iterator

---
 instantiated.go      | 87 +++++++++++++++++++++++++++++---------------
 instantiated_test.go | 23 ++++++++++++
 2 files changed, 81 insertions(+), 29 deletions(-)

diff --git a/instantiated.go b/instantiated.go
index 8c17e3e..3014eb7 100644
--- a/instantiated.go
+++ b/instantiated.go
@@ -5,6 +5,7 @@ import (
 	"context"
 	"errors"
 	"io"
+	"iter"
 	"path"
 	"slices"
 	"strings"
@@ -54,46 +55,74 @@ func (m *MalformedInstantiatedError) Is(err error) bool {
 	return e.Type == m.Type
 }
 
+// InstantiatedDecoder interprets the verbose output of `nix build` and collects instantiated derivations.
+type InstantiatedDecoder struct {
+	err     error
+	scanner *bufio.Scanner
+}
+
+// Err returns the first error encountered by the InstantiatedDecoder.
+func (d *InstantiatedDecoder) Err() error { return errors.Join(d.err, d.scanner.Err()) }
+
+// NewInstantiatedDecoder returns a [InstantiatedDecoder] struct for collecting derivations from r.
+func NewInstantiatedDecoder(r io.Reader) *InstantiatedDecoder {
+	return &InstantiatedDecoder{scanner: bufio.NewScanner(r)}
+}
+
+// Drv returns a non-reusable iterator over instantiated derivations collected from [io.Reader].
+func (d *InstantiatedDecoder) Drv() iter.Seq[string] {
+	return func(yield func(string) bool) {
+		for d.scanner.Scan() {
+			v := d.scanner.Text()
+			if !strings.HasPrefix(v, instantiatedPrefix) {
+				// write skipped lines to stderr
+				if Stderr != nil {
+					_, _ = Stderr.Write([]byte(v + "\n"))
+				}
+				continue
+			}
+
+			fields := strings.SplitN(v, instantiatedSeparator, instantiatedFields)
+			if len(fields) != instantiatedFields {
+				// no more than a single -> is expected
+				d.err = &MalformedInstantiatedError{v, InstantiatedBadFields}
+				break
+			}
+
+			// very basic validation here: the output format is not fully understood
+			if len(fields[1]) < 3 || fields[1][0] != '\'' || fields[1][len(fields[1])-1] != '\'' {
+				d.err = &MalformedInstantiatedError{v, InstantiatedUnexpectedQuotes}
+				break
+			}
+			drv := fields[1][1 : len(fields[1])-1]
+
+			if !path.IsAbs(drv) {
+				d.err = &MalformedInstantiatedError{v, InstantiatedNotAbsolute}
+				break
+			}
+
+			if !yield(drv) {
+				break
+			}
+		}
+	}
+}
+
 // DecodeInstantiated interprets the verbose output of `nix build` and returns a slice of all instantiated derivations.
-func DecodeInstantiated(stderr io.Reader) ([]string, error) {
-	scanner := bufio.NewScanner(stderr)
+func DecodeInstantiated(r io.Reader) ([]string, error) {
+	decoder := NewInstantiatedDecoder(r)
 
 	instantiated := make([]string, 0, instantiatedInitialCap)
-	for scanner.Scan() {
-		// on error, process is killed by deferred context cancel
-
+	for drv := range decoder.Drv() {
 		if len(instantiated) == cap(instantiated) {
 			// grow the buffer exponentially to minimise copies
 			instantiated = slices.Grow(instantiated, cap(instantiated)<<1)
 		}
-
-		v := scanner.Text()
-		if !strings.HasPrefix(v, instantiatedPrefix) {
-			if Stderr != nil {
-				_, _ = Stderr.Write([]byte(v + "\n"))
-			}
-			continue
-		}
-		f := strings.SplitN(v, instantiatedSeparator, instantiatedFields)
-		if len(f) != instantiatedFields {
-			return nil, &MalformedInstantiatedError{v, InstantiatedBadFields}
-		}
-
-		// very basic validation here: the output format is not fully understood
-		if len(f[1]) < 3 || f[1][0] != '\'' || f[1][len(f[1])-1] != '\'' {
-			return nil, &MalformedInstantiatedError{v, InstantiatedUnexpectedQuotes}
-		}
-		drv := f[1][1 : len(f[1])-1]
-
-		if !path.IsAbs(drv) {
-			return nil, &MalformedInstantiatedError{v, InstantiatedNotAbsolute}
-		}
-
 		instantiated = append(instantiated, drv)
 	}
 
 	slices.Sort(instantiated)
-	return slices.Compact(instantiated), scanner.Err()
+	return slices.Compact(instantiated), decoder.Err()
 
 }
 
diff --git a/instantiated_test.go b/instantiated_test.go
index 25c7955..accc054 100644
--- a/instantiated_test.go
+++ b/instantiated_test.go
@@ -64,6 +64,29 @@ func TestDecodeInstantiated(t *testing.T) {
 		})
 	}
 
+	t.Run("stop early", func(t *testing.T) {
+		want := []string{
+			"/nix/store/gyks6vvl7x0gq214ldjhi3w4rg37nh8i-zlib-1.3.1.tar.gz.drv",
+			"/nix/store/bamwxswxacs3cjdcydv0z7bj22d7g2kc-config.guess-948ae97.drv",
+			"/nix/store/nbsdqpfzh1jlpmh95s69b3iivfcvv3lh-config.sub-948ae97.drv",
+		}
+
+		decoder := nixbuild.NewInstantiatedDecoder(strings.NewReader(segmentPrefix + segmentBody + segmentSuffix))
+		counter := 3
+		got := make([]string, 0, counter)
+		for drv := range decoder.Drv() {
+			got = append(got, drv)
+			counter--
+			if counter == 0 {
+				break
+			}
+		}
+
+		if !slices.Equal(got, want) {
+			t.Errorf("Drv: %#v, want %#v", got, want)
+		}
+	})
+
 	t.Run("errors", func(t *testing.T) {
 		t.Run("bad type", func(t *testing.T) {
 			badErr := errors.New("")