Support parser/lexer cancellation

examples/arithmetics/benchmark_test.go

@@ -30,7 +30,8 @@ func BenchmarkParser(b *testing.B) {
 	srv := CreateServices()
 	lexerService := service.MustGet[lexer.Lexer](srv)
 	parserService := service.MustGet[parser.Parser](srv)
-	tokens := lexerService.Exec(content).Tokens
+	lexerResult, _ := lexerService.Exec(b.Context(), content)
+	tokens := lexerResult.Tokens
 	doc, err := fastbelt.NewDocumentFromString("file:///workspace/bench.calc", "arithmetics", content)
 	if err != nil {
 		b.Fatal(err)
@@ -39,7 +40,7 @@ func BenchmarkParser(b *testing.B) {
 	b.SetBytes(int64(len(content)))
 	b.ResetTimer()
 	for b.Loop() {
-		result := parserService.Parse(doc)
+		result, _ := parserService.Parse(b.Context(), doc)
 		doc.Root = result.Node
 	}
 }
@@ -57,8 +58,9 @@ func BenchmarkLexerAndParser(b *testing.B) {
 	b.SetBytes(int64(len(content)))
 	b.ResetTimer()
 	for b.Loop() {
-		doc.Tokens = lexerService.Exec(content).Tokens
-		result := parserService.Parse(doc)
+		lexerResult, _ := lexerService.Exec(b.Context(), content)
+		doc.Tokens = lexerResult.Tokens
+		result, _ := parserService.Parse(b.Context(), doc)
 		doc.Root = result.Node
 	}
 }

examples/statemachine/benchmark_test.go

@@ -11,7 +11,6 @@ import (
 	"testing"
 	"time"
 
-	"github.com/stretchr/testify/assert"
 	"typefox.dev/fastbelt"
 	"typefox.dev/fastbelt/lexer"
 	"typefox.dev/fastbelt/parser"
@@ -68,92 +67,15 @@ func BenchmarkWorkspaceCycle(b *testing.B) {
 	b.ReportMetric(msPerResource, "ms/resource")
 }
 
-func BenchmarkTraverseContentSeq(b *testing.B) {
-	content, _ := generateStatemachineContent(0)
-	srv := CreateServices()
-	documentParser := service.MustGet[workspace.DocumentParser](srv)
-	doc, err := fastbelt.NewDocumentFromString("file:///workspace/statemachine_0.statemachine", "statemachine", content)
-	if err != nil {
-		b.Fatal(err)
-	}
-	documentParser.Parse(doc)
-
-	for b.Loop() {
-		count := 0
-		for range fastbelt.AllChildren(doc.Root) {
-			count++
-		}
-		_ = count
-	}
-}
-
-// BenchmarkTraverseContentSeqHalf breaks out of the loop after visiting half of the
-// nodes. With short-circuiting traversal this should take roughly half the time of
-// BenchmarkTraverseContentSeq.
-func BenchmarkTraverseContentSeqHalf(b *testing.B) {
-	content, elementCount := generateStatemachineContent(0)
-	srv := CreateServices()
-	documentParser := service.MustGet[workspace.DocumentParser](srv)
-	doc, err := fastbelt.NewDocumentFromString("file:///workspace/statemachine_0.statemachine", "statemachine", content)
-	if err != nil {
-		b.Fatal(err)
-	}
-	documentParser.Parse(doc)
-	half := elementCount / 2
-
-	for b.Loop() {
-		count := 0
-		for range fastbelt.AllChildren(doc.Root) {
-			count++
-			if count == half {
-				break
-			}
-		}
-		_ = count
-	}
-}
-
-func TestAllNodesEquivalence(t *testing.T) {
-	content, elementCount := generateStatemachineContent(0)
-	srv := CreateServices()
-	documentParser := service.MustGet[workspace.DocumentParser](srv)
-	doc, err := fastbelt.NewDocumentFromString("file:///workspace/statemachine_0.statemachine", "statemachine", content)
-	if err != nil {
-		t.Fatal(err)
-	}
-	documentParser.Parse(doc)
-	nodeCount := 0
-	for range fastbelt.AllNodes(doc.Root) {
-		nodeCount++
-	}
-	assert.Equal(t, elementCount, nodeCount, "AllNodes should iterate all nodes in the document")
-}
-
-func TestAllChildrenEquivalence(t *testing.T) {
-	content, elementCount := generateStatemachineContent(0)
-	totalCount := elementCount - 1 // AllChildren does not include the root node, so we subtract 1 from the total count
-	srv := CreateServices()
-	documentParser := service.MustGet[workspace.DocumentParser](srv)
-	doc, err := fastbelt.NewDocumentFromString("file:///workspace/statemachine_0.statemachine", "statemachine", content)
-	if err != nil {
-		t.Fatal(err)
-	}
-	documentParser.Parse(doc)
-	childCount := 0
-	for range fastbelt.AllChildren(doc.Root) {
-		childCount++
-	}
-	assert.Equal(t, totalCount, childCount, "AllChildren should iterate all child nodes in the document")
-}
-
 // BenchmarkParser benchmarks parsing a single generated statemachine document,
 // reusing the pre-lexed token slice every iteration.
 func BenchmarkParser(b *testing.B) {
 	content, _ := generateStatemachineContent(0)
 	srv := CreateServices()
 	lexerService := service.MustGet[lexer.Lexer](srv)
 	parserService := service.MustGet[parser.Parser](srv)
-	tokens := lexerService.Exec(content).Tokens
+	lexerResult, _ := lexerService.Exec(b.Context(), content)
+	tokens := lexerResult.Tokens
 	doc, err := fastbelt.NewDocumentFromString("file:///workspace/statemachine_0.statemachine", "statemachine", content)
 	if err != nil {
 		b.Fatal(err)
@@ -162,7 +84,7 @@ func BenchmarkParser(b *testing.B) {
 	b.SetBytes(int64(len(content)))
 	b.ResetTimer()
 	for b.Loop() {
-		result := parserService.Parse(doc)
+		result, _ := parserService.Parse(b.Context(), doc)
 		doc.Root = result.Node
 	}
 }
@@ -174,7 +96,7 @@ func BenchmarkLexer(b *testing.B) {
 	b.SetBytes(int64(len(content)))
 	b.ResetTimer()
 	for b.Loop() {
-		_ = l.Exec(content)
+		_, _ = l.Exec(b.Context(), content)
 	}
 }
 
@@ -191,8 +113,9 @@ func BenchmarkLexerAndParser(b *testing.B) {
 	b.SetBytes(int64(len(content)))
 	b.ResetTimer()
 	for b.Loop() {
-		doc.Tokens = lexerService.Exec(content).Tokens
-		result := parserService.Parse(doc)
+		lexerResult, _ := lexerService.Exec(b.Context(), content)
+		doc.Tokens = lexerResult.Tokens
+		result, _ := parserService.Parse(b.Context(), doc)
 		doc.Root = result.Node
 	}
 }
@@ -226,9 +149,14 @@ func BenchmarkLocalLinking(b *testing.B) {
 //   - 3 commands
 //   - 50 states, each with transitions that cycle through events/states
 func generateStatemachineContent(index int) (string, int) {
+	return generateStatemachineContentWithStates(index, 50)
+}
+
+// generateStatemachineContentWithStates is like [generateStatemachineContent]
+// with a configurable number of states.
+func generateStatemachineContentWithStates(index, numStates int) (string, int) {
 	const numEvents = 4
 	const numCommands = 3
-	const numStates = 50
 	elementCount := 1
 
 	var sb strings.Builder
@@ -269,23 +197,3 @@ func generateStatemachineContent(index int) (string, int) {
 
 	return sb.String(), elementCount
 }
-
-func TestAllNodesEarlyExitShortCircuits(t *testing.T) {
-	content, elementCount := generateStatemachineContent(0)
-	srv := CreateServices()
-	documentParser := service.MustGet[workspace.DocumentParser](srv)
-	doc, err := fastbelt.NewDocumentFromString("file:///workspace/statemachine_0.statemachine", "statemachine", content)
-	if err != nil {
-		t.Fatal(err)
-	}
-	documentParser.Parse(doc)
-	visited := 0
-	for range fastbelt.AllNodes(doc.Root) {
-		visited++
-		if visited == 3 {
-			break
-		}
-	}
-	assert.Equal(t, 3, visited)
-	assert.Greater(t, elementCount, 3)
-}

examples/statemachine/cancellation_test.go

@@ -0,0 +1,98 @@
+package statemachine
+
+import (
+	"context"
+	"runtime"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"typefox.dev/fastbelt"
+	"typefox.dev/fastbelt/lexer"
+	"typefox.dev/fastbelt/parser"
+	"typefox.dev/fastbelt/util/service"
+)
+
+func TestParseCancelledContextTerminates(t *testing.T) {
+	content, _ := generateStatemachineContent(0)
+	srv := CreateServices()
+	doc, err := fastbelt.NewDocumentFromString("file:///workspace/statemachine_0.statemachine", "statemachine", content)
+	if err != nil {
+		t.Fatal(err)
+	}
+	lexer := service.MustGet[lexer.Lexer](srv)
+	parser := service.MustGet[parser.Parser](srv)
+	lexerResult, _ := lexer.Exec(t.Context(), content)
+	doc.Tokens = lexerResult.Tokens
+	ctx, cancel := context.WithCancel(t.Context())
+	// Cancel ahead of time
+	cancel()
+	_, err = parser.Parse(ctx, doc)
+	assert.Error(t, err)
+}
+
+// hugeStateCount yields a ~20 MB document, large enough that timer and
+// scheduling jitter stay small relative to a full lexer or parser run.
+const hugeStateCount = 200_000
+
+func TestLexerCancellationStopsEarly(t *testing.T) {
+	if testing.Short() {
+		t.Skip("skip long running lexer cancellation test")
+	}
+	content, _ := generateStatemachineContentWithStates(0, hugeStateCount)
+	lexer := service.MustGet[lexer.Lexer](CreateServices())
+	assertCancellationStopsEarly(t, func(ctx context.Context) error {
+		_, err := lexer.Exec(ctx, content)
+		return err
+	})
+}
+
+func TestParserCancellationStopsEarly(t *testing.T) {
+	if testing.Short() {
+		t.Skip("skip long running parser cancellation test")
+	}
+	content, _ := generateStatemachineContentWithStates(0, hugeStateCount)
+	srv := CreateServices()
+	doc, err := fastbelt.NewDocumentFromString("file:///workspace/statemachine_0.statemachine", "statemachine", content)
+	if err != nil {
+		t.Fatal(err)
+	}
+	lexer := service.MustGet[lexer.Lexer](srv)
+	lexerResult, err := lexer.Exec(t.Context(), content)
+	require.NoError(t, err)
+	doc.Tokens = lexerResult.Tokens
+	parser := service.MustGet[parser.Parser](srv)
+	assertCancellationStopsEarly(t, func(ctx context.Context) error {
+		_, err := parser.Parse(ctx, doc)
+		return err
+	})
+}
+
+// assertCancellationStopsEarly times a full run, then reruns with
+// cancellation fired after 20% of that time and asserts the cancelled run
+// finishes within 30% of it.
+func assertCancellationStopsEarly(t *testing.T, run func(ctx context.Context) error) {
+	t.Helper()
+	err := run(t.Context()) // warm up
+	require.NoError(t, err)
+
+	runtime.GC()
+	start := time.Now()
+	err = run(t.Context())
+	require.NoError(t, err)
+	full := time.Since(start)
+
+	ctx, cancel := context.WithCancel(t.Context())
+	defer cancel()
+	runtime.GC()
+	start = time.Now()
+	timer := time.AfterFunc(full/5, cancel)
+	defer timer.Stop()
+	err = run(ctx)
+	assert.Error(t, err)
+	cancelled := time.Since(start)
+
+	t.Logf("full: %v, cancelled: %v", full, cancelled)
+	assert.Less(t, cancelled, full*3/10)
+}