1package chroma
 2
 3import "strings"
 4
 5// An Iterator across tokens.
 6//
 7// EOF will be returned at the end of the Token stream.
 8//
 9// If an error occurs within an Iterator, it may propagate this in a panic. Formatters should recover.
10type Iterator func() Token
11
12// Tokens consumes all tokens from the iterator and returns them as a slice.
13func (i Iterator) Tokens() []Token {
14	var out []Token
15	for t := i(); t != EOF; t = i() {
16		out = append(out, t)
17	}
18	return out
19}
20
21// Stdlib converts a Chroma iterator to a Go 1.23-compatible iterator.
22func (i Iterator) Stdlib() func(yield func(Token) bool) {
23	return func(yield func(Token) bool) {
24		for t := i(); t != EOF; t = i() {
25			if !yield(t) {
26				return
27			}
28		}
29	}
30}
31
32// Concaterator concatenates tokens from a series of iterators.
33func Concaterator(iterators ...Iterator) Iterator {
34	return func() Token {
35		for len(iterators) > 0 {
36			t := iterators[0]()
37			if t != EOF {
38				return t
39			}
40			iterators = iterators[1:]
41		}
42		return EOF
43	}
44}
45
46// Literator converts a sequence of literal Tokens into an Iterator.
47func Literator(tokens ...Token) Iterator {
48	return func() Token {
49		if len(tokens) == 0 {
50			return EOF
51		}
52		token := tokens[0]
53		tokens = tokens[1:]
54		return token
55	}
56}
57
58// SplitTokensIntoLines splits tokens containing newlines in two.
59func SplitTokensIntoLines(tokens []Token) (out [][]Token) {
60	var line []Token // nolint: prealloc
61tokenLoop:
62	for _, token := range tokens {
63		for strings.Contains(token.Value, "\n") {
64			parts := strings.SplitAfterN(token.Value, "\n", 2)
65			// Token becomes the tail.
66			token.Value = parts[1]
67
68			// Append the head to the line and flush the line.
69			clone := token.Clone()
70			clone.Value = parts[0]
71			line = append(line, clone)
72			out = append(out, line)
73			line = nil
74
75			// If the tail token is empty, don't emit it.
76			if len(token.Value) == 0 {
77				continue tokenLoop
78			}
79		}
80		line = append(line, token)
81	}
82	if len(line) > 0 {
83		out = append(out, line)
84	}
85	// Strip empty trailing token line.
86	if len(out) > 0 {
87		last := out[len(out)-1]
88		if len(last) == 1 && last[0].Value == "" {
89			out = out[:len(out)-1]
90		}
91	}
92	return out
93}