Actually find and load the original location projects

Author: sheetalkamat

internal/project/dirty/set.go

@@ -0,0 +1,126 @@
+package dirty
+
+import "maps"
+
+type SetEntry[K comparable] struct {
+	m *Set[K]
+	mapEntry[K, struct{}]
+}
+
+func (e *SetEntry[K]) Delete() {
+	if !e.dirty {
+		e.m.dirty[e.key] = e
+	}
+	e.delete = true
+}
+
+type Set[K comparable] struct {
+	base  map[K]struct{}
+	dirty map[K]*SetEntry[K]
+}
+
+func NewSet[K comparable](base map[K]struct{}) *Set[K] {
+	return &Set[K]{
+		base:  base,
+		dirty: make(map[K]*SetEntry[K]),
+	}
+}
+
+func (m *Set[K]) Get(key K) (*SetEntry[K], bool) {
+	if entry, ok := m.dirty[key]; ok {
+		if entry.delete {
+			return nil, false
+		}
+		return entry, true
+	}
+	value, ok := m.base[key]
+	if !ok {
+		return nil, false
+	}
+	return &SetEntry[K]{
+		m: m,
+		mapEntry: mapEntry[K, struct{}]{
+			key:      key,
+			original: value,
+			value:    value,
+			dirty:    false,
+		},
+	}, true
+}
+
+func (m *Set[K]) Has(key K) bool {
+	if entry, ok := m.dirty[key]; ok {
+		if entry.delete {
+			return false
+		}
+		return true
+	}
+	_, ok := m.base[key]
+	return ok
+}
+
+// Add sets a new entry in the dirty map without checking if it exists
+// in the base map. The entry added is considered dirty, so it should
+// be a fresh value, mutable until finalized (i.e., it will not be cloned
+// before changing if a change is made). If modifying an entry that may
+// exist in the base map, use `Change` instead.
+func (m *Set[K]) Add(key K) {
+	m.dirty[key] = &SetEntry[K]{
+		m: m,
+		mapEntry: mapEntry[K, struct{}]{
+			key:   key,
+			value: struct{}{},
+			dirty: true,
+		},
+	}
+}
+
+func (m *Set[K]) Delete(key K) {
+	if entry, ok := m.Get(key); ok {
+		entry.Delete()
+	} else {
+		panic("tried to delete a non-existent entry")
+	}
+}
+
+func (m *Set[K]) Range(fn func(*SetEntry[K]) bool) {
+	seenInDirty := make(map[K]struct{})
+	for _, entry := range m.dirty {
+		seenInDirty[entry.key] = struct{}{}
+		if !entry.delete && !fn(entry) {
+			break
+		}
+	}
+	for key, value := range m.base {
+		if _, ok := seenInDirty[key]; ok {
+			continue // already processed in dirty entries
+		}
+		if !fn(&SetEntry[K]{m: m, mapEntry: mapEntry[K, struct{}]{
+			key:      key,
+			original: value,
+			value:    value,
+			dirty:    false,
+		}}) {
+			break
+		}
+	}
+}
+
+func (m *Set[K]) Finalize() (result map[K]struct{}, changed bool) {
+	if len(m.dirty) == 0 {
+		return m.base, false // no changes, return base map
+	}
+	if m.base == nil {
+		result = make(map[K]struct{}, len(m.dirty))
+	} else {
+		result = maps.Clone(m.base)
+	}
+	for key, entry := range m.dirty {
+		if entry.delete {
+			delete(result, key)
+		} else {
+			result[key] = entry.value
+		}
+	}
+	return result, true
+}

internal/project/projectcollection.go

@@ -27,6 +27,8 @@ type ProjectCollection struct {
 	// apiOpenedProjects is the set of projects that should be kept open for
 	// API clients.
 	apiOpenedProjects map[tspath.Path]struct{}
+	// Files that are open
+	openedFiles map[tspath.Path]struct{}
 }
 
 func (c *ProjectCollection) ConfigFileRegistry() *ConfigFileRegistry { return c.configFileRegistry }
@@ -229,6 +231,7 @@ func (c *ProjectCollection) clone() *ProjectCollection {
 		configuredProjects:  c.configuredProjects,
 		inferredProject:     c.inferredProject,
 		fileDefaultProjects: c.fileDefaultProjects,
+		openedFiles:         c.openedFiles,
 	}
 }
 

internal/project/projectcollectionbuilder.go

@@ -38,9 +38,11 @@ type ProjectCollectionBuilder struct {
 
 	newSnapshotID           uint64
 	programStructureChanged bool
-	fileDefaultProjects     map[tspath.Path]tspath.Path
-	configuredProjects      *dirty.SyncMap[tspath.Path, *Project]
-	inferredProject         *dirty.Box[*Project]
+
+	fileDefaultProjects map[tspath.Path]tspath.Path
+	configuredProjects  *dirty.SyncMap[tspath.Path, *Project]
+	inferredProject     *dirty.Box[*Project]
+	openedFiles         *dirty.Set[tspath.Path]
 
 	apiOpenedProjects map[tspath.Path]struct{}
 }
@@ -69,6 +71,7 @@ func newProjectCollectionBuilder(
 		newSnapshotID:                      newSnapshotID,
 		configuredProjects:                 dirty.NewSyncMap(oldProjectCollection.configuredProjects, nil),
 		inferredProject:                    dirty.NewBox(oldProjectCollection.inferredProject),
+		openedFiles:                        dirty.NewSet(oldProjectCollection.openedFiles),
 		apiOpenedProjects:                  maps.Clone(oldAPIOpenedProjects),
 	}
 }
