more idiomatic iterator use

cmd/chroma/main.go

@@ -394,9 +394,6 @@ func format(ctx *kong.Context, w io.Writer, style *chroma.Style, it chroma.Itera
 func check(filename string, it chroma.Iterator) {
 	line, col := 1, 0
 	for token := range it {
-		if token == chroma.EOF {
-			break
-		}
 		if token.Type == chroma.Error {
 			fmt.Printf("%s:%d:%d %q\n", filename, line, col, token.String())
 		}

coalesce.go

@@ -13,13 +13,10 @@ func (d *coalescer) Tokenise(options *TokeniseOptions, text string) (Iterator, e
 	return func(yield func(Token) bool) {
 		var prev Token
 		for token := range it {
-			if token == EOF {
-				break
-			}
 			if len(token.Value) == 0 {
 				continue
 			}
-			if prev == EOF {
+			if prev.IsZero() {
 				prev = token
 			} else {
 				if prev.Type == token.Type && len(prev.Value) < 8192 {
@@ -32,7 +29,7 @@ func (d *coalescer) Tokenise(options *TokeniseOptions, text string) (Iterator, e
 				}
 			}
 		}
-		if prev != EOF {
+		if !prev.IsZero() {
 			yield(prev)
 		}
 	}, nil

delegate.go

@@ -71,12 +71,12 @@ func (d *delegatingLexer) Tokenise(options *TokeniseOptions, text string) (Itera
 	var last Token
 	for _, t := range tokens {
 		if t.Type == Other {
-			if last != EOF && insert != nil && last.Type != Other {
+			if !last.IsZero() && insert != nil && last.Type != Other {
 				insert.end = offset
 			}
 			others.WriteString(t.Value)
 		} else {
-			if last == EOF || last.Type == Other {
+			if last.IsZero() || last.Type == Other {
 				insert = &insertion{start: offset}
 				insertions = append(insertions, insert)
 			}
@@ -97,64 +97,71 @@ func (d *delegatingLexer) Tokenise(options *TokeniseOptions, text string) (Itera
 	}
 
 	// Interleave the two sets of tokens.
-	var out []Token
-	offset = 0 // Offset into text.
-	tokenIndex := 0
-	nextToken := func() Token {
-		if tokenIndex >= len(rootTokens) {
-			return EOF
-		}
-		t := rootTokens[tokenIndex]
-		tokenIndex++
-		return t
-	}
-	insertionIndex := 0
-	nextInsertion := func() *insertion {
-		if insertionIndex >= len(insertions) {
-			return nil
+	return func(yield func(Token) bool) {
+		offset := 0 // Offset into text.
+		tokenIndex := 0
+		nextToken := func() Token {
+			if tokenIndex >= len(rootTokens) {
+				return Token{}
+			}
+			t := rootTokens[tokenIndex]
+			tokenIndex++
+			return t
 		}
-		i := insertions[insertionIndex]
-		insertionIndex++
-		return i
-	}
-	t := nextToken()
-	i := nextInsertion()
-	for t != EOF || i != nil {
-		// fmt.Printf("%d->%d:%q   %d->%d:%q\n", offset, offset+len(t.Value), t.Value, i.start, i.end, Stringify(i.tokens...))
-		if t == EOF || (i != nil && i.start < offset+len(t.Value)) {
-			var l Token
-			l, t = splitToken(t, i.start-offset)
-			if l != EOF {
-				out = append(out, l)
-				offset += len(l.Value)
+		insertionIndex := 0
+		nextInsertion := func() *insertion {
+			if insertionIndex >= len(insertions) {
+				return nil
 			}
-			out = append(out, i.tokens...)
-			offset += i.end - i.start
-			if t == EOF {
+			i := insertions[insertionIndex]
+			insertionIndex++
+			return i
+		}
+		t := nextToken()
+		i := nextInsertion()
+		for !t.IsZero() || i != nil {
+			// fmt.Printf("%d->%d:%q   %d->%d:%q\n", offset, offset+len(t.Value), t.Value, i.start, i.end, Stringify(i.tokens...))
+			if t.IsZero() || (i != nil && i.start < offset+len(t.Value)) {
+				var l Token
+				l, t = splitToken(t, i.start-offset)
+				if !l.IsZero() {
+					if !yield(l) {
+						return
+					}
+					offset += len(l.Value)
+				}
+				for _, tok := range i.tokens {
+					if !yield(tok) {
+						return
+					}
+				}
+				offset += i.end - i.start
+				if t.IsZero() {
+					t = nextToken()
+				}
+				i = nextInsertion()
+			} else {
+				if !yield(t) {
+					return
+				}
+				offset += len(t.Value)
 				t = nextToken()
 			}
-			i = nextInsertion()
-		} else {
-			out = append(out, t)
-			offset += len(t.Value)
-			t = nextToken()
 		}
-	}
-	return Literator(out...), nil
+	}, nil
 }
 
 func splitToken(t Token, offset int) (l Token, r Token) {
-	if t == EOF {
-		return EOF, EOF
+	if t.IsZero() {
+		return t, t
 	}
 	if offset == 0 {
-		return EOF, t
+		return Token{}, t
 	}
 	if offset == len(t.Value) {
-		return t, EOF
+		return t, Token{}
 	}
-	l = t.Clone()
-	r = t.Clone()
+	l, r = t, t
 	l.Value = l.Value[:offset]
 	r.Value = r.Value[offset:]
 	return

delegate_test.go

@@ -1,6 +1,7 @@
 package chroma
 
 import (
+	"slices"
 	"testing"
 
 	assert "github.com/alecthomas/assert/v2"
@@ -104,7 +105,7 @@ func TestDelegate(t *testing.T) {
 		t.Run(test.name, func(t *testing.T) {
 			it, err := delegate.Tokenise(nil, test.source)
 			assert.NoError(t, err)
-			actual := it.Tokens()
+			actual := slices.Collect(it)
 			assert.Equal(t, test.expected, actual)
 		})
 	}

formatters/api.go

@@ -13,9 +13,6 @@ var (
 	// NoOp formatter.
 	NoOp = Register("noop", chroma.FormatterFunc(func(w io.Writer, s *chroma.Style, iterator chroma.Iterator) error {
 		for t := range iterator {
-			if t == chroma.EOF {
-				break
-			}
 			if _, err := io.WriteString(w, t.Value); err != nil {
 				return err
 			}

formatters/html/html.go

@@ -4,6 +4,8 @@ import (
 	"fmt"
 	"html"
 	"io"
+	"iter"
+	"slices"
 	"sort"
 	"strconv"
 	"strings"
@@ -222,13 +224,13 @@ func (h highlightRanges) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
 func (h highlightRanges) Less(i, j int) bool { return h[i][0] < h[j][0] }
 
 func (f *Formatter) Format(w io.Writer, style *chroma.Style, iterator chroma.Iterator) (err error) {
-	return f.writeHTML(w, style, iterator.Tokens())
+	return f.writeHTML(w, style, iterator)
 }
 
 // We deliberately don't use html/template here because it is two orders of magnitude slower (benchmarked).
 //
 // OTOH we need to be super careful about correct escaping...
-func (f *Formatter) writeHTML(w io.Writer, style *chroma.Style, tokens []chroma.Token) (err error) { // nolint: gocyclo
+func (f *Formatter) writeHTML(w io.Writer, style *chroma.Style, tokens iter.Seq[chroma.Token]) (err error) { // nolint: gocyclo
 	css := f.styleCache.get(style, true)
 	if f.standalone {
 		fmt.Fprint(w, "<html>\n")
@@ -246,7 +248,7 @@ func (f *Formatter) writeHTML(w io.Writer, style *chroma.Style, tokens []chroma.
 
 	wrapInTable := f.lineNumbers && f.lineNumbersInTable
 
-	lines := chroma.SplitTokensIntoLines(tokens)
+	lines := slices.Collect(chroma.SplitTokensIntoLines(tokens))
 	lineDigits := len(strconv.Itoa(f.baseLineNumber + len(lines) - 1))
 	highlightIndex := 0
 

formatters/html/html_test.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"io"
 	"regexp"
+	"slices"
 	"strings"
 	"testing"
 
@@ -47,7 +48,9 @@ func TestSplitTokensIntoLines(t *testing.T) {
 			{Type: chroma.NameKeyword, Value: "what?\n"},
 		},
 	}
-	actual := chroma.SplitTokensIntoLines(in)
+	actual := slices.Collect(chroma.SplitTokensIntoLines(slices.Values(in)))
+	t.Logf("got[%d]: %q", len(actual), actual)
+	t.Logf("want[%d]: %q", len(expected), expected)
 	assert.Equal(t, expected, actual)
 
 	in = []chroma.Token{
@@ -89,7 +92,7 @@ func TestSplitTokensIntoLines(t *testing.T) {
 			{Type: chroma.TextWhitespace, Value: "\n"},
 		},
 	}
-	actual = chroma.SplitTokensIntoLines(in)
+	actual = slices.Collect(chroma.SplitTokensIntoLines(slices.Values(in)))
 	assert.Equal(t, expected, actual)
 }
 

formatters/json.go

@@ -15,9 +15,6 @@ var JSON = Register("json", chroma.FormatterFunc(func(w io.Writer, s *chroma.Sty
 	}
 	i := 0
 	for t := range it {
-		if t == chroma.EOF {
-			break
-		}
 		if i > 0 {
 			if _, err := fmt.Fprintln(w, ","); err != nil {
 				return err

formatters/svg/svg.go

@@ -8,6 +8,7 @@ import (
 	"io"
 	"os"
 	"path"
+	"slices"
 	"strings"
 
 	"github.com/alecthomas/chroma/v2"
@@ -62,7 +63,7 @@ type Formatter struct {
 }
 
 func (f *Formatter) Format(w io.Writer, style *chroma.Style, iterator chroma.Iterator) (err error) {
-	f.writeSVG(w, style, iterator.Tokens())
+	f.writeSVG(w, style, iterator)
 	return err
 }
 
@@ -80,9 +81,9 @@ func escapeString(s string) string {
 	return svgEscaper.Replace(s)
 }
 
-func (f *Formatter) writeSVG(w io.Writer, style *chroma.Style, tokens []chroma.Token) { // nolint: gocyclo
+func (f *Formatter) writeSVG(w io.Writer, style *chroma.Style, tokens chroma.Iterator) { // nolint: gocyclo
 	svgStyles := f.styleToSVG(style)
-	lines := chroma.SplitTokensIntoLines(tokens)
+	lines := slices.Collect(chroma.SplitTokensIntoLines(tokens))
 
 	fmt.Fprint(w, "<?xml version=\"1.0\" encoding=\"UTF-8\"?>\n")
 	fmt.Fprint(w, "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.0//EN\" \"http://www.w3.org/TR/2001/REC-SVG-20010904/DTD/svg10.dtd\">\n")

formatters/tokens.go

@@ -10,9 +10,6 @@ import (
 // Tokens formatter outputs the raw token structures.
 var Tokens = Register("tokens", chroma.FormatterFunc(func(w io.Writer, s *chroma.Style, it chroma.Iterator) error {
 	for t := range it {
-		if t == chroma.EOF {
-			break
-		}
 		if _, err := fmt.Fprintln(w, t.GoString()); err != nil {
 			return err
 		}

formatters/tty_indexed.go

@@ -240,9 +240,6 @@ type indexedTTYFormatter struct {
 func (c *indexedTTYFormatter) Format(w io.Writer, style *chroma.Style, it chroma.Iterator) (err error) {
 	theme := styleToEscapeSequence(c.table, style)
 	for token := range it {
-		if token == chroma.EOF {
-			break
-		}
 		clr, ok := theme[token.Type]
 
 		// This search mimics how styles.Get() is used in tty_truecolour.go.

formatters/tty_truecolour.go

@@ -47,9 +47,6 @@ func writeToken(w io.Writer, formatting string, text string) {
 func trueColourFormatter(w io.Writer, style *chroma.Style, it chroma.Iterator) error {
 	style = clearBackground(style)
 	for token := range it {
-		if token == chroma.EOF {
-			break
-		}
 		entry := style.Get(token.Type)
 		if entry.IsZero() {
 			fmt.Fprint(w, token.Value)