Compute the snapshot changes in parallel (#1475)

Author: sheetalkamat

internal/collections/manytomanyset.go

@@ -0,0 +1,34 @@
+package collections
+
+import "fmt"
+
+type SyncManyToManySet[K comparable, V comparable] struct {
+	keyToValueSet SyncMap[K, *Set[V]]
+	valueToKeySet SyncMap[V, *SyncSet[K]]
+}
+
+func (m *SyncManyToManySet[K, V]) GetKeys(value V) (*SyncSet[K], bool) {
+	keys, present := m.valueToKeySet.Load(value)
+	return keys, present
+}
+
+func (m *SyncManyToManySet[K, V]) GetValues(key K) (*Set[V], bool) {
+	values, present := m.keyToValueSet.Load(key)
+	return values, present
+}
+
+func (m *SyncManyToManySet[K, V]) Keys() *SyncMap[K, *Set[V]] {
+	return &m.keyToValueSet
+}
+
+func (m *SyncManyToManySet[K, V]) Store(key K, valueSet *Set[V]) {
+	_, hasExisting := m.keyToValueSet.LoadOrStore(key, valueSet)
+	if hasExisting {
+		panic("ManyToManySet.Set: key already exists: " + fmt.Sprintf("%v", key))
+	}
+	for value := range valueSet.Keys() {
+		// Add to valueToKeySet
+		keySetForValue, _ := m.valueToKeySet.LoadOrStore(value, &SyncSet[K]{})
+		keySetForValue.Add(key)
+	}
+}

internal/collections/syncmanytomanyset.go

@@ -1,6 +1,9 @@
 package collections
 
-import "sync"
+import (
+	"iter"
+	"sync"
+)
 
 type SyncMap[K comparable, V any] struct {
 	m sync.Map
@@ -57,3 +60,14 @@ func (s *SyncMap[K, V]) ToMap() map[K]V {
 	})
 	return m
 }
+
+func (s *SyncMap[K, V]) Keys() iter.Seq[K] {
+	return func(yield func(K) bool) {
+		s.m.Range(func(key, value any) bool {
+			if !yield(key.(K)) {
+				return false
+			}
+			return true
+		})
+	}
+}

internal/collections/syncmap.go

@@ -1,5 +1,7 @@
 package collections
 
+import "iter"
+
 type SyncSet[T comparable] struct {
 	m SyncMap[T, struct{}]
 }
@@ -23,6 +25,18 @@ func (s *SyncSet[T]) Range(fn func(key T) bool) {
 	})
 }
 
+// Size returns the approximate number of items in the map.
+// Note that this is not a precise count, as the map may be modified
+// concurrently while this method is running.
+func (s *SyncSet[T]) Size() int {
+	count := 0
+	s.m.Range(func(_ T, _ struct{}) bool {
+		count++
+		return true
+	})
+	return count
+}
+
 func (s *SyncSet[T]) ToSlice() []T {
 	var arr []T
 	arr = make([]T, 0, s.m.Size())
@@ -32,3 +46,14 @@ func (s *SyncSet[T]) ToSlice() []T {
 	})
 	return arr
 }
+
+func (s *SyncSet[T]) Keys() iter.Seq[T] {
+	return func(yield func(T) bool) {
+		s.m.Range(func(key T, value struct{}) bool {
+			if !yield(key) {
+				return false
+			}
+			return true
+		})
+	}
+}

internal/collections/syncset.go