@@ -98,6 +101,11 @@ func (b *ProjectCollectionBuilder) Finalize(logger *logging.LogTree) (*ProjectCo
 		newProjectCollection.inferredProject = newInferredProject
 	}
 
+	if openedFiles, openedFilesChanged := b.openedFiles.Finalize(); openedFilesChanged {
+		ensureCloned()
+		newProjectCollection.openedFiles = openedFiles
+	}
+
 	configFileRegistry := b.configFileRegistryBuilder.Finalize()
 	newProjectCollection.configFileRegistry = configFileRegistry
 	return newProjectCollection, configFileRegistry
@@ -174,6 +182,7 @@ func (b *ProjectCollectionBuilder) DidChangeFiles(summary FileChangeSummary, log
 		if fh := b.fs.GetFileByPath(fileName, path); fh == nil || fh.Hash() != hash {
 			changedFiles = append(changedFiles, path)
 		}
+		b.openedFiles.Delete(path)
 	}
 	for uri := range summary.Changed.Keys() {
 		fileName := uri.FileName()
@@ -242,8 +251,9 @@ func (b *ProjectCollectionBuilder) DidChangeFiles(summary FileChangeSummary, log
 	if summary.Opened != "" {
 		fileName := summary.Opened.FileName()
 		path := b.toPath(fileName)
+		b.openedFiles.Add(path)
 		var toRemoveProjects collections.Set[tspath.Path]
-		openFileResult := b.ensureConfiguredProjectAndAncestorsForOpenFile(fileName, path, logger)
+		openFileResult := b.ensureConfiguredProjectAndAncestorsForFile(fileName, path, logger)
 		b.configuredProjects.Range(func(entry *dirty.SyncMapEntry[tspath.Path, *Project]) bool {
 			toRemoveProjects.Add(entry.Value().configFilePath)
 			b.updateProgram(entry, logger)
@@ -357,6 +367,23 @@ func (b *ProjectCollectionBuilder) DidRequestProject(projectId tspath.Path, logg
 	}
 }
 
+func (b *ProjectCollectionBuilder) DidRequestEnsureDefaultProject(uri lsproto.DocumentUri, logger *logging.LogTree) {
+	fileName := uri.FileName()
+	path := b.toPath(fileName)
+	if b.openedFiles.Has(path) {
+		b.DidRequestFile(uri, logger)
+		return
+	}
+
+	startTime := time.Now()
+	b.ensureConfiguredProjectAndAncestorsForFile(fileName, path, logger)
+
+	if logger != nil {
+		elapsed := time.Since(startTime)
+		logger.Log(fmt.Sprintf("Completed file request for %s in %v", fileName, elapsed))
+	}
+}
+
 func (b *ProjectCollectionBuilder) DidUpdateATAState(ataChanges map[tspath.Path]*ATAStateChange, logger *logging.LogTree) {
 	updateProject := func(project dirty.Value[*Project], ataChange *ATAStateChange) {
 		project.ChangeIf(
@@ -427,7 +454,7 @@ func (b *ProjectCollectionBuilder) markProjectsAffectedByConfigChanges(
 	var hasChanges bool
 	for path := range configChangeResult.affectedFiles {
 		fileName := b.fs.overlays[path].FileName()
-		_ = b.ensureConfiguredProjectAndAncestorsForOpenFile(fileName, path, logger)
+		_ = b.ensureConfiguredProjectAndAncestorsForFile(fileName, path, logger)
 		hasChanges = true
 	}
 
@@ -468,17 +495,17 @@ func (b *ProjectCollectionBuilder) findDefaultConfiguredProject(fileName string,
 	})
 
 	if multipleCandidates {
-		if p := b.findOrCreateDefaultConfiguredProjectForOpenScriptInfo(fileName, path, projectLoadKindFind, nil).project; p != nil {
+		if p := b.findOrCreateDefaultConfiguredProjectForFile(fileName, path, projectLoadKindFind, nil).project; p != nil {
 			return p
 		}
 	}
 
 	return configuredProjects[project]
 }
 
-func (b *ProjectCollectionBuilder) ensureConfiguredProjectAndAncestorsForOpenFile(fileName string, path tspath.Path, logger *logging.LogTree) searchResult {
-	result := b.findOrCreateDefaultConfiguredProjectForOpenScriptInfo(fileName, path, projectLoadKindCreate, logger)
-	if result.project != nil {
+func (b *ProjectCollectionBuilder) ensureConfiguredProjectAndAncestorsForFile(fileName string, path tspath.Path, logger *logging.LogTree) searchResult {
+	result := b.findOrCreateDefaultConfiguredProjectForFile(fileName, path, projectLoadKindCreate, logger)
+	if result.project != nil && b.openedFiles.Has(path) {
 		// !!! sheetal todo this later
 		// // Create ancestor tree for findAllRefs (dont load them right away)
 		// forEachAncestorProjectLoad(
@@ -672,7 +699,7 @@ func (b *ProjectCollectionBuilder) findOrCreateDefaultConfiguredProjectWorker(
 	return searchResult{retain: retain}
 }
 
-func (b *ProjectCollectionBuilder) findOrCreateDefaultConfiguredProjectForOpenScriptInfo(
+func (b *ProjectCollectionBuilder) findOrCreateDefaultConfiguredProjectForFile(
 	fileName string,
 	path tspath.Path,
 	loadKind projectLoadKind,