Fix getReferencedSymbolsForSymbol, implement getReferencesAtExportSpecifier (#2304)

Author: jakebailey

internal/fourslash/_scripts/failingTests.txt

@@ -539,7 +539,6 @@ TestQuickinfoWrongComment
 TestRecursiveInternalModuleImport
 TestReferencesInEmptyFile
 TestRegexDetection
-TestRenameForAliasingExport02
 TestRenameFromNodeModulesDep1
 TestRenameFromNodeModulesDep2
 TestRenameFromNodeModulesDep3

internal/fourslash/fourslash.go

@@ -616,8 +616,12 @@ func sendRequest[Params, Resp any](t *testing.T, f *FourslashTest, info lsproto.
 	if resMsg == nil {
 		t.Fatalf(prefix+"Nil response received for %s request", info.Method)
 	}
+	resp := resMsg.AsResponse()
+	if resp.Error != nil {
+		t.Fatalf(prefix+"%s request returned error: %s", info.Method, resp.Error.String())
+	}
 	if !resultOk {
-		t.Fatalf(prefix+"Unexpected %s response type: %T", info.Method, resMsg.AsResponse().Result)
+		t.Fatalf(prefix+"Unexpected %s response type: %T", info.Method, resp.Result)
 	}
 	return result
 }

internal/fourslash/tests/gen/renameForAliasingExport02_test.go

@@ -9,7 +9,7 @@ import (
 
 func TestRenameForAliasingExport02(t *testing.T) {
 	t.Parallel()
-	t.Skip()
+
 	defer testutil.RecoverAndFail(t, "Panic on fourslash test")
 	const content = `// @Filename: foo.ts
 let x = 1;

internal/fourslash/tests/statefindallrefs_test.go

@@ -1041,6 +1041,78 @@ import * as shared from "../../shared/dist"
 	}
 }
 
+func TestFindAllRefsReExportInMultiProjectSolution(t *testing.T) {
+	t.Parallel()
+	defer testutil.RecoverAndFail(t, "Panic on fourslash test")
+	content := `
+// @stateBaseline: true
+// @Filename: /tsconfig.base.json
+{
+	"compilerOptions": {
+		"rootDir": ".",
+		"outDir": "target",
+		"module": "ESNext",
+		"moduleResolution": "bundler",
+		"composite": true,
+		"declaration": true,
+		"strict": true
+	},
+	"include": []
+}
+// @Filename: /tsconfig.json
+{
+	"extends": "./tsconfig.base.json",
+	"references": [
+		{ "path": "project-a" },
+		{ "path": "project-b" },
+		{ "path": "project-c" },
+	]
+}
+// @Filename: /project-a/tsconfig.json
+{
+	"extends": "../tsconfig.base.json",
+	"include": ["*"]
+}
+// @Filename: /project-a/private.ts
+export const /*symbolA*/symbolA = 'some-symbol';
+console.log(symbolA);
+// @Filename: /project-a/public.ts
+export { symbolA } from './private';
+// @Filename: /project-b/tsconfig.json
+{
+	"extends": "../tsconfig.base.json",
+	"include": ["*"]
+}
+// @Filename: /project-b/public.ts
+export const /*symbolB*/symbolB = 'symbol-b';
+// @Filename: /project-c/tsconfig.json
+{
+	"extends": "../tsconfig.base.json",
+	"include": ["*"],
+	"references": [
+		{ "path": "../project-a" },
+		{ "path": "../project-b" },
+	]
+}
+// @Filename: /project-c/index.ts
+import { symbolB } from '../project-b/public';
+import { /*symbolAUsage*/symbolA } from '../project-a/public';
+console.log(symbolB);
+console.log(symbolA);
+`
+	f, done := fourslash.NewFourslash(t, nil /*capabilities*/, content)
+	defer done()
+
+	// Find all refs for symbolA - should find definition in private.ts, re-export in public.ts, and usage in project-c/index.ts
+	f.VerifyBaselineFindAllReferences(t, "symbolA")
+
+	// Find all refs for symbolB - should find definition and usage (no re-export involved)
+	f.VerifyBaselineFindAllReferences(t, "symbolB")
+
+	// Find all refs from the usage site - should also work
+	f.VerifyBaselineFindAllReferences(t, "symbolAUsage")
+}
+
 func TestFindAllRefsDeclarationInOtherProject(t *testing.T) {
 	t.Parallel()
 	type testCase struct {

internal/ls/findallreferences.go

@@ -1528,14 +1528,12 @@ func getReferencedSymbolsForSymbol(originalSymbol *ast.Symbol, node *ast.Node, s
 	state := newState(sourceFiles, sourceFilesSet, node, checker /*, cancellationToken*/, searchMeaning, options)
 
 	var exportSpecifier *ast.Node
-	if !isForRenameWithPrefixAndSuffixText(options) || len(symbol.Declarations) == 0 {
+	if isForRenameWithPrefixAndSuffixText(options) && len(symbol.Declarations) != 0 {
 		exportSpecifier = core.Find(symbol.Declarations, ast.IsExportSpecifier)
 	}
 	if exportSpecifier != nil {
-		// !!! not implemented
-
 		// When renaming at an export specifier, rename the export and not the thing being exported.
-		// state.getReferencesAtExportSpecifier(exportSpecifier.Name(), symbol, exportSpecifier.AsExportSpecifier(), state.createSearch(node, originalSymbol, comingFromUnknown /*comingFrom*/, "", nil), true /*addReferencesHere*/, true /*alwaysGetReferences*/)
+		state.getReferencesAtExportSpecifier(exportSpecifier.Name(), symbol, exportSpecifier.AsExportSpecifier(), state.createSearch(node, originalSymbol, ImpExpKindUnknown /*comingFrom*/, "", nil), true /*addReferencesHere*/, true /*alwaysGetReferences*/)
 	} else if node != nil && node.Kind == ast.KindDefaultKeyword && symbol.Name == ast.InternalSymbolNameDefault && symbol.Parent != nil {
 		state.addReference(node, symbol, entryKindNode)
 		state.searchForImportsOfExport(node, symbol, &ExportInfo{exportingModuleSymbol: symbol.Parent, exportKind: ExportKindDefault})
@@ -1583,22 +1581,21 @@ type refState struct {
 	result                       []*SymbolAndEntries
 	inheritsFromCache            map[inheritKey]bool
 	seenContainingTypeReferences collections.Set[*ast.Node] // node seen tracker
-	// seenReExportRHS           *collections.Set[*ast.Node] // node seen tracker
-	importTracker           ImportTracker
-	symbolToReferences      map[*ast.Symbol]*SymbolAndEntries
-	sourceFileToSeenSymbols map[*ast.SourceFile]*collections.Set[*ast.Symbol]
+	seenReExportRHS              collections.Set[*ast.Node] // node seen tracker
+	importTracker                ImportTracker
+	symbolToReferences           map[*ast.Symbol]*SymbolAndEntries
+	sourceFileToSeenSymbols      map[*ast.SourceFile]*collections.Set[*ast.Symbol]
 }
 
 func newState(sourceFiles []*ast.SourceFile, sourceFilesSet *collections.Set[string], node *ast.Node, checker *checker.Checker, searchMeaning ast.SemanticMeaning, options refOptions) *refState {
 	return &refState{
-		sourceFiles:       sourceFiles,
-		sourceFilesSet:    sourceFilesSet,
-		specialSearchKind: getSpecialSearchKind(node),
-		checker:           checker,
-		searchMeaning:     searchMeaning,
-		options:           options,
-		inheritsFromCache: map[inheritKey]bool{},
-		// seenReExportRHS:           &collections.Set[*ast.Node]{},
+		sourceFiles:             sourceFiles,
+		sourceFilesSet:          sourceFilesSet,
+		specialSearchKind:       getSpecialSearchKind(node),
+		checker:                 checker,
+		searchMeaning:           searchMeaning,
+		options:                 options,
+		inheritsFromCache:       map[inheritKey]bool{},
 		symbolToReferences:      map[*ast.Symbol]*SymbolAndEntries{},
 		sourceFileToSeenSymbols: map[*ast.SourceFile]*collections.Set[*ast.Symbol]{},
 	}
@@ -1911,9 +1908,7 @@ func (state *refState) getReferencesAtLocation(sourceFile *ast.SourceFile, posit
 	}
 
 	if parent.Kind == ast.KindExportSpecifier {
-		// !!! not implemented
-		// debug.Assert(referenceLocation.Kind == ast.KindIdentifier || referenceLocation.Kind == ast.KindStringLiteral)
-		// state.getReferencesAtExportSpecifier(referenceLocation /* Identifier | StringLiteral*/, referenceSymbol, parent.AsExportSpecifier(), search, addReferencesHere, false /*alwaysGetReferences*/)
+		state.getReferencesAtExportSpecifier(referenceLocation, referenceSymbol, parent.AsExportSpecifier(), search, addReferencesHere, false /*alwaysGetReferences*/)
 		return
 	}
 
@@ -2034,6 +2029,82 @@ func (state *refState) getImportOrExportReferences(referenceLocation *ast.Node,
 	}
 }
 
+func (state *refState) markSeenReExportRHS(node *ast.Node) bool {
+	return state.seenReExportRHS.AddIfAbsent(node)
+}
+
+func (state *refState) getReferencesAtExportSpecifier(
+	referenceLocation *ast.Node,
+	referenceSymbol *ast.Symbol,
+	exportSpecifier *ast.ExportSpecifier,
+	search *refSearch,
+	addReferencesHere bool,
+	alwaysGetReferences bool,
+) {
+	debug.Assert(!alwaysGetReferences || state.options.useAliasesForRename, "If alwaysGetReferences is true, then prefix/suffix text must be enabled")
+
+	exportDeclaration := exportSpecifier.Parent.Parent.AsExportDeclaration()
+	propertyName := exportSpecifier.PropertyName
+	name := exportSpecifier.Name()
+	localSymbol := getLocalSymbolForExportSpecifier(referenceLocation.AsIdentifier(), referenceSymbol, exportSpecifier, state.checker)
+
+	if !alwaysGetReferences && !search.includes(localSymbol) {
+		return
+	}
+
+	addRef := func() {
+		if addReferencesHere {
+			state.addReference(referenceLocation, localSymbol, entryKindNode)
+		}
+	}
+
+	if propertyName == nil {
+		// Don't rename at `export { default } from "m";`. (but do continue to search for imports of the re-export)
+		if !(state.options.use == referenceUseRename && ast.ModuleExportNameIsDefault(name)) {
+			addRef()
+		}
+	} else if referenceLocation == propertyName.AsNode() {
+		// For `export { foo as bar } from "baz"`, "`foo`" will be added from the singleReferences for import searches of the original export.
+		// For `export { foo as bar };`, where `foo` is a local, so add it now.
+		if exportDeclaration.ModuleSpecifier == nil {
+			addRef()
+		}
+
+		if addReferencesHere && state.options.use != referenceUseRename && state.markSeenReExportRHS(name) {
+			exportSymbol := exportSpecifier.AsNode().Symbol()
+			debug.Assert(exportSymbol != nil, "exportSpecifier.Symbol() should not be nil")
+			state.addReference(name, exportSymbol, entryKindNode)
+		}
+	} else {
+		if state.markSeenReExportRHS(referenceLocation) {
+			addRef()
+		}
+	}
+
+	// For `export { foo as bar }`, rename `foo`, but not `bar`.
+	if !isForRenameWithPrefixAndSuffixText(state.options) || alwaysGetReferences {
+		isDefaultExport := ast.ModuleExportNameIsDefault(referenceLocation) || ast.ModuleExportNameIsDefault(exportSpecifier.Name())
+		exportKind := ExportKindNamed
+		if isDefaultExport {
+			exportKind = ExportKindDefault
+		}
+		exportSymbol := exportSpecifier.AsNode().Symbol()
+		debug.Assert(exportSymbol != nil, "exportSpecifier.Symbol() should not be nil")
+		exportInfo := getExportInfo(exportSymbol, exportKind, state.checker)
+		if exportInfo != nil {
+			state.searchForImportsOfExport(referenceLocation, exportSymbol, exportInfo)
+		}
+	}
+
+	// At `export { x } from "foo"`, also search for the imported symbol `"foo".x`.
+	if search.comingFrom != ImpExpKindExport && exportDeclaration.ModuleSpecifier != nil && propertyName == nil && !isForRenameWithPrefixAndSuffixText(state.options) {
+		imported := state.checker.GetExportSpecifierLocalTargetSymbol(exportSpecifier.AsNode())
+		if imported != nil {
+			state.searchForImportedSymbol(imported)
+		}
+	}
+}
+
 // Go to the symbol we imported from and find references for it.
 func (state *refState) searchForImportedSymbol(symbol *ast.Symbol) {
 	for _, declaration := range symbol.Declarations {

testdata/baselines/reference/fourslash/documentHighlights/documentHighlightInExport1.baseline.jsonc

@@ -1,7 +1,7 @@
 // === documentHighlights ===
 // === /documentHighlightInExport1.ts ===
 // class /*HIGHLIGHTS*/[|C|] {}
-// export { C as D };
+// export { [|C|] as [|D|] };
 
 
 
@@ -15,7 +15,7 @@
 // === documentHighlights ===
 // === /documentHighlightInExport1.ts ===
 // class [|C|] {}
-// export { /*HIGHLIGHTS*/C as D };
+// export { /*HIGHLIGHTS*/[|C|] as [|D|] };
 
 
 
@@ -29,4 +29,4 @@
 // === documentHighlights ===
 // === /documentHighlightInExport1.ts ===
 // class C {}
-// export { C as /*HIGHLIGHTS*/D };
+// export { C as /*HIGHLIGHTS*/[|D|] };

testdata/baselines/reference/fourslash/documentHighlights/documentHighlightInTypeExport.baseline.jsonc

@@ -8,14 +8,14 @@
 // === documentHighlights ===
 // === /1.ts ===
 // type A = 1;
-// export { /*HIGHLIGHTS*/A as B };
+// export { /*HIGHLIGHTS*/[|A|] as [|B|] };
 
 
 
 // === documentHighlights ===
 // === /1.ts ===
 // type A = 1;
-// export { A as /*HIGHLIGHTS*/B };
+// export { A as /*HIGHLIGHTS*/[|B|] };
 
 
 
@@ -31,7 +31,7 @@
 // === /2.ts ===
 // type A = 1;
 // let /*HIGHLIGHTS*/[|A|]: A = 1;
-// export { A as B };
+// export { [|A|] as [|B|] };
 
 
 
@@ -47,15 +47,15 @@
 // === /2.ts ===
 // type A = 1;
 // let [|A|]: A = 1;
-// export { /*HIGHLIGHTS*/A as B };
+// export { /*HIGHLIGHTS*/[|A|] as [|B|] };
 
 
 
 // === documentHighlights ===
 // === /2.ts ===
 // type A = 1;
 // let A: A = 1;
-// export { A as /*HIGHLIGHTS*/B };
+// export { A as /*HIGHLIGHTS*/[|B|] };
 
 
 
@@ -71,7 +71,7 @@
 // === /3.ts ===
 // type A = 1;
 // let /*HIGHLIGHTS*/[|A|]: A = 1;
-// export type { A as B };
+// export type { [|A|] as [|B|] };
 
 
 
@@ -87,12 +87,12 @@
 // === /3.ts ===
 // type A = 1;
 // let [|A|]: A = 1;
-// export type { /*HIGHLIGHTS*/A as B };
+// export type { /*HIGHLIGHTS*/[|A|] as [|B|] };
 
 
 
 // === documentHighlights ===
 // === /3.ts ===
 // type A = 1;
 // let A: A = 1;
-// export type { A as /*HIGHLIGHTS*/B };
+// export type { A as /*HIGHLIGHTS*/[|B|] };

testdata/baselines/reference/fourslash/documentHighlights/documentHighlightsInvalidGlobalThis.baseline.jsonc

@@ -1,5 +1,5 @@
 // === documentHighlights ===
 // === /documentHighlightsInvalidGlobalThis.ts ===
 // declare global {
-//     export { globalThis as /*HIGHLIGHTS*/global }
+//     export { globalThis as /*HIGHLIGHTS*/[|global|] }
 // }

testdata/baselines/reference/fourslash/findAllReferences/findAllReferencesNonExistentExportBinding.baseline.jsonc

@@ -1,3 +1,6 @@
 // === findAllReferences ===
 // === /bar.ts ===
-// import { [|Foo|]/*FIND ALL REFS*/ } from "./foo";
+// import { [|Foo|]/*FIND ALL REFS*/ } from "./foo";
+
+// === /foo.ts ===
+// export { [|Foo|] }

testdata/baselines/reference/fourslash/findAllReferences/findAllRefsCommonJsRequire.baseline.jsonc

@@ -1,7 +1,7 @@
 // === findAllReferences ===
 // === /a.js ===
 // function [|f|]() { }
-// export { f }
+// export { [|f|] }
 
 // === /b.js ===
 // const { [|f|] } = require('./a')