feat: switch to stdlib iterators

.golangci.yml

@@ -48,6 +48,7 @@ linters:
     - recvcheck
     - tenv
     - err113
+    - prealloc
 
 linters-settings:
   gocyclo:
@@ -62,7 +63,6 @@ linters-settings:
     #  - (Must)?NewLexer$
     exclude_godoc_examples: false
 
-
 issues:
   exclude-dirs:
     - _examples
@@ -71,19 +71,19 @@ issues:
   exclude-use-default: false
   exclude:
     # Captured by errcheck.
-    - '^(G104|G204):'
+    - "^(G104|G204):"
     # Very commonly not checked.
     - 'Error return value of .(.*\.Help|.*\.MarkFlagRequired|(os\.)?std(out|err)\..*|.*Close|.*Flush|os\.Remove(All)?|.*printf?|os\.(Un)?Setenv). is not checked'
     - 'exported method (.*\.MarshalJSON|.*\.UnmarshalJSON|.*\.EntityURN|.*\.GoString|.*\.Pos) should have comment or be unexported'
-    - 'composite literal uses unkeyed fields'
+    - "composite literal uses unkeyed fields"
     - 'declaration of "err" shadows declaration'
-    - 'should not use dot imports'
-    - 'Potential file inclusion via variable'
-    - 'should have comment or be unexported'
-    - 'comment on exported var .* should be of the form'
-    - 'at least one file in a package should have a package comment'
-    - 'string literal contains the Unicode'
-    - 'methods on the same type should have the same receiver name'
-    - '_TokenType_name should be _TokenTypeName'
-    - '`_TokenType_map` should be `_TokenTypeMap`'
-    - 'rewrite if-else to switch statement'
+    - "should not use dot imports"
+    - "Potential file inclusion via variable"
+    - "should have comment or be unexported"
+    - "comment on exported var .* should be of the form"
+    - "at least one file in a package should have a package comment"
+    - "string literal contains the Unicode"
+    - "methods on the same type should have the same receiver name"
+    - "_TokenType_name should be _TokenTypeName"
+    - "`_TokenType_map` should be `_TokenTypeMap`"
+    - "rewrite if-else to switch statement"

cmd/chroma/go.mod

@@ -1,6 +1,6 @@
 module github.com/alecthomas/chroma/v2/cmd/chroma
 
-go 1.22
+go 1.23
 
 toolchain go1.25.1
 

cmd/chroma/main.go

@@ -393,7 +393,10 @@ 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 := it(); token != chroma.EOF; token = it() {
+	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())
 		}

cmd/chromad/go.mod

@@ -1,6 +1,6 @@
 module github.com/alecthomas/chroma/v2/cmd/chromad
 
-go 1.22
+go 1.23
 
 toolchain go1.25.1
 

coalesce.go

