Watch type reference locations too and handle auto type refs

Author: sheetalkamat

internal/compiler/fileloader.go

@@ -113,10 +113,16 @@ func processAllProgramFiles(
 	var unsupportedExtensions []string
 
 	loader.parseTasks.collect(&loader, loader.rootTasks, func(task *parseTask, _ []tspath.Path) {
-		file := task.file
 		if task.isRedirected {
 			return
 		}
+
+		if task.isForAutomaticTypeDirective {
+			typeResolutionsInFile[task.path] = task.typeResolutionsInFile
+			return
+		}
+		file := task.file
+		path := task.path
 		if file == nil {
 			missingFiles = append(missingFiles, task.normalizedFilePath)
 			return
@@ -126,7 +132,6 @@ func processAllProgramFiles(
 		} else {
 			files = append(files, file)
 		}
-		path := file.Path()
 
 		filesByPath[path] = file
 		resolvedModules[path] = task.resolutionsInFile
@@ -189,14 +194,31 @@ func (p *fileLoader) addAutomaticTypeDirectiveTasks() {
 		containingDirectory = p.opts.Host.GetCurrentDirectory()
 	}
 	containingFileName := tspath.CombinePaths(containingDirectory, module.InferredTypesContainingFile)
+	p.rootTasks = append(p.rootTasks, &parseTask{normalizedFilePath: containingFileName, isLib: false, isForAutomaticTypeDirective: true})
+}
 
-	automaticTypeDirectiveNames := module.GetAutomaticTypeDirectiveNames(compilerOptions, p.opts.Host)
-	for _, name := range automaticTypeDirectiveNames {
-		resolved := p.resolver.ResolveTypeReferenceDirective(name, containingFileName, core.ModuleKindNodeNext, nil)
-		if resolved.IsResolved() {
-			p.rootTasks = append(p.rootTasks, &parseTask{normalizedFilePath: resolved.ResolvedFileName, isLib: false})
+func (p *fileLoader) resolveAutomaticTypeDirectives(containingFileName string) (
+	toParse []resolvedRef,
+	typeResolutionsInFile module.ModeAwareCache[*module.ResolvedTypeReferenceDirective],
+) {
+	automaticTypeDirectiveNames := module.GetAutomaticTypeDirectiveNames(p.opts.Config.CompilerOptions(), p.opts.Host)
+	if len(automaticTypeDirectiveNames) != 0 {
+		toParse = make([]resolvedRef, 0, len(automaticTypeDirectiveNames))
+		typeResolutionsInFile = make(module.ModeAwareCache[*module.ResolvedTypeReferenceDirective], len(automaticTypeDirectiveNames))
+		for _, name := range automaticTypeDirectiveNames {
+			resolutionMode := core.ModuleKindNodeNext
+			resolved := p.resolver.ResolveTypeReferenceDirective(name, containingFileName, resolutionMode, nil)
+			typeResolutionsInFile[module.ModeAwareCacheKey{Name: name, Mode: resolutionMode}] = resolved
+			if resolved.IsResolved() {
+				toParse = append(toParse, resolvedRef{
+					fileName:      resolved.ResolvedFileName,
+					increaseDepth: resolved.IsExternalLibraryImport,
+					elideOnDepth:  false,
+				})
+			}
 		}
 	}
+	return toParse, typeResolutionsInFile
 }
 
 func (p *fileLoader) addProjectReferenceTasks() {

internal/compiler/parsetask.go

@@ -9,13 +9,14 @@ import (
 )
 
 type parseTask struct {
-	normalizedFilePath string
-	path               tspath.Path
-	file               *ast.SourceFile
-	isLib              bool
-	isRedirected       bool
-	subTasks           []*parseTask
-	loaded             bool
+	normalizedFilePath          string
+	path                        tspath.Path
+	file                        *ast.SourceFile
+	isLib                       bool
+	isRedirected                bool
+	subTasks                    []*parseTask
+	loaded                      bool
+	isForAutomaticTypeDirective bool
 
 	metadata                     ast.SourceFileMetaData
 	resolutionsInFile            module.ModeAwareCache[*module.ResolvedModule]
@@ -36,8 +37,11 @@ func (t *parseTask) Path() tspath.Path {
 
 func (t *parseTask) load(loader *fileLoader) {
 	t.loaded = true
-
 	t.path = loader.toPath(t.normalizedFilePath)
+	if t.isForAutomaticTypeDirective {
+		t.loadAutomaticTypeDirectives(loader)
+		return
+	}
 	redirect := loader.projectReferenceFileMapper.getParseFileRedirect(t)
 	if redirect != "" {
 		t.redirect(loader, redirect)
@@ -97,6 +101,14 @@ func (t *parseTask) redirect(loader *fileLoader, fileName string) {
 	t.subTasks = []*parseTask{{normalizedFilePath: tspath.NormalizePath(fileName), isLib: t.isLib}}
 }
 
+func (t *parseTask) loadAutomaticTypeDirectives(loader *fileLoader) {
+	toParseTypeRefs, typeResolutionsInFile := loader.resolveAutomaticTypeDirectives(t.normalizedFilePath)
+	t.typeResolutionsInFile = typeResolutionsInFile
+	for _, typeResolution := range toParseTypeRefs {
+		t.addSubTask(typeResolution, false)
+	}
+}
+
 type resolvedRef struct {
 	fileName      string
 	increaseDepth bool

internal/compiler/program.go

@@ -813,6 +813,10 @@ func (p *Program) GetResolvedTypeReferenceDirectiveFromTypeReferenceDirective(ty
 	return nil
 }
 
+func (p *Program) GetResolvedTypeReferenceDirectives() map[tspath.Path]module.ModeAwareCache[*module.ResolvedTypeReferenceDirective] {
+	return p.typeResolutionsInFile
+}
+
 func (p *Program) getModeForTypeReferenceDirectiveInFile(ref *ast.FileReference, sourceFile *ast.SourceFile) core.ResolutionMode {
 	if ref.ResolutionMode != core.ResolutionModeNone {
 		return ref.ResolutionMode

internal/module/types.go

@@ -81,6 +81,10 @@ func (r *ResolvedModule) IsResolved() bool {
 	return r != nil && r.ResolvedFileName != ""
 }
 
+func (r *ResolvedModule) GetLookupLocations() *LookupLocations {
+	return &r.LookupLocations
+}
+
 type ResolvedTypeReferenceDirective struct {
 	LookupLocations
 	Primary                 bool
@@ -94,6 +98,10 @@ func (r *ResolvedTypeReferenceDirective) IsResolved() bool {
 	return r.ResolvedFileName != ""
 }
 
+func (r *ResolvedTypeReferenceDirective) GetLookupLocations() *LookupLocations {
+	return &r.LookupLocations
+}
+
 type extensions int32
 
 const (

internal/project/project.go

@@ -350,23 +350,37 @@ func (p *Project) GetLanguageServiceForRequest(ctx context.Context) (*ls.Languag
 func (p *Project) getModuleResolutionWatchGlobs() (failedLookups map[tspath.Path]string, affectingLocaions map[tspath.Path]string) {
 	failedLookups = make(map[tspath.Path]string)
 	affectingLocaions = make(map[tspath.Path]string)
-	for _, resolvedModulesInFile := range p.program.GetResolvedModules() {
+	extractLookups(p, failedLookups, affectingLocaions, p.program.GetResolvedModules())
+	extractLookups(p, failedLookups, affectingLocaions, p.program.GetResolvedTypeReferenceDirectives())
+	return failedLookups, affectingLocaions
+}
+
+type ResolutionWithLookupLocations interface {
+	GetLookupLocations() *module.LookupLocations
+}
+
+func extractLookups[T ResolutionWithLookupLocations](
+	p *Project,
+	failedLookups map[tspath.Path]string,
+	affectingLocaions map[tspath.Path]string,
+	cache map[tspath.Path]module.ModeAwareCache[T],
+) {
+	for _, resolvedModulesInFile := range cache {
 		for _, resolvedModule := range resolvedModulesInFile {
-			for _, failedLookupLocation := range resolvedModule.FailedLookupLocations {
+			for _, failedLookupLocation := range resolvedModule.GetLookupLocations().FailedLookupLocations {
 				path := p.toPath(failedLookupLocation)
 				if _, ok := failedLookups[path]; !ok {
 					failedLookups[path] = failedLookupLocation
 				}
 			}
-			for _, affectingLocation := range resolvedModule.AffectingLocations {
+			for _, affectingLocation := range resolvedModule.GetLookupLocations().AffectingLocations {
 				path := p.toPath(affectingLocation)
 				if _, ok := affectingLocaions[path]; !ok {
 					affectingLocaions[path] = affectingLocation
 				}
 			}
 		}
 	}
-	return failedLookups, affectingLocaions
 }
 
 func (p *Project) updateWatchers(ctx context.Context) {