Handle existing projects

sheetalkamat · sheetalkamat · commit f88b8f278f70 · 2025-11-03T13:45:52.000-08:00
diff --git a/internal/lsp/server.go b/internal/lsp/server.go
@@ -505,8 +505,8 @@ var handlers = sync.OnceValue(func() handlerMap {
 	registerLanguageServiceDocumentRequestHandler(handlers, lsproto.TextDocumentDocumentHighlightInfo, (*Server).handleDocumentHighlight)
 	registerLanguageServiceDocumentRequestHandler(handlers, lsproto.TextDocumentSelectionRangeInfo, (*Server).handleSelectionRange)
 
-	registerMultiProjectDocumentRequestHandler(handlers, lsproto.TextDocumentReferencesInfo, (*Server).handleReferences)
-	registerMultiProjectDocumentRequestHandler(handlers, lsproto.TextDocumentRenameInfo, (*Server).handleRename)
+	registerMultiProjectDocumentRequestHandler(handlers, lsproto.TextDocumentReferencesInfo, (*Server).handleReferences, combineReferences)
+	registerMultiProjectDocumentRequestHandler(handlers, lsproto.TextDocumentRenameInfo, (*Server).handleRename, combineRenameResponse)
 
 	registerRequestHandler(handlers, lsproto.WorkspaceSymbolInfo, (*Server).handleWorkspaceSymbol)
 	registerRequestHandler(handlers, lsproto.CompletionItemResolveInfo, (*Server).handleCompletionItemResolve)
@@ -583,27 +583,70 @@ func registerLanguageServiceDocumentRequestHandler[Req lsproto.HasTextDocumentUR
 	}
 }
 
-func registerMultiProjectDocumentRequestHandler[Req lsproto.HasTextDocumentURI, Resp any](handlers handlerMap, info lsproto.RequestInfo[Req, Resp], fn func(*Server, context.Context, *ls.LanguageService, Req) (Resp, error)) {
+func registerMultiProjectDocumentRequestHandler[Req lsproto.HasTextDocumentURI, Resp any](
+	handlers handlerMap,
+	info lsproto.RequestInfo[Req, Resp],
+	fn func(*Server, context.Context, *ls.LanguageService, Req) (Resp, error),
+	combineResults func(*project.Project, *collections.SyncMap[tspath.Path, Resp]) Resp,
+) {
 	handlers[info.Method] = func(s *Server, ctx context.Context, req *lsproto.RequestMessage) error {
 		var params Req
 		// Ignore empty params.
 		if req.Params != nil {
 			params = req.Params.(Req)
 		}
 		// !!! sheetal: multiple projects that contain the file through symlinks
-		// !!! multiple projects that contain the file directly
-		ls, err := s.session.GetLanguageService(ctx, params.TextDocumentURI())
+		defaultProject, defaultLs, allProjects, err := s.session.GetProjectsForFile(ctx, params.TextDocumentURI())
 		if err != nil {
 			return err
 		}
 		defer s.recover(req)
-		resp, err := fn(s, ctx, ls, params)
-		if err != nil {
-			return err
+		var results collections.SyncMap[tspath.Path, Resp]
+		var workInProgress collections.SyncMap[tspath.Path, bool]
+		workInProgress.Store(defaultProject.Id(), true)
+		wg := core.NewWorkGroup(false)
+		enqueueWorkForNonDefaultProject := func(project *project.Project) {
+			if project == defaultProject {
+				return
+			}
+			wg.Queue(func() {
+				if ctx.Err() != nil {
+					return
+				}
+				if _, loaded := workInProgress.LoadOrStore(project.Id(), true); loaded {
+					return
+				}
+				defer s.recover(req)
+				ls := s.session.GetLanguageServiceForProjectWithFile(ctx, project, params.TextDocumentURI())
+				if ls != nil {
+					resp, err1 := fn(s, ctx, ls, params)
+					if err1 != nil {
+						return
+					}
+					results.Store(project.Id(), resp)
+				}
+			})
 		}
+
+		// Other projects
+		for _, project := range allProjects {
+			enqueueWorkForNonDefaultProject(project)
+		}
+
+		// Default projects
+		wg.Queue(func() {
+			resp, err1 := fn(s, ctx, defaultLs, params)
+			if err1 != nil {
+				return
+			}
+			results.Store(defaultProject.Id(), resp)
+			// TODO:: if default location needs to be add more projects to request
+		})
+		wg.RunAndWait()
 		if ctx.Err() != nil {
 			return ctx.Err()
 		}
+		resp := combineResults(defaultProject, &results)
 		s.sendResult(req.ID, resp)
 		return nil
 	}
@@ -880,9 +923,36 @@ func (s *Server) handleTypeDefinition(ctx context.Context, ls *ls.LanguageServic
 	return ls.ProvideTypeDefinition(ctx, params.TextDocument.Uri, params.Position, getTypeDefinitionClientSupportsLink(s.initializeParams))
 }
 
-func (s *Server) handleReferences(ctx context.Context, ls *ls.LanguageService, params *lsproto.ReferenceParams) (lsproto.ReferencesResponse, error) {
+func (s *Server) handleReferences(ctx context.Context, defaultLs *ls.LanguageService, params *lsproto.ReferenceParams) (lsproto.ReferencesResponse, error) {
 	// findAllReferences
-	return ls.ProvideReferences(ctx, params)
+	return defaultLs.ProvideReferences(ctx, params)
+}
+
+func combineReferences(defaultProject *project.Project, results *collections.SyncMap[tspath.Path, lsproto.ReferencesResponse]) lsproto.ReferencesResponse {
+	var combined []lsproto.Location
+	var seenLocations collections.Set[lsproto.Location]
+	if resp, ok := results.Load(defaultProject.Id()); ok {
+		if resp.Locations != nil {
+			for _, loc := range *resp.Locations {
+				seenLocations.Add(loc)
+				combined = append(combined, loc)
+			}
+		}
+	}
+	results.Range(func(projectID tspath.Path, resp lsproto.ReferencesResponse) bool {
+		if projectID != defaultProject.Id() {
+			if resp.Locations != nil {
+				for _, loc := range *resp.Locations {
+					if !seenLocations.Has(loc) {
+						seenLocations.Add(loc)
+						combined = append(combined, loc)
+					}
+				}
+			}
+		}
+		return true
+	})
+	return lsproto.LocationsOrNull{Locations: &combined}
 }
 
 func (s *Server) handleImplementations(ctx context.Context, ls *ls.LanguageService, params *lsproto.ImplementationParams) (lsproto.ImplementationResponse, error) {
@@ -962,6 +1032,62 @@ func (s *Server) handleRename(ctx context.Context, ls *ls.LanguageService, param
 	return ls.ProvideRename(ctx, params)
 }
 
+func combineRenameResponse(defaultProject *project.Project, results *collections.SyncMap[tspath.Path, lsproto.RenameResponse]) lsproto.RenameResponse {
+	combined := make(map[lsproto.DocumentUri][]*lsproto.TextEdit)
+	seenChanges := make(map[lsproto.DocumentUri]*collections.Set[lsproto.TextEdit])
+	// !!! this is not used any more so we will skip this part of deduplication and combining
+	// 	DocumentChanges *[]TextDocumentEditOrCreateFileOrRenameFileOrDeleteFile `json:"documentChanges,omitzero"`
+	// 	ChangeAnnotations *map[string]*ChangeAnnotation `json:"changeAnnotations,omitzero"`
+
+	if resp, ok := results.Load(defaultProject.Id()); ok {
+		if resp.WorkspaceEdit != nil {
+			for doc, changes := range *resp.WorkspaceEdit.Changes {
+				seenSet := collections.Set[lsproto.TextEdit]{}
+				seenChanges[doc] = &seenSet
+				var changesForDoc []*lsproto.TextEdit
+				for _, change := range changes {
+					seenSet.Add(*change)
+					changesForDoc = append(changesForDoc, change)
+				}
+				combined[doc] = changesForDoc
+			}
+		}
+	}
+	results.Range(func(projectID tspath.Path, resp lsproto.RenameResponse) bool {
+		if projectID != defaultProject.Id() {
+			if resp.WorkspaceEdit != nil {
+				for doc, changes := range *resp.WorkspaceEdit.Changes {
+					seenSet, ok := seenChanges[doc]
+					if !ok {
+						seenSet = &collections.Set[lsproto.TextEdit]{}
+						seenChanges[doc] = seenSet
+					}
+					changesForDoc, exists := combined[doc]
+					if !exists {
+						changesForDoc = []*lsproto.TextEdit{}
+					}
+					for _, change := range changes {
+						if !seenSet.Has(*change) {
+							seenSet.Add(*change)
+							changesForDoc = append(changesForDoc, change)
+						}
+					}
+					combined[doc] = changesForDoc
+				}
+			}
+		}
+		return true
+	})
+	if len(combined) > 0 {
+		return lsproto.RenameResponse{
+			WorkspaceEdit: &lsproto.WorkspaceEdit{
+				Changes: &combined,
+			},
+		}
+	}
+	return lsproto.RenameResponse{}
+}
+
 func (s *Server) handleDocumentHighlight(ctx context.Context, ls *ls.LanguageService, params *lsproto.DocumentHighlightParams) (lsproto.DocumentHighlightResponse, error) {
 	return ls.ProvideDocumentHighlights(ctx, params.TextDocument.Uri, params.Position)
 }
diff --git a/internal/project/project.go b/internal/project/project.go
@@ -192,6 +192,10 @@ func (p *Project) ConfigFilePath() tspath.Path {
 	return p.configFilePath
 }
 
+func (p *Project) Id() tspath.Path {
+	return p.configFilePath
+}
+
 func (p *Project) GetProgram() *compiler.Program {
 	return p.Program
 }
diff --git a/internal/project/projectcollection.go b/internal/project/projectcollection.go
@@ -93,6 +93,19 @@ func (c *ProjectCollection) InferredProject() *Project {
 	return c.inferredProject
 }
 
+func (c *ProjectCollection) GetProjectsContainingFile(path tspath.Path) []*Project {
+	var projects []*Project
+	for _, project := range c.ConfiguredProjects() {
+		if project.containsFile(path) {
+			projects = append(projects, project)
+		}
+	}
+	if c.inferredProject != nil && c.inferredProject.containsFile(path) {
+		projects = append(projects, c.inferredProject)
+	}
+	return projects
+}
+
 // !!! result could be cached
 func (c *ProjectCollection) GetDefaultProject(fileName string, path tspath.Path) *Project {
 	if result, ok := c.fileDefaultProjects[path]; ok {
diff --git a/internal/project/projectcollectionbuilder.go b/internal/project/projectcollectionbuilder.go
@@ -338,6 +338,25 @@ func (b *ProjectCollectionBuilder) DidRequestFile(uri lsproto.DocumentUri, logge
 	}
 }
 
+func (b *ProjectCollectionBuilder) DidRequestProject(projectId tspath.Path, logger *logging.LogTree) {
+	startTime := time.Now()
+	if projectId == inferredProjectName {
+		// Update inferred project
+		if b.inferredProject.Value() != nil {
+			b.updateProgram(b.inferredProject, logger)
+		}
+	} else {
+		if entry, ok := b.configuredProjects.Load(projectId); ok {
+			b.updateProgram(entry, logger)
+		}
+	}
+
+	if logger != nil {
+		elapsed := time.Since(startTime)
+		logger.Log(fmt.Sprintf("Completed project update request for %s in %v", projectId, elapsed))
+	}
+}
+
 func (b *ProjectCollectionBuilder) DidUpdateATAState(ataChanges map[tspath.Path]*ATAStateChange, logger *logging.LogTree) {
 	updateProject := func(project dirty.Value[*Project], ataChange *ATAStateChange) {
 		project.ChangeIf(
diff --git a/internal/project/session.go b/internal/project/session.go
@@ -371,29 +371,34 @@ func (s *Session) Snapshot() (*Snapshot, func()) {
 	}
 }
 
-func (s *Session) GetLanguageService(ctx context.Context, uri lsproto.DocumentUri) (*ls.LanguageService, error) {
+func (s *Session) getSnapshot(ctx context.Context, uri lsproto.DocumentUri, projectId tspath.Path) (*Snapshot, map[tspath.Path]*Overlay, bool) {
 	var snapshot *Snapshot
 	fileChanges, overlays, ataChanges, newConfig := s.flushChanges(ctx)
 	updateSnapshot := !fileChanges.IsEmpty() || len(ataChanges) > 0 || newConfig != nil
 	if updateSnapshot {
 		// If there are pending file changes, we need to update the snapshot.
 		// Sending the requested URI ensures that the project for this URI is loaded.
 		snapshot = s.UpdateSnapshot(ctx, overlays, SnapshotChange{
-			reason:        UpdateReasonRequestedLanguageServicePendingChanges,
-			fileChanges:   fileChanges,
-			ataChanges:    ataChanges,
-			newConfig:     newConfig,
-			requestedURIs: []lsproto.DocumentUri{uri},
+			reason:                  UpdateReasonRequestedLanguageServicePendingChanges,
+			fileChanges:             fileChanges,
+			ataChanges:              ataChanges,
+			newConfig:               newConfig,
+			requestedURIs:           core.IfElse(uri != "", []lsproto.DocumentUri{uri}, nil),
+			requestedProjectUpdates: core.IfElse(projectId != "", []tspath.Path{projectId}, nil),
 		})
 	} else {
 		// If there are no pending file changes, we can try to use the current snapshot.
 		s.snapshotMu.RLock()
 		snapshot = s.snapshot
 		s.snapshotMu.RUnlock()
 	}
+	return snapshot, overlays, updateSnapshot
+}
 
+func (s *Session) getSnapshotAndDefaultProject(ctx context.Context, uri lsproto.DocumentUri) (*Snapshot, *Project, *ls.LanguageService, error) {
+	snapshot, overlays, updatedSnapshot := s.getSnapshot(ctx, uri, "")
 	project := snapshot.GetDefaultProject(uri)
-	if project == nil && !updateSnapshot || project != nil && project.dirty {
+	if project == nil && !updatedSnapshot || project != nil && project.dirty {
 		// The current snapshot does not have an up to date project for the URI,
 		// so we need to update the snapshot to ensure the project is loaded.
 		// !!! Allow multiple projects to update in parallel
@@ -404,9 +409,51 @@ func (s *Session) GetLanguageService(ctx context.Context, uri lsproto.DocumentUr
 		project = snapshot.GetDefaultProject(uri)
 	}
 	if project == nil {
-		return nil, fmt.Errorf("no project found for URI %s", uri)
+		return nil, nil, nil, fmt.Errorf("no project found for URI %s", uri)
+	}
+	return snapshot, project, ls.NewLanguageService(project.GetProgram(), snapshot), nil
+}
+
+func (s *Session) GetLanguageService(ctx context.Context, uri lsproto.DocumentUri) (*ls.LanguageService, error) {
+	_, _, languageService, err := s.getSnapshotAndDefaultProject(ctx, uri)
+	if err != nil {
+		return nil, err
+	}
+	return languageService, nil
+}
+
+func (s *Session) GetProjectsForFile(ctx context.Context, uri lsproto.DocumentUri) (*Project, *ls.LanguageService, []*Project, error) {
+	snapshot, project, defaultLs, err := s.getSnapshotAndDefaultProject(ctx, uri)
+	if err != nil {
+		return nil, nil, nil, err
+	}
+	// !!! TODO: sheetal:  Get other projects that contain the file with symlink
+	allProjects := snapshot.GetProjectsContainingFile(uri)
+	return project, defaultLs, allProjects, nil
+}
+
+func (s *Session) GetLanguageServiceForProjectWithFile(ctx context.Context, project *Project, uri lsproto.DocumentUri) *ls.LanguageService {
+	snapshot, overlays, updatedSnapshot := s.getSnapshot(ctx, "", project.Id())
+	// update the snapshot
+	if !updatedSnapshot && project.dirty {
+		// The current snapshot does not have an up to date project for the URI,
+		// so we need to update the snapshot to ensure the project is loaded.
+		// !!! Allow multiple projects to update in parallel
+		snapshot = s.UpdateSnapshot(ctx, overlays, SnapshotChange{
+			reason:                  UpdateReasonRequestedLanguageServiceProjectDirty,
+			requestedProjectUpdates: []tspath.Path{project.Id()},
+		})
+	}
+	// Ensure we have updated project
+	project = snapshot.ProjectCollection.GetProjectByPath(project.Id())
+	if project == nil {
+		return nil
+	}
+	// if program doesnt contain this file any more ignore it
+	if path := s.toPath(uri.FileName()); !project.containsFile(path) {
+		return nil
 	}
-	return ls.NewLanguageService(project.GetProgram(), snapshot), nil
+	return ls.NewLanguageService(project.GetProgram(), snapshot)
 }
 
 func (s *Session) UpdateSnapshot(ctx context.Context, overlays map[tspath.Path]*Overlay, change SnapshotChange) *Snapshot {
diff --git a/internal/project/snapshot.go b/internal/project/snapshot.go
@@ -76,6 +76,13 @@ func (s *Snapshot) GetDefaultProject(uri lsproto.DocumentUri) *Project {
 	return s.ProjectCollection.GetDefaultProject(fileName, path)
 }
 
+func (s *Snapshot) GetProjectsContainingFile(uri lsproto.DocumentUri) []*Project {
+	fileName := uri.FileName()
+	path := s.toPath(fileName)
+	// TODO!! sheetal may be change this to handle symlinks!!
+	return s.ProjectCollection.GetProjectsContainingFile(path)
+}
+
 func (s *Snapshot) GetFile(fileName string) FileHandle {
 	return s.fs.GetFile(fileName)
 }
@@ -135,6 +142,9 @@ type SnapshotChange struct {
 	// requestedURIs are URIs that were requested by the client.
 	// The new snapshot should ensure projects for these URIs have loaded programs.
 	requestedURIs []lsproto.DocumentUri
+	// Update requested projects.
+	// this is used when we want to get LS and from all the projects the file can be part of
+	requestedProjectUpdates []tspath.Path
 	// compilerOptionsForInferredProjects is the compiler options to use for inferred projects.
 	// It should only be set the value in the next snapshot should be changed. If nil, the
 	// value from the previous snapshot will be copied to the new snapshot.
@@ -185,11 +195,11 @@ func (s *Snapshot) Clone(ctx context.Context, change SnapshotChange, overlays ma
 		case UpdateReasonDidChangeCompilerOptionsForInferredProjects:
 			logger.Logf("Reason: DidChangeCompilerOptionsForInferredProjects")
 		case UpdateReasonRequestedLanguageServicePendingChanges:
-			logger.Logf("Reason: RequestedLanguageService (pending file changes) - %v", change.requestedURIs)
+			logger.Logf("Reason: RequestedLanguageService (pending file changes) - %v, %v", change.requestedURIs, change.requestedProjectUpdates)
 		case UpdateReasonRequestedLanguageServiceProjectNotLoaded:
 			logger.Logf("Reason: RequestedLanguageService (project not loaded) - %v", change.requestedURIs)
 		case UpdateReasonRequestedLanguageServiceProjectDirty:
-			logger.Logf("Reason: RequestedLanguageService (project dirty) - %v", change.requestedURIs)
+			logger.Logf("Reason: RequestedLanguageService (project dirty) - %v, %v", change.requestedURIs, change.requestedProjectUpdates)
 		}
 	}
 
@@ -248,6 +258,10 @@ func (s *Snapshot) Clone(ctx context.Context, change SnapshotChange, overlays ma
 		projectCollectionBuilder.DidRequestFile(uri, logger.Fork("DidRequestFile"))
 	}
 
+	for _, projectId := range change.requestedProjectUpdates {
+		projectCollectionBuilder.DidRequestProject(projectId, logger.Fork("DidRequestProject"))
+	}
+
 	projectCollection, configFileRegistry := projectCollectionBuilder.Finalize(logger)
 
 	// Clean cached disk files not touched by any open project. It's not important that we do this on
diff --git a/testdata/baselines/reference/fourslash/findAllReferences/isDefinitionAcrossModuleProjects.baseline.jsonc b/testdata/baselines/reference/fourslash/findAllReferences/isDefinitionAcrossModuleProjects.baseline.jsonc
diff --git a/testdata/baselines/reference/lspservertests/findAllRefs/does-not-try-to-open-a-file-in-a-project-that-was-updated-and-no-longer-has-the-file.js b/testdata/baselines/reference/lspservertests/findAllRefs/does-not-try-to-open-a-file-in-a-project-that-was-updated-and-no-longer-has-the-file.js
diff --git a/testdata/baselines/reference/lspservertests/findAllRefs/root-file-is-file-from-referenced-project.js b/testdata/baselines/reference/lspservertests/findAllRefs/root-file-is-file-from-referenced-project.js

Original file line number	Diff line number	Diff line change
`@@ -192,6 +192,10 @@ func (p *Project) ConfigFilePath() tspath.Path {`
`192`	`192`	`return p.configFilePath`
`193`	`193`	`}`
`194`	`194`
	`195`	`+func (p *Project) Id() tspath.Path {`
	`196`	`+ return p.configFilePath`
	`197`	`+}`
	`198`	`+`
`195`	`199`	`func (p Project) GetProgram() compiler.Program {`
`196`	`200`	`return p.Program`
`197`	`201`	`}`