Make CheckerPool iteration concurrent by default

internal/compiler/checkerpool.go

@@ -12,10 +12,11 @@ import (
 )
 
 type CheckerPool interface {
+	Count() int
 	GetChecker(ctx context.Context) (*checker.Checker, func())
 	GetCheckerForFile(ctx context.Context, file *ast.SourceFile) (*checker.Checker, func())
 	GetCheckerForFileExclusive(ctx context.Context, file *ast.SourceFile) (*checker.Checker, func())
-	GetAllCheckers(ctx context.Context) ([]*checker.Checker, func())
+	ForEachCheckerParallel(ctx context.Context, cb func(idx int, c *checker.Checker))
 	Files(checker *checker.Checker) iter.Seq[*ast.SourceFile]
 }
 
@@ -42,6 +43,10 @@ func newCheckerPool(checkerCount int, program *Program) *checkerPool {
 	return pool
 }
 
+func (p *checkerPool) Count() int {
+	return p.checkerCount
+}
+
 func (p *checkerPool) GetCheckerForFile(ctx context.Context, file *ast.SourceFile) (*checker.Checker, func()) {
 	p.createCheckers()
 	checker := p.fileAssociations[file]
@@ -82,9 +87,19 @@ func (p *checkerPool) createCheckers() {
 	})
 }
 
