Pooled version of lexer.Upgrade

lexer/peek.go

@@ -1,5 +1,7 @@
 package lexer
 
+import "sync"
+
 // PeekingLexer supports arbitrary lookahead as well as cloning.
 type PeekingLexer struct {
 	Checkpoint
@@ -26,6 +28,10 @@ func Upgrade(lex Lexer, elide ...TokenType) (*PeekingLexer, error) {
 	r := &PeekingLexer{
 		elide: make(map[TokenType]bool, len(elide)),
 	}
+	return fillInPeekingLexer(lex, r, elide...)
+}
+
+func fillInPeekingLexer(lex Lexer, r *PeekingLexer, elide ...TokenType) (*PeekingLexer, error) {
 	for _, rn := range elide {
 		r.elide[rn] = true
 	}
@@ -43,6 +49,43 @@ func Upgrade(lex Lexer, elide ...TokenType) (*PeekingLexer, error) {
 	return r, nil
 }
 
+var peekingLexerPool = sync.Pool{
+	New: func() interface{} {
+		return &PeekingLexer{
+			elide: make(map[TokenType]bool, 4),
+		}
+	},
+}
+
+// UpgradePooled will upgrade a Lexer to a PeekingLexer with arbitrary
+// lookahead. Faster if you need to lex thousands of similar documents.
+//
+// "elide" is a slice of token types to elide from processing.
+//
+// You must call `PutBackPooledPeekingLexer` once done with the
+// returned lexer in all cases (ok or error). If you use the lexer with
+// the parser (`Parser.ParseFromLexer`), note that the parsed results
+// might refer back to lexer tokens, in which case you should not call
+// PutBackPooledPeekingLexer until you have finished with the parser
+// results as well.
+func UpgradePooled(lex Lexer, elide ...TokenType) (*PeekingLexer, error) {
+	r := peekingLexerPool.Get().(*PeekingLexer)
+	// reset the state of the PeekingLexer to empty (preserving any allocated capacity)
+	r.Checkpoint = Checkpoint{}
+	// note: this preserves capacity
+	r.tokens = r.tokens[:0]
+	for k := range r.elide {
+		delete(r.elide, k)
+	}
+	return fillInPeekingLexer(lex, r, elide...)
+}
+
+func PutBackPooledPeekingLexer(r *PeekingLexer) {
+	if r != nil {
+		peekingLexerPool.Put(r)
+	}
+}
+
 // Range returns the slice of tokens between the two cursor points.
 func (p *PeekingLexer) Range(rawStart, rawEnd RawCursor) []Token {
 	return p.tokens[rawStart:rawEnd]

lexer/peek_test.go

@@ -1,6 +1,7 @@
 package lexer_test
 
 import (
+	"math/rand"
 	"testing"
 
 	require "github.com/alecthomas/assert/v2"
@@ -33,6 +34,19 @@ func TestUpgrade(t *testing.T) {
 	require.Equal(t, tokens, l.Range(0, 3))
 }
 
+func TestUpgradePooled(t *testing.T) {
+	t0 := lexer.Token{Type: 1, Value: "moo"}
+	ts := lexer.Token{Type: 3, Value: " "}
+	t1 := lexer.Token{Type: 2, Value: "blah"}
+	tokens := []lexer.Token{t0, ts, t1}
+	l, err := lexer.UpgradePooled(&staticLexer{tokens: tokens}, 3)
+	defer lexer.PutBackPooledPeekingLexer(l)
+	require.NoError(t, err)
+	require.Equal(t, t0, *l.Peek())
+	require.Equal(t, t0, *l.Peek())
+	require.Equal(t, tokens, l.Range(0, 3))
+}
+
 func TestPeekingLexer_Peek_Next_Checkpoint(t *testing.T) {
 	slexdef := lexer.MustSimple([]lexer.SimpleRule{
 		{"Ident", `\w+`},
@@ -71,3 +85,37 @@ func BenchmarkPeekingLexer_Peek(b *testing.B) {
 	}
 	require.Equal(b, lexer.Token{Type: 2, Value: "y"}, *t)
 }
+
+func generateTokenArray(sz int) []lexer.Token {
+	tokens := make([]lexer.Token, sz)
+	for i := range tokens {
+		tokens[i].Type = lexer.TokenType(rand.Int() % 3)                /* #nosec */
+		tokens[i].Value = string([]byte{byte('a' + (rand.Int() % 26))}) /* #nosec */
+	}
+	return tokens
+}
+
+func BenchmarkPeekingLexer_Upgrade(b *testing.B) {
+	tokens := generateTokenArray(10000)
+	testLexer := &staticLexer{tokens: tokens}
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		l, err := lexer.Upgrade(testLexer)
+		require.NoError(b, err)
+		l.Next()
+		l.Peek()
+	}
+}
+
+func BenchmarkPeekingLexer_UpgradePooled(b *testing.B) {
+	tokens := generateTokenArray(10000)
+	testLexer := &staticLexer{tokens: tokens}
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		l, err := lexer.UpgradePooled(testLexer)
+		require.NoError(b, err)
+		l.Next()
+		l.Peek()
+		lexer.PutBackPooledPeekingLexer(l)
+	}
+}