Support cancellation in checker

cmd/tsgo/main.go

@@ -2,6 +2,7 @@
 package main
 
 import (
+	"context"
 	"encoding/json"
 	"flag"
 	"fmt"
@@ -223,19 +224,19 @@ func runMain() int {
 		return 1
 	}
 
-	diagnostics = program.GetSyntacticDiagnostics(nil)
+	diagnostics = program.GetSyntacticDiagnostics(context.Background(), nil)
 	if len(diagnostics) == 0 {
 		if opts.devel.printTypes {
 			program.PrintSourceFileWithTypes()
 		} else {
 			bindStart := time.Now()
-			_ = program.GetBindDiagnostics(nil)
+			_ = program.GetBindDiagnostics(context.Background(), nil)
 			bindTime = time.Since(bindStart)
 
 			// !!! the checker already reads noCheck, but do it here just for stats printing for now
 			if compilerOptions.NoCheck.IsFalseOrUnknown() {
 				checkStart := time.Now()
-				diagnostics = slices.Concat(program.GetGlobalDiagnostics(), program.GetSemanticDiagnostics(nil))
+				diagnostics = slices.Concat(program.GetGlobalDiagnostics(), program.GetSemanticDiagnostics(context.Background(), nil))
 				checkTime = time.Since(checkStart)
 			}
 		}

internal/ast/ast.go

@@ -472,6 +472,26 @@ func (n *Node) Members() []*Node {
 	return nil
 }
 
+func (n *Node) StatementList() *NodeList {
+	switch n.Kind {
+	case KindSourceFile:
+		return n.AsSourceFile().Statements
+	case KindBlock:
+		return n.AsBlock().Statements
+	case KindModuleBlock:
+		return n.AsModuleBlock().Statements
+	}
+	panic("Unhandled case in Node.StatementList: " + n.Kind.String())
+}
+
+func (n *Node) Statements() []*Node {
+	list := n.StatementList()
+	if list != nil {
+		return list.Nodes
+	}
+	return nil
+}
+
 func (n *Node) ModifierFlags() ModifierFlags {
 	modifiers := n.Modifiers()
 	if modifiers != nil {

internal/checker/checker.go

@@ -1,6 +1,7 @@
 package checker
 
 import (
+	"context"
 	"fmt"
 	"iter"
 	"maps"
@@ -567,6 +568,7 @@ type Checker struct {
 	useUnknownInCatchVariables                 bool
 	exactOptionalPropertyTypes                 bool
 	canCollectSymbolAliasAccessibilityData     bool
+	wasCanceled                                bool
 	arrayVariances                             []VarianceFlags
 	globals                                    ast.SymbolTable
 	globalSymbols                              []*ast.Symbol
@@ -824,6 +826,7 @@ type Checker struct {
 	_jsxNamespace                              string
 	_jsxFactoryEntity                          *ast.Node
 	skipDirectInferenceNodes                   core.Set[*ast.Node]
+	ctx                                        context.Context
 }
 
 func NewChecker(program Program) *Checker {
@@ -2045,16 +2048,18 @@ func (c *Checker) getSymbol(symbols ast.SymbolTable, name string, meaning ast.Sy
 	return nil
 }
 
-func (c *Checker) CheckSourceFile(sourceFile *ast.SourceFile) {
+func (c *Checker) CheckSourceFile(ctx context.Context, sourceFile *ast.SourceFile) {
 	if SkipTypeChecking(sourceFile, c.compilerOptions) {
 		return
 	}
-	c.checkSourceFile(sourceFile)
+	c.checkSourceFile(ctx, sourceFile)
 }
 
-func (c *Checker) checkSourceFile(sourceFile *ast.SourceFile) {
+func (c *Checker) checkSourceFile(ctx context.Context, sourceFile *ast.SourceFile) {
+	c.checkNotCanceled()
 	links := c.sourceFileLinks.Get(sourceFile)
 	if !links.typeChecked {
+		c.ctx = ctx
 		// Grammar checking
 		c.checkGrammarSourceFile(sourceFile)
 		c.renamedBindingElementsInTypes = nil
@@ -2065,19 +2070,27 @@ func (c *Checker) checkSourceFile(sourceFile *ast.SourceFile) {
 			c.checkExternalModuleExports(sourceFile.AsNode())
 			c.registerForUnusedIdentifiersCheck(sourceFile.AsNode())
 		}
-		// This relies on the results of other lazy diagnostics, so must be computed after them
-		if !sourceFile.IsDeclarationFile && (c.compilerOptions.NoUnusedLocals.IsTrue() || c.compilerOptions.NoUnusedParameters.IsTrue()) {
-			c.checkUnusedIdentifiers(links.identifierCheckNodes)
-		}
-		if !sourceFile.IsDeclarationFile {
-			c.checkUnusedRenamedBindingElements()
+		if ctx.Err() == nil {
+			// This relies on the results of other lazy diagnostics, so must be computed after them
+			if !sourceFile.IsDeclarationFile && (c.compilerOptions.NoUnusedLocals.IsTrue() || c.compilerOptions.NoUnusedParameters.IsTrue()) {
+				c.checkUnusedIdentifiers(links.identifierCheckNodes)
+			}
+			if !sourceFile.IsDeclarationFile {
+				c.checkUnusedRenamedBindingElements()
+			}
+		} else {
+			c.wasCanceled = true
 		}
+		c.ctx = nil
 		links.typeChecked = true
 	}
 }
 
 func (c *Checker) checkSourceElements(nodes []*ast.Node) {
 	for _, node := range nodes {
+		if c.isCanceled() {
+			break
+		}
 		c.checkSourceElement(node)
 	}
 }
@@ -2094,7 +2107,6 @@ func (c *Checker) checkSourceElement(node *ast.Node) bool {
 }
 
 func (c *Checker) checkSourceElementWorker(node *ast.Node) {
-	// !!! Cancellation
 	kind := node.Kind
 	if kind >= ast.KindFirstStatement && kind <= ast.KindLastStatement {
 		flowNode := node.FlowNodeData().FlowNode
@@ -2246,6 +2258,9 @@ func (c *Checker) checkNodeDeferred(node *ast.Node) {
 func (c *Checker) checkDeferredNodes(context *ast.SourceFile) {
 	links := c.sourceFileLinks.Get(context)
 	for node := range links.deferredNodes.Values() {
+		if c.isCanceled() {
+			break
+		}
 		c.checkDeferredNode(node)
 	}
 	links.deferredNodes.Clear()
@@ -2291,6 +2306,9 @@ func (c *Checker) checkJSDocNodes(sourceFile *ast.SourceFile) {
 	// set parent references in JSDoc nodes.
 	for location, jsdocs := range sourceFile.JSDocCache() {
 		for _, jsdoc := range jsdocs {
+			if c.isCanceled() {
+				return
+			}
 			c.checkJSDocComments(jsdoc, location)
 			tags := jsdoc.AsJSDoc().Tags
 			if tags != nil {
@@ -3510,10 +3528,10 @@ func (c *Checker) checkBlock(node *ast.Node) {
 	}
 	if ast.IsFunctionOrModuleBlock(node) {
 		saveFlowAnalysisDisabled := c.flowAnalysisDisabled
-		node.ForEachChild(c.checkSourceElement)
+		c.checkSourceElements(node.Statements())
 		c.flowAnalysisDisabled = saveFlowAnalysisDisabled
 	} else {
-		node.ForEachChild(c.checkSourceElement)
+		c.checkSourceElements(node.Statements())
 	}
 	if len(node.Locals()) != 0 {
 		c.registerForUnusedIdentifiersCheck(node)
@@ -13134,22 +13152,31 @@ func (c *Checker) getCannotFindNameDiagnosticForName(node *ast.Node) *diagnostic
 	}
 }
 
-func (c *Checker) GetDiagnostics(sourceFile *ast.SourceFile) []*ast.Diagnostic {
+func (c *Checker) GetDiagnostics(ctx context.Context, sourceFile *ast.SourceFile) []*ast.Diagnostic {
+	c.checkNotCanceled()
 	if sourceFile != nil {
-		c.CheckSourceFile(sourceFile)
+		c.CheckSourceFile(ctx, sourceFile)
+		if c.wasCanceled {
+			return nil
+		}
 		return c.diagnostics.GetDiagnosticsForFile(sourceFile.FileName())
 	}
 	for _, file := range c.files {
-		c.CheckSourceFile(file)
+		c.CheckSourceFile(ctx, file)
+		if c.wasCanceled {
+			return nil
+		}
 	}
 	return c.diagnostics.GetDiagnostics()
 }
 
 func (c *Checker) GetDiagnosticsWithoutCheck(sourceFile *ast.SourceFile) []*ast.Diagnostic {
+	c.checkNotCanceled()
 	return c.diagnostics.GetDiagnosticsForFile(sourceFile.FileName())
 }
 
 func (c *Checker) GetGlobalDiagnostics() []*ast.Diagnostic {
+	c.checkNotCanceled()
 	return c.diagnostics.GetGlobalDiagnostics()
 }
 
@@ -29800,7 +29827,7 @@ func (c *Checker) GetEmitResolver(file *ast.SourceFile, skipDiagnostics bool) pr
 		// emitter questions of this resolver will return the right information.
 		c.emitResolver.checkerMu.Lock()
 		defer c.emitResolver.checkerMu.Unlock()
-		c.checkSourceFile(file)
+		c.checkSourceFile(context.Background(), file)
 	}
 	return c.emitResolver
 }

internal/checker/checker_test.go

@@ -68,7 +68,7 @@ func TestCheckSrcCompiler(t *testing.T) {
 		ConfigFileName: tspath.CombinePaths(rootPath, "tsconfig.json"),
 	}
 	p := compiler.NewProgram(opts)
-	p.CheckSourceFiles()
+	p.CheckSourceFiles(t.Context())
 }
 
 func BenchmarkNewChecker(b *testing.B) {

internal/checker/utilities.go

@@ -2063,3 +2063,13 @@ func IsExternalModuleSymbol(moduleSymbol *ast.Symbol) bool {
 	firstRune, _ := utf8.DecodeRuneInString(moduleSymbol.Name)
 	return moduleSymbol.Flags&ast.SymbolFlagsModule != 0 && firstRune == '"'
 }
+
+func (c *Checker) isCanceled() bool {
+	return c.ctx != nil && c.ctx.Err() != nil
+}
+
+func (c *Checker) checkNotCanceled() {
+	if c.wasCanceled {
+		panic("Checker was previously cancelled")
+	}
+}

internal/compiler/program.go

@@ -1,6 +1,7 @@
 package compiler
 
 import (
+	"context"
 	"fmt"
 	"slices"
 	"sync"
@@ -206,13 +207,13 @@ func (p *Program) BindSourceFiles() {
 	wg.RunAndWait()
 }
 
-func (p *Program) CheckSourceFiles() {
+func (p *Program) CheckSourceFiles(ctx context.Context) {
 	p.createCheckers()
 	wg := core.NewWorkGroup(p.programOptions.SingleThreaded)
 	for index, checker := range p.checkers {
 		wg.Queue(func() {
 			for i := index; i < len(p.files); i += len(p.checkers) {
-				checker.CheckSourceFile(p.files[i])
+				checker.CheckSourceFile(ctx, p.files[i])
 			}
 		})
 	}
@@ -277,16 +278,16 @@ func (p *Program) findSourceFile(candidate string, reason FileIncludeReason) *as
 	return p.filesByPath[path]
 }
 
-func (p *Program) GetSyntacticDiagnostics(sourceFile *ast.SourceFile) []*ast.Diagnostic {
-	return p.getDiagnosticsHelper(sourceFile, false /*ensureBound*/, false /*ensureChecked*/, p.getSyntacticDiagnosticsForFile)
+func (p *Program) GetSyntacticDiagnostics(ctx context.Context, sourceFile *ast.SourceFile) []*ast.Diagnostic {
+	return p.getDiagnosticsHelper(ctx, sourceFile, false /*ensureBound*/, false /*ensureChecked*/, p.getSyntacticDiagnosticsForFile)
 }
 
-func (p *Program) GetBindDiagnostics(sourceFile *ast.SourceFile) []*ast.Diagnostic {
-	return p.getDiagnosticsHelper(sourceFile, true /*ensureBound*/, false /*ensureChecked*/, p.getBindDiagnosticsForFile)
+func (p *Program) GetBindDiagnostics(ctx context.Context, sourceFile *ast.SourceFile) []*ast.Diagnostic {
+	return p.getDiagnosticsHelper(ctx, sourceFile, true /*ensureBound*/, false /*ensureChecked*/, p.getBindDiagnosticsForFile)
 }
 
-func (p *Program) GetSemanticDiagnostics(sourceFile *ast.SourceFile) []*ast.Diagnostic {
-	return p.getDiagnosticsHelper(sourceFile, true /*ensureBound*/, true /*ensureChecked*/, p.getSemanticDiagnosticsForFile)
+func (p *Program) GetSemanticDiagnostics(ctx context.Context, sourceFile *ast.SourceFile) []*ast.Diagnostic {
+	return p.getDiagnosticsHelper(ctx, sourceFile, true /*ensureBound*/, true /*ensureChecked*/, p.getSemanticDiagnosticsForFile)
 }
 
 func (p *Program) GetGlobalDiagnostics() []*ast.Diagnostic {
@@ -310,11 +311,11 @@ func (p *Program) getOptionsDiagnosticsOfConfigFile() []*ast.Diagnostic {
 	return p.configFileParsingDiagnostics // TODO: actually call getDiagnosticsHelper on config path
 }
 
-func (p *Program) getSyntacticDiagnosticsForFile(sourceFile *ast.SourceFile) []*ast.Diagnostic {
+func (p *Program) getSyntacticDiagnosticsForFile(ctx context.Context, sourceFile *ast.SourceFile) []*ast.Diagnostic {
 	return sourceFile.Diagnostics()
 }
 
-func (p *Program) getBindDiagnosticsForFile(sourceFile *ast.SourceFile) []*ast.Diagnostic {
+func (p *Program) getBindDiagnosticsForFile(ctx context.Context, sourceFile *ast.SourceFile) []*ast.Diagnostic {
 	// TODO: restore this; tsgo's main depends on this function binding all files for timing.
 	// if checker.SkipTypeChecking(sourceFile, p.compilerOptions) {
 	// 	return nil
@@ -323,26 +324,27 @@ func (p *Program) getBindDiagnosticsForFile(sourceFile *ast.SourceFile) []*ast.D
 	return sourceFile.BindDiagnostics()
 }
 
-func (p *Program) getSemanticDiagnosticsForFile(sourceFile *ast.SourceFile) []*ast.Diagnostic {
+func (p *Program) getSemanticDiagnosticsForFile(ctx context.Context, sourceFile *ast.SourceFile) []*ast.Diagnostic {
 	if checker.SkipTypeChecking(sourceFile, p.compilerOptions) {
 		return nil
 	}
-
 	var fileChecker *checker.Checker
 	if sourceFile != nil {
 		fileChecker = p.GetTypeCheckerForFile(sourceFile)
 	}
-
 	diags := slices.Clip(sourceFile.BindDiagnostics())
 	// Ask for diags from all checkers; checking one file may add diagnostics to other files.
 	// These are deduplicated later.
 	for _, checker := range p.checkers {
 		if sourceFile == nil || checker == fileChecker {
-			diags = append(diags, checker.GetDiagnostics(sourceFile)...)
+			diags = append(diags, checker.GetDiagnostics(ctx, sourceFile)...)
 		} else {
 			diags = append(diags, checker.GetDiagnosticsWithoutCheck(sourceFile)...)
 		}
 	}
+	if ctx.Err() != nil {
+		return nil
+	}
 	if len(sourceFile.CommentDirectives) == 0 {
 		return diags
 	}
@@ -432,22 +434,25 @@ func compactAndMergeRelatedInfos(diagnostics []*ast.Diagnostic) []*ast.Diagnosti
 	return diagnostics[:j]
 }
 
-func (p *Program) getDiagnosticsHelper(sourceFile *ast.SourceFile, ensureBound bool, ensureChecked bool, getDiagnostics func(*ast.SourceFile) []*ast.Diagnostic) []*ast.Diagnostic {
+func (p *Program) getDiagnosticsHelper(ctx context.Context, sourceFile *ast.SourceFile, ensureBound bool, ensureChecked bool, getDiagnostics func(context.Context, *ast.SourceFile) []*ast.Diagnostic) []*ast.Diagnostic {
 	if sourceFile != nil {
 		if ensureBound {
 			binder.BindSourceFile(sourceFile, p.getSourceAffectingCompilerOptions())
 		}
-		return SortAndDeduplicateDiagnostics(getDiagnostics(sourceFile))
+		return SortAndDeduplicateDiagnostics(getDiagnostics(ctx, sourceFile))
 	}
 	if ensureBound {
 		p.BindSourceFiles()
 	}
 	if ensureChecked {
-		p.CheckSourceFiles()
+		p.CheckSourceFiles(ctx)
+		if ctx.Err() != nil {
+			return nil
+		}
 	}
 	var result []*ast.Diagnostic
 	for _, file := range p.files {
-		result = append(result, getDiagnostics(file)...)
+		result = append(result, getDiagnostics(ctx, file)...)
 	}
 	return SortAndDeduplicateDiagnostics(result)
 }

internal/execute/tsc.go

@@ -1,6 +1,7 @@
 package execute
 
 import (
+	"context"
 	"fmt"
 
 	"github.com/microsoft/typescript-go/internal/ast"
@@ -205,7 +206,7 @@ func compileAndEmit(sys System, program *compiler.Program, reportDiagnostic diag
 	allDiagnostics := program.GetConfigFileParsingDiagnostics()
 
 	// todo: early exit logic and append diagnostics
-	diagnostics := program.GetSyntacticDiagnostics(nil)
+	diagnostics := program.GetSyntacticDiagnostics(context.Background(), nil)
 	if len(diagnostics) == 0 {
 		diagnostics = append(diagnostics, program.GetOptionsDiagnostics()...)
 		if options.ListFilesOnly.IsFalse() {
@@ -214,7 +215,7 @@ func compileAndEmit(sys System, program *compiler.Program, reportDiagnostic diag
 		}
 	}
 	if len(diagnostics) == 0 {
-		diagnostics = append(diagnostics, program.GetSemanticDiagnostics(nil)...)
+		diagnostics = append(diagnostics, program.GetSemanticDiagnostics(context.Background(), nil)...)
 	}
 	// TODO: declaration diagnostics
 	if len(diagnostics) == 0 && options.NoEmit == core.TSTrue && (options.Declaration.IsTrue() && options.Composite.IsTrue()) {