refactor: stream delegated tokens

DelegatingLexer buffered every insertion, root token, and output token, causing token memory to grow with input size. Merge iterators lazily while retaining the concatenated Other input so embedded spans cannot corrupt stateful root lexers.

Author: alecthomas

delegate.go

@@ -1,9 +1,8 @@
 package chroma
 
 import (
-	"bytes"
 	"iter"
-	"slices"
+	"strings"
 )
 
 type delegatingLexer struct {
@@ -14,7 +13,7 @@ type delegatingLexer struct {
 // DelegatingLexer combines two lexers to handle the common case of a language embedded inside another, such as PHP
 // inside HTML or PHP inside plain text.
 //
-// It takes two lexer as arguments: a root lexer and a language lexer.  First everything is scanned using the language
+// It takes two lexers as arguments: a root lexer and a language lexer. First everything is scanned using the language
 // lexer, which must return "Other" for unrecognised tokens. Then all "Other" tokens are lexed using the root lexer.
 // Finally, these two sets of tokens are merged.
 //
@@ -54,79 +53,86 @@ func (d *delegatingLexer) Config() *Config {
 	return d.language.Config()
 }
 
-// An insertion is the character range where language tokens should be inserted.
-type insertion struct {
-	start, end int
-	tokens     []Token
-}
-
 func (d *delegatingLexer) Tokenise(options *TokeniseOptions, text string) (iter.Seq[Token], error) { // nolint: gocognit
-	tokens, err := Tokenise(Coalesce(d.language), options, text)
+	languageTokens, err := Tokenise(Coalesce(d.language), options, text)
 	if err != nil {
 		return nil, err
 	}
-	// Compute insertions and gather "Other" tokens.
-	others := &bytes.Buffer{}
-	insertions := []*insertion{}
-	var insert *insertion
-	offset := 0
-	first := true
-	var lastType TokenType
-	for t := range tokens {
-		if t.Type == Other {
-			if !first && insert != nil && lastType != Other {
-				insert.end = offset
-			}
-			others.WriteString(t.Value)
+	var rootText strings.Builder
+	hasLanguageTokens := false
+	for token := range languageTokens {
+		if token.Type == Other {
+			rootText.WriteString(token.Value)
 		} else {
-			if first || lastType == Other {
-				insert = &insertion{start: offset}
-				insertions = append(insertions, insert)
-			}
-			insert.tokens = append(insert.tokens, t)
+			hasLanguageTokens = true
 		}
-		first = false
-		lastType = t.Type
-		offset += len(t.Value)
 	}
-
-	if len(insertions) == 0 {
+	if !hasLanguageTokens {
 		return d.root.Tokenise(options, text)
 	}
 
-	// Lex the other tokens.
-	rootIt, err := Tokenise(Coalesce(d.root), options, others.String())
+	rootTokens, err := Tokenise(Coalesce(d.root), options, rootText.String())
 	if err != nil {
 		return nil, err
 	}
-	rootTokens := slices.Collect(rootIt)
+	languageTokens, err = Tokenise(Coalesce(d.language), options, text)
+	if err != nil {
+		return nil, err
+	}
+	return func(yield func(Token) bool) {
+		nextRoot, stopRoot := iter.Pull(rootTokens)
+		defer stopRoot()
+
+		var root Token
+		var pendingRoot Token
+		for language := range languageTokens {
+			if language.Value == "" {
+				continue
+			}
+			if language.Type != Other {
+				if pendingRoot.Value != "" {
+					if !yield(pendingRoot) {
+						return
+					}
+					pendingRoot = Token{}
+				}
+				if !yield(language) {
+					return
+				}
+				continue
+			}
 
-	// Interleave the two sets of tokens.
-	var out []Token
-	offset = 0
-	ti := 0
-	ii := 0
-	for ti < len(rootTokens) || ii < len(insertions) {
-		if ti >= len(rootTokens) || (ii < len(insertions) && insertions[ii].start < offset+len(rootTokens[ti].Value)) {
-			ins := insertions[ii]
-			ii++
-			if ti < len(rootTokens) {
-				l, r := splitToken(rootTokens[ti], ins.start-offset)
-				if l.Value != "" {
-					out = append(out, l)
-					offset += len(l.Value)
+			// Consume only Other spans because embedded language spans were removed from the root input.
+			remaining := len(language.Value)
+			for remaining > 0 {
+				if root.Value == "" {
+					var ok bool
+					root, ok = nextRoot()
+					if !ok {
+						break
+					}
+					continue
+				}
+				consumed, rest := splitToken(root, min(remaining, len(root.Value)))
+				root = rest
+				remaining -= len(consumed.Value)
+				if pendingRoot.Value == "" {
+					pendingRoot = consumed
+				} else {
+					pendingRoot.Value += consumed.Value
+				}
+				if root.Value == "" {
+					if !yield(pendingRoot) {
+						return
+					}
+					pendingRoot = Token{}
 				}
-				rootTokens[ti] = r
 			}
-			out = append(out, ins.tokens...)
-			offset += ins.end - ins.start
-		} else {
-			out = append(out, rootTokens[ti])
-			offset += len(rootTokens[ti].Value)
-			ti++
 		}
-	}
-	return slices.Values(out), nil
+		if pendingRoot.Value != "" {
+			yield(pendingRoot)
+		}
+	}, nil
 }
 
 func splitToken(t Token, offset int) (l Token, r Token) {

delegate_test.go

@@ -1,12 +1,32 @@
 package chroma
 
 import (
+	"iter"
 	"slices"
 	"testing"
 
 	assert "github.com/alecthomas/assert/v2"
 )
 
+type countingLexer struct {
+	Lexer
+	tokens   []Token
+	consumed int
+	texts    []string
+}
+
+func (l *countingLexer) Tokenise(_ *TokeniseOptions, text string) (iter.Seq[Token], error) {
+	l.texts = append(l.texts, text)
+	return func(yield func(Token) bool) {
+		for _, token := range l.tokens {
+			l.consumed++
+			if !yield(token) {
+				return
+			}
+		}
+	}, nil
+}
+
 func makeDelegationTestLexers(t *testing.T) (lang Lexer, root Lexer) {
 	t.Helper()
 	return mustNewLexer(t, nil, Rules{ // nolint: forbidigo
@@ -110,3 +130,48 @@ func TestDelegate(t *testing.T) {
 		})
 	}
 }
+
+func TestDelegateStreamsRootTokens(t *testing.T) {
+	root := &countingLexer{tokens: []Token{
+		{Name, "a"},
+		{Text, "c"},
+		{Name, "e"},
+	}}
+	language := &countingLexer{tokens: []Token{
+		{Other, "a"},
+		{Keyword, "b"},
+		{Other, "c"},
+		{Keyword, "d"},
+		{Other, "e"},
+	}}
+	it, err := DelegatingLexer(root, language).Tokenise(nil, "abcde")
+	assert.NoError(t, err)
+	assert.Equal(t, 0, root.consumed)
+	assert.Equal(t, 5, language.consumed)
+
+	var actual Token
+	for token := range it {
+		actual = token
+		break
+	}
+	assert.Equal(t, Token{Name, "a"}, actual)
+	assert.Equal(t, 2, root.consumed)
+	assert.Equal(t, 7, language.consumed)
+}
+
+func TestDelegateSplitsRootTokens(t *testing.T) {
+	root := &countingLexer{tokens: []Token{{Text, "before  after"}}}
+	language := &countingLexer{tokens: []Token{
+		{Other, "before "},
+		{Keyword, "embedded"},
+		{Other, " after"},
+	}}
+	it, err := DelegatingLexer(root, language).Tokenise(nil, "before embedded after")
+	assert.NoError(t, err)
+	assert.Equal(t, []string{"before  after"}, root.texts)
+	assert.Equal(t, []Token{
+		{Text, "before "},
+		{Keyword, "embedded"},
+		{Text, " after"},
+	}, slices.Collect(it))
+}