@@ -6,13 +6,16 @@ func Coalesce(lexer Lexer) Lexer { return &coalescer{lexer} }
 type coalescer struct{ Lexer }
 
 func (d *coalescer) Tokenise(options *TokeniseOptions, text string) (Iterator, error) {
-	var prev Token
 	it, err := d.Lexer.Tokenise(options, text)
 	if err != nil {
 		return nil, err
 	}
-	return func() Token {
-		for token := it(); token != (EOF); token = it() {
+	return func(yield func(Token) bool) {
+		var prev Token
+		for token := range it {
+			if token == EOF {
+				break
+			}
 			if len(token.Value) == 0 {
 				continue
 			}
@@ -22,14 +25,15 @@ func (d *coalescer) Tokenise(options *TokeniseOptions, text string) (Iterator, e
 				if prev.Type == token.Type && len(prev.Value) < 8192 {
 					prev.Value += token.Value
 				} else {
-					out := prev
+					if !yield(prev) {
+						return
+					}
 					prev = token
-					return out
 				}
 			}
 		}
-		out := prev
-		prev = EOF
-		return out
+		if prev != EOF {
+			yield(prev)
+		}
 	}, nil
 }

formatters/api.go

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

formatters/json.go

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

formatters/tokens.go

@@ -9,7 +9,10 @@ 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 := it(); t != chroma.EOF; t = it() {
+	for t := range it {
+		if t == chroma.EOF {
+			break
+		}
 		if _, err := fmt.Fprintln(w, t.GoString()); err != nil {
 			return err
 		}

formatters/tty_indexed.go

@@ -239,7 +239,10 @@ 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 := it(); token != chroma.EOF; token = it() {
+	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

@@ -46,7 +46,10 @@ 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 := it(); token != chroma.EOF; token = it() {
+	for token := range it {
+		if token == chroma.EOF {
+			break
+		}
 		entry := style.Get(token.Type)
 		if entry.IsZero() {
 			fmt.Fprint(w, token.Value)

go.mod

@@ -1,6 +1,6 @@
 module github.com/alecthomas/chroma/v2
 
-go 1.22
+go 1.23
 
 require (
 	github.com/alecthomas/assert/v2 v2.11.0

iterator.go

@@ -1,57 +1,53 @@
 package chroma
 
-import "strings"
+import (
+	"iter"
+	"strings"
+)
 
 // An Iterator across tokens.
 //
 // EOF will be returned at the end of the Token stream.
 //
 // If an error occurs within an Iterator, it may propagate this in a panic. Formatters should recover.
-type Iterator func() Token
+type Iterator iter.Seq[Token]
 
 // Tokens consumes all tokens from the iterator and returns them as a slice.
 func (i Iterator) Tokens() []Token {
 	var out []Token
-	for t := i(); t != EOF; t = i() {
+	for t := range i {
+		if t == EOF {
+			break
+		}
 		out = append(out, t)
 	}
 	return out
 }
 
-// Stdlib converts a Chroma iterator to a Go 1.23-compatible iterator.
-func (i Iterator) Stdlib() func(yield func(Token) bool) {
-	return func(yield func(Token) bool) {
-		for t := i(); t != EOF; t = i() {
-			if !yield(t) {
-				return
-			}
-		}
-	}
-}
-
 // Concaterator concatenates tokens from a series of iterators.
 func Concaterator(iterators ...Iterator) Iterator {
-	return func() Token {
-		for len(iterators) > 0 {
-			t := iterators[0]()
-			if t != EOF {
-				return t
+	return func(yield func(Token) bool) {
+		for _, it := range iterators {
+			for t := range it {
+				if t == EOF {
+					break
+				}
+				if !yield(t) {
+					return
+				}
 			}
-			iterators = iterators[1:]
 		}
-		return EOF
 	}
 }
 
 // Literator converts a sequence of literal Tokens into an Iterator.
 func Literator(tokens ...Token) Iterator {
-	return func() Token {
-		if len(tokens) == 0 {
-			return EOF
+	return func(yield func(Token) bool) {
+		for _, token := range tokens {
+			if !yield(token) {
+				return
+			}
 		}
-		token := tokens[0]
-		tokens = tokens[1:]
-		return token
 	}
 }
 

lexers/http.go

@@ -69,63 +69,71 @@ func httpBodyContentTypeLexer(lexer Lexer) Lexer { return &httpBodyContentTyper{
 type httpBodyContentTyper struct{ Lexer }
 
 func (d *httpBodyContentTyper) Tokenise(options *TokeniseOptions, text string) (Iterator, error) { // nolint: gocognit
-	var contentType string
-	var isContentType bool
-	var subIterator Iterator
-
 	it, err := d.Lexer.Tokenise(options, text)
 	if err != nil {
 		return nil, err
 	}
 
-	return func() Token {
-		token := it()
+	return func(yield func(Token) bool) {
+		var contentType string
+		var isContentType bool
+		var subIterator Iterator
 
-		if token == EOF {
-			if subIterator != nil {
-				return subIterator()
+		for token := range it {
+			if token == EOF {
+				break
 			}
-			return EOF
-		}
 
-		switch {
-		case token.Type == Name && strings.ToLower(token.Value) == "content-type":
-			{
-				isContentType = true
-			}
-		case token.Type == Literal && isContentType:
-			{
-				isContentType = false
-				contentType = strings.TrimSpace(token.Value)
-				pos := strings.Index(contentType, ";")
-				if pos > 0 {
-					contentType = strings.TrimSpace(contentType[:pos])
+			switch {
+			case token.Type == Name && strings.ToLower(token.Value) == "content-type":
+				{
+					isContentType = true
 				}
-			}
-		case token.Type == Generic && contentType != "":
-			{
-				lexer := MatchMimeType(contentType)
-
-				// application/calendar+xml can be treated as application/xml
-				// if there's not a better match.
-				if lexer == nil && strings.Contains(contentType, "+") {
-					slashPos := strings.Index(contentType, "/")
-					plusPos := strings.LastIndex(contentType, "+")
-					contentType = contentType[:slashPos+1] + contentType[plusPos+1:]
-					lexer = MatchMimeType(contentType)
+			case token.Type == Literal && isContentType:
+				{
+					isContentType = false
+					contentType = strings.TrimSpace(token.Value)
+					pos := strings.Index(contentType, ";")
+					if pos > 0 {
+						contentType = strings.TrimSpace(contentType[:pos])
+					}
 				}
+			case token.Type == Generic && contentType != "":
+				{
+					lexer := MatchMimeType(contentType)
+
+					// application/calendar+xml can be treated as application/xml
+					// if there's not a better match.
+					if lexer == nil && strings.Contains(contentType, "+") {
+						slashPos := strings.Index(contentType, "/")
+						plusPos := strings.LastIndex(contentType, "+")
+						contentType = contentType[:slashPos+1] + contentType[plusPos+1:]
+						lexer = MatchMimeType(contentType)
+					}
 
-				if lexer == nil {
-					token.Type = Text
-				} else {
-					subIterator, err = lexer.Tokenise(nil, token.Value)
-					if err != nil {
-						panic(err)
+					if lexer == nil {
+						token.Type = Text
+					} else {
+						subIterator, err = lexer.Tokenise(nil, token.Value)
+						if err != nil {
+							panic(err)
+						}
+						// Emit tokens from the sub-iterator
+						for st := range subIterator {
+							if st == EOF {
+								break
+							}
+							if !yield(st) {
+								return
+							}
+						}
+						continue
 					}
-					return EOF
 				}
 			}
+			if !yield(token) {
+				return
+			}
 		}
-		return token
 	}, nil
 }

lexers/lexer_benchmark_test.go

@@ -212,7 +212,10 @@ func Benchmark(b *testing.B) {
 	for range b.N {
 		it, err := lexers.GlobalLexerRegistry.Get("Java").Tokenise(nil, lexerBenchSource)
 		assert.NoError(b, err)
-		for t := it(); t != chroma.EOF; t = it() {
+		for t := range it {
+			if t == chroma.EOF {
+				break
+			}
 		}
 	}
 }