-func (p *checkerPool) GetAllCheckers(ctx context.Context) ([]*checker.Checker, func()) {
+// Runs `cb` for each checker in the pool concurrently, locking and unlocking checker mutexes as it goes,
+// making it safe to call `ForEachCheckerParallel` from many threads simultaneously.
+func (p *checkerPool) ForEachCheckerParallel(ctx context.Context, cb func(idx int, c *checker.Checker)) {
 	p.createCheckers()
-	return p.checkers, noop
+	wg := core.NewWorkGroup(p.program.SingleThreaded())
+	for idx, checker := range p.checkers {
+		wg.Queue(func() {
+			p.locks[idx].Lock()
+			defer p.locks[idx].Unlock()
+			cb(idx, checker)
+		})
+	}
+	wg.RunAndWait()
 }
 
 func (p *checkerPool) Files(checker *checker.Checker) iter.Seq[*ast.SourceFile] {

internal/compiler/program.go

@@ -8,6 +8,7 @@ import (
 	"slices"
 	"strings"
 	"sync"
+	"sync/atomic"
 
 	"github.com/go-json-experiment/json"
 	"github.com/microsoft/typescript-go/internal/ast"
@@ -351,28 +352,22 @@ func (p *Program) BindSourceFiles() {
 }
 
 func (p *Program) CheckSourceFiles(ctx context.Context, files []*ast.SourceFile) {
-	wg := core.NewWorkGroup(p.SingleThreaded())
-	checkers, done := p.checkerPool.GetAllCheckers(ctx)
-	defer done()
-	for _, checker := range checkers {
-		wg.Queue(func() {
-			for file := range p.checkerPool.Files(checker) {
-				if files == nil || slices.Contains(files, file) {
-					checker.CheckSourceFile(ctx, file)
-				}
+	p.checkerPool.ForEachCheckerParallel(ctx, func(_ int, checker *checker.Checker) {
+		for file := range p.checkerPool.Files(checker) {
+			if files == nil || slices.Contains(files, file) {
+				checker.CheckSourceFile(ctx, file)
 			}
-		})
-	}
-	wg.RunAndWait()
+		}
+	})
 }
 
 // Return the type checker associated with the program.
 func (p *Program) GetTypeChecker(ctx context.Context) (*checker.Checker, func()) {
 	return p.checkerPool.GetChecker(ctx)
 }
 
-func (p *Program) GetTypeCheckers(ctx context.Context) ([]*checker.Checker, func()) {
-	return p.checkerPool.GetAllCheckers(ctx)
+func (p *Program) ForEachCheckerParallel(ctx context.Context, cb func(idx int, c *checker.Checker)) {
+	p.checkerPool.ForEachCheckerParallel(ctx, cb)
 }
 
 // Return a checker for the given file. We may have multiple checkers in concurrent scenarios and this
@@ -965,14 +960,12 @@ func (p *Program) GetGlobalDiagnostics(ctx context.Context) []*ast.Diagnostic {
 		return nil
 	}
 
-	var globalDiagnostics []*ast.Diagnostic
-	checkers, done := p.checkerPool.GetAllCheckers(ctx)
-	defer done()
-	for _, checker := range checkers {
-		globalDiagnostics = append(globalDiagnostics, checker.GetGlobalDiagnostics()...)
-	}
+	globalDiagnostics := make([][]*ast.Diagnostic, p.checkerPool.Count())
+	p.checkerPool.ForEachCheckerParallel(ctx, func(idx int, checker *checker.Checker) {
+		globalDiagnostics[idx] = checker.GetGlobalDiagnostics()
+	})
 
-	return SortAndDeduplicateDiagnostics(globalDiagnostics)
+	return SortAndDeduplicateDiagnostics(slices.Concat(globalDiagnostics...))
 }
 
 func (p *Program) GetDeclarationDiagnostics(ctx context.Context, sourceFile *ast.SourceFile) []*ast.Diagnostic {
@@ -1033,22 +1026,23 @@ func (p *Program) getSemanticDiagnosticsForFileNotFilter(ctx context.Context, so
 		defer done()
 	}
 	diags := slices.Clip(sourceFile.BindDiagnostics())
-	checkers, closeCheckers := p.checkerPool.GetAllCheckers(ctx)
-	defer closeCheckers()
 
 	// Ask for diags from all checkers; checking one file may add diagnostics to other files.
 	// These are deduplicated later.
-	for _, checker := range checkers {
+	checkerDiags := make([][]*ast.Diagnostic, p.checkerPool.Count())
+	p.checkerPool.ForEachCheckerParallel(ctx, func(idx int, checker *checker.Checker) {
 		if sourceFile == nil || checker == fileChecker {
-			diags = append(diags, checker.GetDiagnostics(ctx, sourceFile)...)
+			checkerDiags[idx] = checker.GetDiagnostics(ctx, sourceFile)
 		} else {
-			diags = append(diags, checker.GetDiagnosticsWithoutCheck(sourceFile)...)
+			checkerDiags[idx] = checker.GetDiagnosticsWithoutCheck(sourceFile)
 		}
-	}
+	})
 	if ctx.Err() != nil {
 		return nil
 	}
 
+	diags = append(diags, slices.Concat(checkerDiags...)...)
+
 	// !!! This should be rewritten to work like getBindAndCheckDiagnosticsForFileNoCache.
 
 	isPlainJS := ast.IsPlainJSFile(sourceFile, compilerOptions.CheckJs)
@@ -1140,22 +1134,20 @@ func (p *Program) getSuggestionDiagnosticsForFile(ctx context.Context, sourceFil
 
 	diags := slices.Clip(sourceFile.BindSuggestionDiagnostics)
 
-	checkers, closeCheckers := p.checkerPool.GetAllCheckers(ctx)
-	defer closeCheckers()
-
-	// Ask for diags from all checkers; checking one file may add diagnostics to other files.
-	// These are deduplicated later.
-	for _, checker := range checkers {
+	checkerDiags := make([][]*ast.Diagnostic, p.checkerPool.Count())
+	p.checkerPool.ForEachCheckerParallel(ctx, func(idx int, checker *checker.Checker) {
 		if sourceFile == nil || checker == fileChecker {
-			diags = append(diags, checker.GetSuggestionDiagnostics(ctx, sourceFile)...)
+			checkerDiags[idx] = checker.GetSuggestionDiagnostics(ctx, sourceFile)
 		} else {
 			// !!! is there any case where suggestion diagnostics are produced in other checkers?
 		}
-	}
+	})
 	if ctx.Err() != nil {
 		return nil
 	}
 
+	diags = append(diags, slices.Concat(checkerDiags...)...)
+
 	return diags
 }
 
@@ -1251,32 +1243,28 @@ func (p *Program) SymbolCount() int {
 	for _, file := range p.files {
 		count += file.SymbolCount
 	}
-	checkers, done := p.checkerPool.GetAllCheckers(context.Background())
-	defer done()
-	for _, checker := range checkers {
-		count += int(checker.SymbolCount)
-	}
-	return count
+	var val atomic.Uint32
+	val.Store(uint32(count))
+	p.checkerPool.ForEachCheckerParallel(context.Background(), func(idx int, c *checker.Checker) {
+		val.Add(c.SymbolCount)
+	})
+	return int(val.Load())
 }
 
 func (p *Program) TypeCount() int {
-	var count int
-	checkers, done := p.checkerPool.GetAllCheckers(context.Background())
-	defer done()
-	for _, checker := range checkers {
-		count += int(checker.TypeCount)
-	}
-	return count
+	var val atomic.Uint32
+	p.checkerPool.ForEachCheckerParallel(context.Background(), func(idx int, c *checker.Checker) {
+		val.Add(c.TypeCount)
+	})
+	return int(val.Load())
 }
 
 func (p *Program) InstantiationCount() int {
-	var count int
-	checkers, done := p.checkerPool.GetAllCheckers(context.Background())
-	defer done()
-	for _, checker := range checkers {
-		count += int(checker.TotalInstantiationCount)
-	}
-	return count
+	var val atomic.Uint32
+	p.checkerPool.ForEachCheckerParallel(context.Background(), func(idx int, c *checker.Checker) {
+		val.Add(c.TotalInstantiationCount)
+	})
+	return int(val.Load())
 }
 
 func (p *Program) Program() *Program {

internal/project/checkerpool.go

@@ -94,22 +94,32 @@ func (p *CheckerPool) Files(checker *checker.Checker) iter.Seq[*ast.SourceFile]
 	panic("unimplemented")
 }
 
-func (p *CheckerPool) GetAllCheckers(ctx context.Context) ([]*checker.Checker, func()) {
+func (p *CheckerPool) Count() int {
+	return p.maxCheckers
+}
+
+func (p *CheckerPool) ForEachCheckerParallel(ctx context.Context, cb func(idx int, c *checker.Checker)) {
 	p.mu.Lock()
 	defer p.mu.Unlock()
 
 	requestID := core.GetRequestID(ctx)
 	if requestID == "" {
-		panic("cannot call GetAllCheckers on a project.checkerPool without a request ID")
+		panic("cannot call ForEachCheckerParallel on a project.checkerPool without a request ID")
 	}
 
 	// A request can only access one checker
 	if c, release := p.getRequestCheckerLocked(requestID); c != nil {
-		return []*checker.Checker{c}, release
+		defer release()
+		cb(0, c)
+		return
 	}
 
+	// TODO: Does this ever work without deadlocking? `p.GetChecker` also tries to lock this mutex.
+	// Should this just be a panic?
 	c, release := p.GetChecker(ctx)
-	return []*checker.Checker{c}, release
+	defer release()
+	cb(0, c)
+	return
 }
 
 func (p *CheckerPool) getCheckerLocked(requestID string) (*checker.Checker, int) {

internal/testrunner/compiler_runner.go

@@ -529,9 +529,7 @@ func createHarnessTestFile(unit *testUnit, currentDirectory string) *harnessutil
 func (c *compilerTest) verifyUnionOrdering(t *testing.T) {
 	t.Run("union ordering", func(t *testing.T) {
 		p := c.result.Program.Program()
-		checkers, done := p.GetTypeCheckers(t.Context())
-		defer done()
-		for _, c := range checkers {
+		p.ForEachCheckerParallel(t.Context(), func(_ int, c *checker.Checker) {
 			for union := range c.UnionTypes() {
 				types := union.Types()
 
@@ -549,7 +547,7 @@ func (c *compilerTest) verifyUnionOrdering(t *testing.T) {
 					assert.Assert(t, slices.Equal(shuffled, types), "compareTypes does not sort union types consistently")
 				}
 			}
-		}
+		})
 	})
 }