Add various new standard library extensions (#457)

This PR aggregates new additions to the `exts` directory form various
in-progress PRs, aiming to try to get those changes out of in-progress
PRs and into their own, easier-to-review PR.

Author: mcy

experimental/ast/path.go

@@ -22,7 +22,6 @@ import (
 	"github.com/bufbuild/protocompile/experimental/report"
 	"github.com/bufbuild/protocompile/experimental/token"
 	"github.com/bufbuild/protocompile/internal/ext/iterx"
-	"github.com/bufbuild/protocompile/internal/ext/slicesx"
 )
 
 // Path represents a multi-part identifier.
@@ -70,12 +69,11 @@ func (p Path) ToRelative() Path {
 // AsIdent returns the single identifier that comprises this path, or
 // the zero token.
 func (p Path) AsIdent() token.Token {
-	var buf [2]PathComponent
-	prefix := slicesx.AppendSeq(buf[:0], iterx.Limit(2, p.Components))
-	if len(prefix) != 1 || !prefix[0].Separator().IsZero() {
+	first, _ := iterx.OnlyOne(p.Components)
+	if !first.Separator().IsZero() {
 		return token.Zero
 	}
-	return prefix[0].AsIdent()
+	return first.AsIdent()
 }
 
 // AsPredeclared returns the [predeclared.Name] that this path represents.

experimental/report/diff.go

@@ -87,24 +87,19 @@ func unifiedDiff(span Span, edits []Edit) (Span, []hunk) {
 
 	// Partition offsets into overlapping lines. That is, this connects together
 	// all edit spans whose end and start are not separated by a newline.
-	prev := 0
-	parts := slicesx.SplitFunc(edits, func(i int, next Edit) bool {
-		if i == prev {
-			return false
-		}
-
-		chunk := src[edits[i-1].End:next.Start]
-		if !strings.Contains(chunk, "\n") {
-			return false
-		}
-
-		prev = i
-		return true
+	parts := slicesx.SplitAfterFunc(edits, func(i int, edit Edit) bool {
+		next, ok := slicesx.Get(edits, i+1)
+		return ok && edit.End < next.Start && // Go treats str[x:y] for x > y as an error.
+			strings.Contains(src[edit.End:next.Start], "\n")
 	})
 
 	var out []hunk
 	var prevHunk int
 	parts(func(edits []Edit) bool {
+		if len(edits) == 0 {
+			return true
+		}
+
 		// First, figure out the start and end of the modified region.
 		start, end := edits[0].Start, edits[0].End
 		for _, edit := range edits[1:] {

internal/ext/iterx/consume.go

@@ -0,0 +1,63 @@
+// Copyright 2020-2025 Buf Technologies, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Package iterx contains extensions to Go's package iter.
+package iterx
+
+import (
+	"fmt"
+	"strings"
+
+	"github.com/bufbuild/protocompile/internal/iter"
+)
+
+// Count counts the number of elements in seq that match the given predicate.
+//
+// If p is nil, it is treated as func(_ T) bool { return true }.
+func Count[T any](seq iter.Seq[T]) int {
+	var total int
+	seq(func(_ T) bool {
+		total++
+		return true
+	})
+	return total
+}
+
+// Join is like [strings.Join], but works on an iterator. Elements are
+// stringified as if by [fmt.Print].
+func Join[T any](seq iter.Seq[T], sep string) string {
+	var out strings.Builder
+	first := true
+	seq(func(v T) bool {
+		if !first {
+			out.WriteString(sep)
+		}
+		first = false
+
+		fmt.Fprint(&out, v)
+		return true
+	})
+	return out.String()
+}
+
+// Every returns whether every element of an iterator satisfies the given
+// predicate. Returns true if seq yields no values.
+func Every[T any](seq iter.Seq[T], p func(T) bool) bool {
+	all := true
+	seq(func(v T) bool {
+		all = p(v)
+		return all
+	})
+	return all
+}

internal/ext/iterx/filtermap.go

@@ -0,0 +1,104 @@
+// Copyright 2020-2025 Buf Technologies, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package iterx
+
+import (
+	"github.com/bufbuild/protocompile/internal/iter"
+)
+
+// This file contains the matrix of {Map, Filter, FilterMap} x {1, 2, 1to2, 2to1},
+// except that Filter1to2 and Filter2to1 don't really make sense.
+
+// Map returns a new iterator applying f to each element of seq.
+func Map[T, U any](seq iter.Seq[T], f func(T) U) iter.Seq[U] {
+	return FilterMap(seq, func(v T) (U, bool) { return f(v), true })
+}
+
+// Filter returns a new iterator that only includes values satisfying p.
+func Filter[T any](seq iter.Seq[T], p func(T) bool) iter.Seq[T] {
+	return FilterMap(seq, func(v T) (T, bool) { return v, p(v) })
+}
+
+// FilterMap combines the operations of [Map] and [Filter].
+func FilterMap[T, U any](seq iter.Seq[T], f func(T) (U, bool)) iter.Seq[U] {
+	return func(yield func(U) bool) {
+		seq(func(v T) bool {
+			v2, ok := f(v)
+			return !ok || yield(v2)
+		})
+	}
+}
+
+// Map2 returns a new iterator applying f to each element of seq.
+func Map2[T, U, V, W any](seq iter.Seq2[T, U], f func(T, U) (V, W)) iter.Seq2[V, W] {
+	return FilterMap2(seq, func(v1 T, v2 U) (V, W, bool) {
+		x1, x2 := f(v1, v2)
+		return x1, x2, true
+	})
+}
+
+// Filter2 returns a new iterator that only includes values satisfying p.
+func Filter2[T, U any](seq iter.Seq2[T, U], p func(T, U) bool) iter.Seq2[T, U] {
+	return FilterMap2(seq, func(v1 T, v2 U) (T, U, bool) { return v1, v2, p(v1, v2) })
+}
+
+// FilterMap2 combines the operations of [Map] and [Filter].
+func FilterMap2[T, U, V, W any](seq iter.Seq2[T, U], f func(T, U) (V, W, bool)) iter.Seq2[V, W] {
+	return func(yield func(V, W) bool) {
+		seq(func(v1 T, v2 U) bool {
+			x1, x2, ok := f(v1, v2)
+			return !ok || yield(x1, x2)
+		})
+	}
+}
+
+// Map2to1 is like [Map], but it also acts a Y pipe for converting a two-element
+// iterator into a one-element iterator.
+func Map2to1[T, U, V any](seq iter.Seq2[T, U], f func(T, U) V) iter.Seq[V] {
+	return FilterMap2to1(seq, func(v1 T, v2 U) (V, bool) {
+		return f(v1, v2), true
+	})
+}
+
+// FilterMap2to1 is like [FilterMap], but it also acts a Y pipe for converting
+// a two-element iterator into a one-element iterator.
+func FilterMap2to1[T, U, V any](seq iter.Seq2[T, U], f func(T, U) (V, bool)) iter.Seq[V] {
+	return func(yield func(V) bool) {
+		seq(func(v1 T, v2 U) bool {
+			v, ok := f(v1, v2)
+			return !ok || yield(v)
+		})
+	}
+}
+
+// Map1To2 is like [Map], but it also acts a Y pipe for converting a one-element
+// iterator into a two-element iterator.
+func Map1To2[T, U, V any](seq iter.Seq[T], f func(T) (U, V)) iter.Seq2[U, V] {
+	return FilterMap1To2(seq, func(v T) (U, V, bool) {
+		x1, x2 := f(v)
+		return x1, x2, true
+	})
+}
+
+// FilterMap1To2 is like [FilterMap], but it also acts a Y pipe for converting
+// a one-element iterator into a two-element iterator.
+func FilterMap1To2[T, U, V any](seq iter.Seq[T], f func(T) (U, V, bool)) iter.Seq2[U, V] {
+	return func(yield func(U, V) bool) {
+		seq(func(v T) bool {
+			x1, x2, ok := f(v)
+			return !ok || yield(x1, x2)
+		})
+	}
+}

internal/ext/iterx/get.go

@@ -0,0 +1,105 @@
+// Copyright 2020-2025 Buf Technologies, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package iterx
+
+import (
+	"github.com/bufbuild/protocompile/internal/iter"
+)
+
+// First retrieves the first element of an iterator.
+func First[T any](seq iter.Seq[T]) (v T, ok bool) {
+	seq(func(x T) bool {
+		v = x
+		ok = true
+		return false
+	})
+	return v, ok
+}
+
+// OnlyOne retrieves the only element of an iterator.
+func OnlyOne[T any](seq iter.Seq[T]) (v T, ok bool) {
+	var found T
+	seq(func(x T) bool {
+		if !ok {
+			found = x
+		}
+		ok = !ok
+		return ok
+	})
+	if ok {
+		// Ensure we return the zero value if there is more
+		// than one element.
+		v = found
+	}
+	return v, ok
+}
+
+// Find returns the first element that matches a predicate.
+//
+// Returns the value and the index at which it was found, or -1 if it wasn't
+// found.
+func Find[T any](seq iter.Seq[T], p func(T) bool) (int, T) {
+	var v T
+	var idx int
+	var found bool
+	seq(func(x T) bool {
+		if p(x) {
+			v = x
+			found = true
+			return false
+		}
+		idx++
+		return true
+	})
+	if !found {
+		idx = -1
+	}
+	return idx, v
+}
+
+// Find2 is like [Find] but for two-element iterators.
+func Find2[T, U any](seq iter.Seq2[T, U], p func(T, U) bool) (int, T, U) {
+	var v1 T
+	var v2 U
+	var idx int
+	var found bool
+	seq(func(x1 T, x2 U) bool {
+		if p(x1, x2) {
+			v1, v2 = x1, x2
+			found = true
+			return false
+		}
+		idx++
+		return true
+	})
+	if !found {
+		idx = -1
+	}
+	return idx, v1, v2
+}
+
+// Index returns the index of the first element of seq that satisfies p.
+//
+// if not found, returns -1.
+func Index[T any](seq iter.Seq[T], p func(T) bool) int {
+	idx, _ := Find(seq, p)
+	return idx
+}
+
+// Index2 is like [Index], but for two-element iterators.
+func Index2[T, U any](seq iter.Seq2[T, U], p func(T, U) bool) int {
+	idx, _, _ := Find2(seq, p)
+	return idx
+}