Fix BFS loop during config lookup (#1648)

Author: gabritto

internal/core/bfs.go

@@ -8,43 +8,43 @@ import (
 	"github.com/microsoft/typescript-go/internal/collections"
 )
 
-type BreadthFirstSearchResult[N comparable] struct {
+type BreadthFirstSearchResult[N any] struct {
 	Stopped bool
 	Path    []N
 }
 
-type breadthFirstSearchJob[N comparable] struct {
+type breadthFirstSearchJob[N any] struct {
 	node   N
 	parent *breadthFirstSearchJob[N]
 }
 
-type BreadthFirstSearchLevel[N comparable] struct {
-	jobs *collections.OrderedMap[N, *breadthFirstSearchJob[N]]
+type BreadthFirstSearchLevel[K comparable, N any] struct {
+	jobs *collections.OrderedMap[K, *breadthFirstSearchJob[N]]
 }
 
-func (l *BreadthFirstSearchLevel[N]) Has(node N) bool {
-	return l.jobs.Has(node)
+func (l *BreadthFirstSearchLevel[K, N]) Has(key K) bool {
+	return l.jobs.Has(key)
 }
 
-func (l *BreadthFirstSearchLevel[N]) Delete(node N) {
-	l.jobs.Delete(node)
+func (l *BreadthFirstSearchLevel[K, N]) Delete(key K) {
+	l.jobs.Delete(key)
 }
 
-func (l *BreadthFirstSearchLevel[N]) Range(f func(node N) bool) {
-	for node := range l.jobs.Keys() {
-		if !f(node) {
+func (l *BreadthFirstSearchLevel[K, N]) Range(f func(node N) bool) {
+	for job := range l.jobs.Values() {
+		if !f(job.node) {
 			return
 		}
 	}
 }
 
-type BreadthFirstSearchOptions[N comparable] struct {
+type BreadthFirstSearchOptions[K comparable, N any] struct {
 	// Visited is a set of nodes that have already been visited.
 	// If nil, a new set will be created.
-	Visited *collections.SyncSet[N]
+	Visited *collections.SyncSet[K]
 	// PreprocessLevel is a function that, if provided, will be called
 	// before each level, giving the caller an opportunity to remove nodes.
-	PreprocessLevel func(*BreadthFirstSearchLevel[N])
+	PreprocessLevel func(*BreadthFirstSearchLevel[K, N])
 }
 
 // BreadthFirstSearchParallel performs a breadth-first search on a graph
@@ -55,41 +55,42 @@ func BreadthFirstSearchParallel[N comparable](
 	neighbors func(N) []N,
 	visit func(node N) (isResult bool, stop bool),
 ) BreadthFirstSearchResult[N] {
-	return BreadthFirstSearchParallelEx(start, neighbors, visit, BreadthFirstSearchOptions[N]{})
+	return BreadthFirstSearchParallelEx(start, neighbors, visit, BreadthFirstSearchOptions[N, N]{}, Identity)
 }
 
 // BreadthFirstSearchParallelEx is an extension of BreadthFirstSearchParallel that allows
 // the caller to pass a pre-seeded set of already-visited nodes and a preprocessing function
 // that can be used to remove nodes from each level before parallel processing.
-func BreadthFirstSearchParallelEx[N comparable](
+func BreadthFirstSearchParallelEx[K comparable, N any](
 	start N,
 	neighbors func(N) []N,
 	visit func(node N) (isResult bool, stop bool),
-	options BreadthFirstSearchOptions[N],
+	options BreadthFirstSearchOptions[K, N],
+	getKey func(N) K,
 ) BreadthFirstSearchResult[N] {
 	visited := options.Visited
 	if visited == nil {
-		visited = &collections.SyncSet[N]{}
+		visited = &collections.SyncSet[K]{}
 	}
 
 	type result struct {
 		stop bool
 		job  *breadthFirstSearchJob[N]
-		next *collections.OrderedMap[N, *breadthFirstSearchJob[N]]
+		next *collections.OrderedMap[K, *breadthFirstSearchJob[N]]
 	}
 
 	var fallback *breadthFirstSearchJob[N]
 	// processLevel processes each node at the current level in parallel.
 	// It produces either a list of jobs to be processed in the next level,
 	// or a result if the visit function returns true for any node.
-	processLevel := func(index int, jobs *collections.OrderedMap[N, *breadthFirstSearchJob[N]]) result {
+	processLevel := func(index int, jobs *collections.OrderedMap[K, *breadthFirstSearchJob[N]]) result {
 		var lowestFallback atomic.Int64
 		var lowestGoal atomic.Int64
 		var nextJobCount atomic.Int64
 		lowestGoal.Store(math.MaxInt64)
 		lowestFallback.Store(math.MaxInt64)
 		if options.PreprocessLevel != nil {
-			options.PreprocessLevel(&BreadthFirstSearchLevel[N]{jobs: jobs})
+			options.PreprocessLevel(&BreadthFirstSearchLevel[K, N]{jobs: jobs})
 		}
 		next := make([][]*breadthFirstSearchJob[N], jobs.Size())
 		var wg sync.WaitGroup
@@ -103,7 +104,7 @@ func BreadthFirstSearchParallelEx[N comparable](
 				}
 
 				// If we have already visited this node, skip it.
-				if !visited.AddIfAbsent(j.node) {
+				if !visited.AddIfAbsent(getKey(j.node)) {
 					// Note that if we are here, we already visited this node at a
 					// previous *level*, which means `visit` must have returned false,
 					// so we don't need to update our result indices. This holds true
@@ -152,13 +153,13 @@ func BreadthFirstSearchParallelEx[N comparable](
 				_, fallback, _ = jobs.EntryAt(int(index))
 			}
 		}
-		nextJobs := collections.NewOrderedMapWithSizeHint[N, *breadthFirstSearchJob[N]](int(nextJobCount.Load()))
+		nextJobs := collections.NewOrderedMapWithSizeHint[K, *breadthFirstSearchJob[N]](int(nextJobCount.Load()))
 		for _, jobs := range next {
 			for _, j := range jobs {
-				if !nextJobs.Has(j.node) {
+				if !nextJobs.Has(getKey(j.node)) {
 					// Deduplicate synchronously to avoid messy locks and spawning
 					// unnecessary goroutines.
-					nextJobs.Set(j.node, j)
+					nextJobs.Set(getKey(j.node), j)
 				}
 			}
 		}
@@ -175,8 +176,8 @@ func BreadthFirstSearchParallelEx[N comparable](
 	}
 
 	levelIndex := 0
-	level := collections.NewOrderedMapFromList([]collections.MapEntry[N, *breadthFirstSearchJob[N]]{
-		{Key: start, Value: &breadthFirstSearchJob[N]{node: start}},
+	level := collections.NewOrderedMapFromList([]collections.MapEntry[K, *breadthFirstSearchJob[N]]{
+		{Key: getKey(start), Value: &breadthFirstSearchJob[N]{node: start}},
 	})
 	for level.Size() > 0 {
 		result := processLevel(levelIndex, level)

internal/core/bfs_test.go

@@ -78,9 +78,10 @@ func TestBreadthFirstSearchParallel(t *testing.T) {
 		var visited collections.SyncSet[string]
 		core.BreadthFirstSearchParallelEx("Root", children, func(node string) (bool, bool) {
 			return node == "L2B", true // Stop at level 2
-		}, core.BreadthFirstSearchOptions[string]{
+		}, core.BreadthFirstSearchOptions[string, string]{
 			Visited: &visited,
-		})
+		},
+			core.Identity)
 
 		assert.Assert(t, visited.Has("Root"), "Expected to visit Root")
 		assert.Assert(t, visited.Has("L1A"), "Expected to visit L1A")
@@ -108,9 +109,10 @@ func TestBreadthFirstSearchParallel(t *testing.T) {
 		var visited collections.SyncSet[string]
 		result := core.BreadthFirstSearchParallelEx("A", children, func(node string) (bool, bool) {
 			return node == "A", false // Record A as a fallback, but do not stop
-		}, core.BreadthFirstSearchOptions[string]{
+		}, core.BreadthFirstSearchOptions[string, string]{
 			Visited: &visited,
-		})
+		},
+			core.Identity)
 
 		assert.Equal(t, result.Stopped, false, "Expected search to not stop early")
 		assert.DeepEqual(t, result.Path, []string{"A"})

internal/project/projectcollection.go

@@ -179,9 +179,10 @@ func (c *ProjectCollection) findDefaultConfiguredProjectWorker(fileName string,
 			}
 			return false, false
 		},
-		core.BreadthFirstSearchOptions[*Project]{
+		core.BreadthFirstSearchOptions[*Project, *Project]{
 			Visited: visited,
 		},
+		core.Identity,
 	)
 
 	if search.Stopped {

internal/project/projectcollectionbuilder.go

@@ -473,6 +473,11 @@ type searchNode struct {
 	logger         *logging.LogTree
 }
 
+type searchNodeKey struct {
+	configFileName string
+	loadKind       projectLoadKind
+}
+
 type searchResult struct {
 	project *dirty.SyncMapEntry[tspath.Path, *Project]
 	retain  collections.Set[tspath.Path]
@@ -483,13 +488,13 @@ func (b *projectCollectionBuilder) findOrCreateDefaultConfiguredProjectWorker(
 	path tspath.Path,
 	configFileName string,
 	loadKind projectLoadKind,
-	visited *collections.SyncSet[searchNode],
+	visited *collections.SyncSet[searchNodeKey],
 	fallback *searchResult,
 	logger *logging.LogTree,
 ) searchResult {
 	var configs collections.SyncMap[tspath.Path, *tsoptions.ParsedCommandLine]
 	if visited == nil {
-		visited = &collections.SyncSet[searchNode]{}
+		visited = &collections.SyncSet[searchNodeKey]{}
 	}
 
 	search := core.BreadthFirstSearchParallelEx(
@@ -558,18 +563,21 @@ func (b *projectCollectionBuilder) findOrCreateDefaultConfiguredProjectWorker(
 			node.logger.Log("Project does not contain file")
 			return false, false
 		},
-		core.BreadthFirstSearchOptions[searchNode]{
+		core.BreadthFirstSearchOptions[searchNodeKey, searchNode]{
 			Visited: visited,
-			PreprocessLevel: func(level *core.BreadthFirstSearchLevel[searchNode]) {
+			PreprocessLevel: func(level *core.BreadthFirstSearchLevel[searchNodeKey, searchNode]) {
 				level.Range(func(node searchNode) bool {
-					if node.loadKind == projectLoadKindFind && level.Has(searchNode{configFileName: node.configFileName, loadKind: projectLoadKindCreate, logger: node.logger}) {
+					if node.loadKind == projectLoadKindFind && level.Has(searchNodeKey{configFileName: node.configFileName, loadKind: projectLoadKindCreate}) {
 						// Remove find requests when a create request for the same project is already present.
-						level.Delete(node)
+						level.Delete(searchNodeKey{configFileName: node.configFileName, loadKind: node.loadKind})
 					}
 					return true
 				})
 			},
 		},
+		func(node searchNode) searchNodeKey {
+			return searchNodeKey{configFileName: node.configFileName, loadKind: node.loadKind}
+		},
 	)
 
 	var retain collections.Set[tspath.Path]
@@ -626,7 +634,7 @@ func (b *projectCollectionBuilder) findOrCreateDefaultConfiguredProjectWorker(
 	// If we didn't find anything, we can retain everything we visited,
 	// since the whole graph must have been traversed (i.e., the set of
 	// retained projects is guaranteed to be deterministic).
-	visited.Range(func(node searchNode) bool {
+	visited.Range(func(node searchNodeKey) bool {
 		retain.Add(b.toPath(node.configFileName))
 		return true
 	})

internal/project/projectcollectionbuilder_test.go

@@ -469,6 +469,49 @@ func TestProjectCollectionBuilder(t *testing.T) {
 			"/project/c.ts",
 		})
 	})
+
+	t.Run("project lookup terminates", func(t *testing.T) {
+		t.Parallel()
+		files := map[string]any{
+			"/tsconfig.json": `{
+				"files": [],
+				"references": [
+					{
+						"path": "./packages/pkg1"
+					},
+					{
+						"path": "./packages/pkg2"
+					},
+				]
+			}`,
+			"/packages/pkg1/tsconfig.json": `{
+				"include": ["src/**/*.ts"],
+				"compilerOptions": {
+					"composite": true,
+				},
+				"references": [
+					{
+						"path": "../pkg2"
+					},
+				]
+			}`,
+			"/packages/pkg2/tsconfig.json": `{
+				"include": ["src/**/*.ts"],
+				"compilerOptions": {
+					"composite": true,
+				},
+				"references": [
+					{
+						"path": "../pkg1"
+					},
+				]
+			}`,
+			"/script.ts": `export const a = 1;`,
+		}
+		session, _ := projecttestutil.Setup(files)
+		session.DidOpenFile(context.Background(), "file:///script.ts", 1, files["/script.ts"].(string), lsproto.LanguageKindTypeScript)
+		// Test should terminate
+	})
 }
 
 func filesForSolutionConfigFile(solutionRefs []string, compilerOptions string, ownFiles []string) map[string]any {