@@ -183,8 +183,8 @@ func (s *testSys) baselineProgram(baseline *strings.Builder, program *incrementa
 	baseline.WriteString("SemanticDiagnostics::\n")
 	testingData := program.GetTestingData(program.GetProgram())
 	for _, file := range program.GetProgram().GetSourceFiles() {
-		if diagnostics, ok := testingData.SemanticDiagnosticsPerFile[file.Path()]; ok {
-			if oldDiagnostics, ok := testingData.OldProgramSemanticDiagnosticsPerFile[file.Path()]; !ok || oldDiagnostics != diagnostics {
+		if diagnostics, ok := testingData.SemanticDiagnosticsPerFile.Load(file.Path()); ok {
+			if oldDiagnostics, ok := testingData.OldProgramSemanticDiagnosticsPerFile.Load(file.Path()); !ok || oldDiagnostics != diagnostics {
 				baseline.WriteString("*refresh*    " + file.FileName() + "\n")
 			}
 		} else {

internal/execute/testsys_test.go

@@ -2,6 +2,7 @@ package incremental
 
 import (
 	"context"
+	"iter"
 	"maps"
 	"slices"
 	"sync"
@@ -41,8 +42,8 @@ func (h *affectedFilesHandler) getDtsMayChange(affectedFilePath tspath.Path, aff
 
 func (h *affectedFilesHandler) isChangedSignature(path tspath.Path) bool {
 	newSignature, _ := h.updatedSignatures.Load(path)
-	oldSignature := h.program.snapshot.fileInfos[path].signature
-	return newSignature != oldSignature
+	oldInfo, _ := h.program.snapshot.fileInfos.Load(path)
+	return newSignature != oldInfo.signature
 }
 
 func (h *affectedFilesHandler) removeSemanticDiagnosticsOf(path tspath.Path) {
@@ -81,7 +82,7 @@ func (h *affectedFilesHandler) updateShapeSignature(file *ast.SourceFile, useFil
 		return false
 	}
 
-	info := h.program.snapshot.fileInfos[file.Path()]
+	info, _ := h.program.snapshot.fileInfos.Load(file.Path())
 	prevSignature := info.signature
 	var latestSignature string
 	var updateKind SignatureUpdateKind
@@ -110,7 +111,7 @@ func (h *affectedFilesHandler) getFilesAffectedBy(path tspath.Path) []*ast.Sourc
 		return []*ast.SourceFile{file}
 	}
 
-	if info := h.program.snapshot.fileInfos[file.Path()]; info.affectsGlobalScope {
+	if info, _ := h.program.snapshot.fileInfos.Load(file.Path()); info.affectsGlobalScope {
 		h.hasAllFilesExcludingDefaultLibraryFile.Store(true)
 		h.program.snapshot.getAllFilesExcludingDefaultLibraryFile(h.program.program, file)
 	}
@@ -139,10 +140,10 @@ func (h *affectedFilesHandler) getFilesAffectedBy(path tspath.Path) []*ast.Sourc
 }
 
 // Gets the files referenced by the the file path
-func (h *affectedFilesHandler) getReferencedByPaths(file tspath.Path) map[tspath.Path]struct{} {
+func (h *affectedFilesHandler) getReferencedByPaths(file tspath.Path) iter.Seq[tspath.Path] {
 	keys, ok := h.program.snapshot.referencedMap.GetKeys(file)
 	if !ok {
-		return nil
+		return func(yield func(tspath.Path) bool) {}
 	}
 	return keys.Keys()
 }
@@ -154,8 +155,7 @@ func (h *affectedFilesHandler) forEachFileReferencedBy(file *ast.SourceFile, fn
 	seenFileNamesMap := map[tspath.Path]*ast.SourceFile{}
 	// Start with the paths this file was referenced by
 	seenFileNamesMap[file.Path()] = file
-	references := h.getReferencedByPaths(file.Path())
-	queue := slices.Collect(maps.Keys(references))
+	queue := slices.Collect(h.getReferencedByPaths(file.Path()))
 	for len(queue) > 0 {
 		currentPath := queue[len(queue)-1]
 		queue = queue[:len(queue)-1]
@@ -290,7 +290,7 @@ func (h *affectedFilesHandler) handleDtsMayChangeOfFileAndExportsOfFile(dtsMayCh
 }
 
 func (h *affectedFilesHandler) handleDtsMayChangeOfGlobalScope(dtsMayChange dtsMayChange, filePath tspath.Path, invalidateJsFiles bool) bool {
-	if info, ok := h.program.snapshot.fileInfos[filePath]; !ok || !info.affectsGlobalScope {
+	if info, ok := h.program.snapshot.fileInfos.Load(filePath); !ok || !info.affectsGlobalScope {
 		return false
 	}
 	// Every file needs to be handled
@@ -331,7 +331,9 @@ func (h *affectedFilesHandler) updateSnapshot() {
 		return
 	}
 	h.updatedSignatures.Range(func(filePath tspath.Path, signature string) bool {
-		h.program.snapshot.fileInfos[filePath].signature = signature
+		if info, ok := h.program.snapshot.fileInfos.Load(filePath); ok {
+			info.signature = signature
+		}
 		return true
 	})
 	if h.updatedSignatureKinds != nil {
@@ -341,33 +343,34 @@ func (h *affectedFilesHandler) updateSnapshot() {
 		})
 	}
 	h.filesToRemoveDiagnostics.Range(func(file tspath.Path) bool {
-		delete(h.program.snapshot.semanticDiagnosticsPerFile, file)
+		h.program.snapshot.semanticDiagnosticsPerFile.Delete(file)
 		return true
 	})
 	for _, change := range h.dtsMayChange {
 		for filePath, emitKind := range change {
 			h.program.snapshot.addFileToAffectedFilesPendingEmit(filePath, emitKind)
 		}
 	}
-	h.program.snapshot.changedFilesSet = &collections.Set[tspath.Path]{}
-	h.program.snapshot.buildInfoEmitPending = true
+	h.program.snapshot.changedFilesSet = collections.SyncSet[tspath.Path]{}
+	h.program.snapshot.buildInfoEmitPending.Store(true)
 }
 
 func collectAllAffectedFiles(ctx context.Context, program *Program) {
-	if program.snapshot.changedFilesSet.Len() == 0 {
+	if program.snapshot.changedFilesSet.Size() == 0 {
 		return
 	}
 
 	handler := affectedFilesHandler{ctx: ctx, program: program, updatedSignatureKinds: core.IfElse(program.updatedSignatureKinds == nil, nil, &collections.SyncMap[tspath.Path, SignatureUpdateKind]{})}
 	wg := core.NewWorkGroup(handler.program.program.SingleThreaded())
 	var result collections.SyncSet[*ast.SourceFile]
-	for file := range program.snapshot.changedFilesSet.Keys() {
+	program.snapshot.changedFilesSet.Range(func(file tspath.Path) bool {
 		wg.Queue(func() {
 			for _, affectedFile := range handler.getFilesAffectedBy(file) {
 				result.Add(affectedFile)
 			}
 		})
-	}
+		return true
+	})
 	wg.RunAndWait()
 
 	if ctx.Err() != nil {

internal/incremental/affectedfileshandler.go

@@ -328,12 +328,13 @@ func (b *BuildInfoEmitSignature) noEmitSignature() bool {
 	return b.Signature == "" && !b.DiffersOnlyInDtsMap && !b.DiffersInOptions
 }
 
-func (b *BuildInfoEmitSignature) toEmitSignature(path tspath.Path, emitSignatures map[tspath.Path]*emitSignature) *emitSignature {
+func (b *BuildInfoEmitSignature) toEmitSignature(path tspath.Path, emitSignatures *collections.SyncMap[tspath.Path, *emitSignature]) *emitSignature {
 	var signature string
 	var signatureWithDifferentOptions []string
 	if b.DiffersOnlyInDtsMap {
 		signatureWithDifferentOptions = make([]string, 0, 1)
-		signatureWithDifferentOptions = append(signatureWithDifferentOptions, emitSignatures[path].signature)
+		info, _ := emitSignatures.Load(path)
+		signatureWithDifferentOptions = append(signatureWithDifferentOptions, info.signature)
 	} else if b.DiffersInOptions {
 		signatureWithDifferentOptions = make([]string, 0, 1)
 		signatureWithDifferentOptions = append(signatureWithDifferentOptions, b.Signature)

internal/incremental/buildInfo.go

@@ -93,66 +93,63 @@ func (t *toSnapshot) setCompilerOptions() {
 }
 
 func (t *toSnapshot) setFileInfoAndEmitSignatures() {
-	t.snapshot.fileInfos = make(map[tspath.Path]*fileInfo, len(t.buildInfo.FileInfos))
-	t.snapshot.createEmitSignaturesMap()
+	isComposite := t.snapshot.options.Composite.IsTrue()
 	for index, buildInfoFileInfo := range t.buildInfo.FileInfos {
 		path := t.toFilePath(BuildInfoFileId(index + 1))
 		info := buildInfoFileInfo.GetFileInfo()
-		t.snapshot.fileInfos[path] = info
+		t.snapshot.fileInfos.Store(path, info)
 		// Add default emit signature as file's signature
-		if info.signature != "" && t.snapshot.emitSignatures != nil {
-			t.snapshot.emitSignatures[path] = &emitSignature{signature: info.signature}
+		if info.signature != "" && isComposite {
+			t.snapshot.emitSignatures.Store(path, &emitSignature{signature: info.signature})
 		}
 	}
 	// Fix up emit signatures
 	for _, value := range t.buildInfo.EmitSignatures {
 		if value.noEmitSignature() {
-			delete(t.snapshot.emitSignatures, t.toFilePath(value.FileId))
+			t.snapshot.emitSignatures.Delete(t.toFilePath(value.FileId))
 		} else {
 			path := t.toFilePath(value.FileId)
-			t.snapshot.emitSignatures[path] = value.toEmitSignature(path, t.snapshot.emitSignatures)
+			t.snapshot.emitSignatures.Store(path, value.toEmitSignature(path, &t.snapshot.emitSignatures))
 		}
 	}
 }
 
 func (t *toSnapshot) setReferencedMap() {
 	for _, entry := range t.buildInfo.ReferencedMap {
-		t.snapshot.referencedMap.Set(t.toFilePath(entry.FileId), t.toFilePathSet(entry.FileIdListId))
+		t.snapshot.referencedMap.Store(t.toFilePath(entry.FileId), t.toFilePathSet(entry.FileIdListId))
 	}
 }
 
 func (t *toSnapshot) setChangeFileSet() {
-	t.snapshot.changedFilesSet = collections.NewSetWithSizeHint[tspath.Path](len(t.buildInfo.ChangeFileSet))
 	for _, fileId := range t.buildInfo.ChangeFileSet {
 		filePath := t.toFilePath(fileId)
 		t.snapshot.changedFilesSet.Add(filePath)
 	}
 }
 
 func (t *toSnapshot) setSemanticDiagnostics() {
-	t.snapshot.semanticDiagnosticsPerFile = make(map[tspath.Path]*diagnosticsOrBuildInfoDiagnosticsWithFileName, len(t.snapshot.fileInfos))
-	for path := range t.snapshot.fileInfos {
+	t.snapshot.fileInfos.Range(func(path tspath.Path, info *fileInfo) bool {
 		// Initialize to have no diagnostics if its not changed file
 		if !t.snapshot.changedFilesSet.Has(path) {
-			t.snapshot.semanticDiagnosticsPerFile[path] = &diagnosticsOrBuildInfoDiagnosticsWithFileName{}
+			t.snapshot.semanticDiagnosticsPerFile.Store(path, &diagnosticsOrBuildInfoDiagnosticsWithFileName{})
 		}
-	}
+		return true
+	})
 	for _, diagnostic := range t.buildInfo.SemanticDiagnosticsPerFile {
 		if diagnostic.FileId != 0 {
 			filePath := t.toFilePath(diagnostic.FileId)
-			delete(t.snapshot.semanticDiagnosticsPerFile, filePath) // does not have cached diagnostics
+			t.snapshot.semanticDiagnosticsPerFile.Delete(filePath) // does not have cached diagnostics
 		} else {
 			filePath := t.toFilePath(diagnostic.Diagnostics.FileId)
-			t.snapshot.semanticDiagnosticsPerFile[filePath] = t.toDiagnosticsOrBuildInfoDiagnosticsWithFileName(diagnostic.Diagnostics)
+			t.snapshot.semanticDiagnosticsPerFile.Store(filePath, t.toDiagnosticsOrBuildInfoDiagnosticsWithFileName(diagnostic.Diagnostics))
 		}
 	}
 }
 
 func (t *toSnapshot) setEmitDiagnostics() {
-	t.snapshot.emitDiagnosticsPerFile = make(map[tspath.Path]*diagnosticsOrBuildInfoDiagnosticsWithFileName, len(t.snapshot.fileInfos))
 	for _, diagnostic := range t.buildInfo.EmitDiagnosticsPerFile {
 		filePath := t.toFilePath(diagnostic.FileId)
-		t.snapshot.emitDiagnosticsPerFile[filePath] = t.toDiagnosticsOrBuildInfoDiagnosticsWithFileName(diagnostic)
+		t.snapshot.emitDiagnosticsPerFile.Store(filePath, t.toDiagnosticsOrBuildInfoDiagnosticsWithFileName(diagnostic))
 	}
 }
 
@@ -161,8 +158,7 @@ func (t *toSnapshot) setAffectedFilesPendingEmit() {
 		return
 	}
 	ownOptionsEmitKind := GetFileEmitKind(t.snapshot.options)
-	t.snapshot.affectedFilesPendingEmit = make(map[tspath.Path]FileEmitKind, len(t.buildInfo.AffectedFilesPendingEmit))
 	for _, pendingEmit := range t.buildInfo.AffectedFilesPendingEmit {
-		t.snapshot.affectedFilesPendingEmit[t.toFilePath(pendingEmit.FileId)] = core.IfElse(pendingEmit.EmitKind == 0, ownOptionsEmitKind, pendingEmit.EmitKind)
+		t.snapshot.affectedFilesPendingEmit.Store(t.toFilePath(pendingEmit.FileId), core.IfElse(pendingEmit.EmitKind == 0, ownOptionsEmitKind, pendingEmit.EmitKind))
 	}
 }