Compute referencedBy when its needed

internal/incremental/affectedfileshandler.go

@@ -2,7 +2,6 @@ package incremental
 
 import (
 	"context"
-	"iter"
 	"maps"
 	"slices"
 	"sync"
@@ -139,23 +138,14 @@ 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) iter.Seq[tspath.Path] {
-	keys, ok := h.program.snapshot.referencedMap.GetKeys(file)
-	if !ok {
-		return func(yield func(tspath.Path) bool) {}
-	}
-	return keys.Keys()
-}
-
 func (h *affectedFilesHandler) forEachFileReferencedBy(file *ast.SourceFile, fn func(currentFile *ast.SourceFile, currentPath tspath.Path) (queueForFile bool, fastReturn bool)) map[tspath.Path]*ast.SourceFile {
 	// Now we need to if each file in the referencedBy list has a shape change as well.
 	// Because if so, its own referencedBy files need to be saved as well to make the
 	// emitting result consistent with files on disk.
 	seenFileNamesMap := map[tspath.Path]*ast.SourceFile{}
 	// Start with the paths this file was referenced by
 	seenFileNamesMap[file.Path()] = file
-	queue := slices.Collect(h.getReferencedByPaths(file.Path()))
+	queue := slices.Collect(h.program.snapshot.referencedMap.getReferencedBy(file.Path()))
 	for len(queue) > 0 {
 		currentPath := queue[len(queue)-1]
 		queue = queue[:len(queue)-1]
@@ -167,7 +157,7 @@ func (h *affectedFilesHandler) forEachFileReferencedBy(file *ast.SourceFile, fn
 				return seenFileNamesMap
 			}
 			if queueForFile {
-				for ref := range h.getReferencedByPaths(currentFile.Path()) {
+				for ref := range h.program.snapshot.referencedMap.getReferencedBy(currentFile.Path()) {
 					queue = append(queue, ref)
 				}
 			}
@@ -253,18 +243,14 @@ func (h *affectedFilesHandler) handleDtsMayChangeOfAffectedFile(dtsMayChange dts
 	}
 
 	// Go through files that reference affected file and handle dts emit and semantic diagnostics for them and their references
-	if keys, ok := h.program.snapshot.referencedMap.GetKeys(affectedFile.Path()); ok {
-		for exportedFromPath := range keys.Keys() {
-			if h.handleDtsMayChangeOfGlobalScope(dtsMayChange, exportedFromPath, invalidateJsFiles) {
+	for exportedFromPath := range h.program.snapshot.referencedMap.getReferencedBy(affectedFile.Path()) {
+		if h.handleDtsMayChangeOfGlobalScope(dtsMayChange, exportedFromPath, invalidateJsFiles) {
+			return
+		}
+		for filePath := range h.program.snapshot.referencedMap.getReferencedBy(exportedFromPath) {
+			if h.handleDtsMayChangeOfFileAndExportsOfFile(dtsMayChange, filePath, invalidateJsFiles) {
 				return
 			}
-			if references, ok := h.program.snapshot.referencedMap.GetKeys(exportedFromPath); ok {
-				for filePath := range references.Keys() {
-					if h.handleDtsMayChangeOfFileAndExportsOfFile(dtsMayChange, filePath, invalidateJsFiles) {
-						return
-					}
-				}
-			}
 		}
 	}
 }
@@ -279,11 +265,9 @@ func (h *affectedFilesHandler) handleDtsMayChangeOfFileAndExportsOfFile(dtsMayCh
 	h.handleDtsMayChangeOf(dtsMayChange, filePath, invalidateJsFiles)
 
 	// Remove the diagnostics of files that import this file and handle all its exports too
-	if keys, ok := h.program.snapshot.referencedMap.GetKeys(filePath); ok {
-		for referencingFilePath := range keys.Keys() {
-			if h.handleDtsMayChangeOfFileAndExportsOfFile(dtsMayChange, referencingFilePath, invalidateJsFiles) {
-				return true
-			}
+	for referencingFilePath := range h.program.snapshot.referencedMap.getReferencedBy(filePath) {
+		if h.handleDtsMayChangeOfFileAndExportsOfFile(dtsMayChange, referencingFilePath, invalidateJsFiles) {
+			return true
 		}
 	}
 	return false

internal/incremental/buildinfotosnapshot.go

@@ -122,7 +122,7 @@ func (t *toSnapshot) setFileInfoAndEmitSignatures() {
 
 func (t *toSnapshot) setReferencedMap() {
 	for _, entry := range t.buildInfo.ReferencedMap {
-		t.snapshot.referencedMap.Store(t.toFilePath(entry.FileId), t.toFilePathSet(entry.FileIdListId))
+		t.snapshot.referencedMap.storeReferences(t.toFilePath(entry.FileId), t.toFilePathSet(entry.FileIdListId))
 	}
 }
 

internal/incremental/programtosnapshot.go

@@ -93,14 +93,14 @@ func (t *toProgramSnapshot) computeProgramFileChanges() {
 			var signature string
 			newReferences := getReferencedFiles(t.program, file)
 			if newReferences != nil {
-				t.snapshot.referencedMap.Store(file.Path(), newReferences)
+				t.snapshot.referencedMap.storeReferences(file.Path(), newReferences)
 			}
 			if t.oldProgram != nil {
 				if oldFileInfo, ok := t.oldProgram.snapshot.fileInfos.Load(file.Path()); ok {
 					signature = oldFileInfo.signature
 					if oldFileInfo.version != version || oldFileInfo.affectsGlobalScope != affectsGlobalScope || oldFileInfo.impliedNodeFormat != impliedNodeFormat {
 						t.snapshot.addFileToChangeSet(file.Path())
-					} else if oldReferences, _ := t.oldProgram.snapshot.referencedMap.GetValues(file.Path()); !newReferences.Equals(oldReferences) {
+					} else if oldReferences, _ := t.oldProgram.snapshot.referencedMap.getReferences(file.Path()); !newReferences.Equals(oldReferences) {
 						// Referenced files changed
 						t.snapshot.addFileToChangeSet(file.Path())
 					} else if newReferences != nil {

internal/incremental/referencemap.go

@@ -0,0 +1,52 @@
+package incremental
+
+import (
+	"iter"
+	"maps"
+	"slices"
+	"sync"
+
+	"github.com/microsoft/typescript-go/internal/collections"
+	"github.com/microsoft/typescript-go/internal/tspath"
+)
+
+type referenceMap struct {
+	references   collections.SyncMap[tspath.Path, *collections.Set[tspath.Path]]
+	referencedBy map[tspath.Path]*collections.Set[tspath.Path]
+	referenceBy  sync.Once
+}
+
+func (r *referenceMap) storeReferences(path tspath.Path, refs *collections.Set[tspath.Path]) {
+	r.references.Store(path, refs)
+}
+
+func (r *referenceMap) getReferences(path tspath.Path) (*collections.Set[tspath.Path], bool) {
+	refs, ok := r.references.Load(path)
+	return refs, ok
+}
+
+func (r *referenceMap) getPathsWithReferences() []tspath.Path {
+	return slices.Collect(r.references.Keys())
+}
+
+func (r *referenceMap) getReferencedBy(path tspath.Path) iter.Seq[tspath.Path] {
+	r.referenceBy.Do(func() {
+		r.referencedBy = make(map[tspath.Path]*collections.Set[tspath.Path])
+		r.references.Range(func(key tspath.Path, value *collections.Set[tspath.Path]) bool {
+			for ref := range value.Keys() {
+				set, ok := r.referencedBy[ref]
+				if !ok {
+					set = &collections.Set[tspath.Path]{}
+					r.referencedBy[ref] = set
+				}
+				set.Add(key)
+			}
+			return true
+		})
+	})
+	refs, ok := r.referencedBy[path]
+	if ok {
+		return maps.Keys(refs.Keys())
+	}
+	return func(yield func(tspath.Path) bool) {}
+}

internal/incremental/snapshot.go

@@ -188,7 +188,7 @@ type snapshot struct {
 	fileInfos collections.SyncMap[tspath.Path, *fileInfo]
 	options   *core.CompilerOptions
 	//  Contains the map of ReferencedSet=Referenced files of the file if module emit is enabled
-	referencedMap collections.SyncManyToManySet[tspath.Path, tspath.Path]
+	referencedMap referenceMap
 	// Cache of semantic diagnostics for files with their Path being the key
 	semanticDiagnosticsPerFile collections.SyncMap[tspath.Path, *diagnosticsOrBuildInfoDiagnosticsWithFileName]
 	// Cache of dts emit diagnostics for files with their Path being the key

internal/incremental/snapshottobuildinfo.go

@@ -231,10 +231,10 @@ func (t *toBuildInfo) setCompilerOptions() {
 }
 
 func (t *toBuildInfo) setReferencedMap() {
-	keys := slices.Collect(t.snapshot.referencedMap.Keys().Keys())
+	keys := t.snapshot.referencedMap.getPathsWithReferences()
 	slices.Sort(keys)
 	t.buildInfo.ReferencedMap = core.Map(keys, func(filePath tspath.Path) *BuildInfoReferenceMapEntry {
-		references, _ := t.snapshot.referencedMap.GetValues(filePath)
+		references, _ := t.snapshot.referencedMap.getReferences(filePath)
 		return &BuildInfoReferenceMapEntry{
 			FileId:       t.toFileId(filePath),
 			FileIdListId: t.toFileIdListId(references),