Handle document Registry to distinguish between files with same name and document registry key(compiler options affecting source file) but different ScriptKind (microsoft#43474)

sheetalkamat · web-flow · commit 4b556e37db4b · 2021-04-05T13:27:34.000-07:00
* Test that fails because of change in scriptKind of untitled file * buckets are keyed with DocumentRegistryBucketKey * Use scriptKind in document Registry to distinguish between files Fixes microsoft#42613
diff --git a/src/services/documentRegistry.ts b/src/services/documentRegistry.ts
@@ -83,12 +83,26 @@ namespace ts {
          * @param fileName The name of the file to be released
          * @param compilationSettings The compilation settings used to acquire the file
          */
+        /**@deprecated pass scriptKind for correctness */
         releaseDocument(fileName: string, compilationSettings: CompilerOptions): void;
-
+        /**
+         * Informs the DocumentRegistry that a file is not needed any longer.
+         *
+         * Note: It is not allowed to call release on a SourceFile that was not acquired from
+         * this registry originally.
+         *
+         * @param fileName The name of the file to be released
+         * @param compilationSettings The compilation settings used to acquire the file
+         * @param scriptKind The script kind of the file to be released
+         */
+        releaseDocument(fileName: string, compilationSettings: CompilerOptions, scriptKind: ScriptKind): void; // eslint-disable-line @typescript-eslint/unified-signatures
+        /**
+         * @deprecated pass scriptKind for correctness */
         releaseDocumentWithKey(path: Path, key: DocumentRegistryBucketKey): void;
+        releaseDocumentWithKey(path: Path, key: DocumentRegistryBucketKey, scriptKind: ScriptKind): void; // eslint-disable-line @typescript-eslint/unified-signatures
 
         /*@internal*/
-        getLanguageServiceRefCounts(path: Path): [string, number | undefined][];
+        getLanguageServiceRefCounts(path: Path, scriptKind: ScriptKind): [string, number | undefined][];
 
         reportStats(): string;
     }
@@ -110,6 +124,11 @@ namespace ts {
         languageServiceRefCount: number;
     }
 
+    type BucketEntry = DocumentRegistryEntry | ESMap<ScriptKind, DocumentRegistryEntry>;
+    function isDocumentRegistryEntry(entry: BucketEntry): entry is DocumentRegistryEntry {
+        return !!(entry as DocumentRegistryEntry).sourceFile;
+    }
+
     export function createDocumentRegistry(useCaseSensitiveFileNames?: boolean, currentDirectory?: string): DocumentRegistry {
         return createDocumentRegistryInternal(useCaseSensitiveFileNames, currentDirectory);
     }
@@ -118,18 +137,24 @@ namespace ts {
     export function createDocumentRegistryInternal(useCaseSensitiveFileNames?: boolean, currentDirectory = "", externalCache?: ExternalDocumentCache): DocumentRegistry {
         // Maps from compiler setting target (ES3, ES5, etc.) to all the cached documents we have
         // for those settings.
-        const buckets = new Map<string, ESMap<Path, DocumentRegistryEntry>>();
+        const buckets = new Map<DocumentRegistryBucketKey, ESMap<Path, BucketEntry>>();
         const getCanonicalFileName = createGetCanonicalFileName(!!useCaseSensitiveFileNames);
 
         function reportStats() {
             const bucketInfoArray = arrayFrom(buckets.keys()).filter(name => name && name.charAt(0) === "_").map(name => {
                 const entries = buckets.get(name)!;
-                const sourceFiles: { name: string; refCount: number; }[] = [];
+                const sourceFiles: { name: string; scriptKind: ScriptKind, refCount: number; }[] = [];
                 entries.forEach((entry, name) => {
-                    sourceFiles.push({
-                        name,
-                        refCount: entry.languageServiceRefCount
-                    });
+                    if (isDocumentRegistryEntry(entry)) {
+                        sourceFiles.push({
+                            name,
+                            scriptKind: entry.sourceFile.scriptKind,
+                            refCount: entry.languageServiceRefCount
+                        });
+                    }
+                    else {
+                        entry.forEach((value, scriptKind) => sourceFiles.push({ name, scriptKind, refCount: value.languageServiceRefCount }));
+                    }
                 });
                 sourceFiles.sort((x, y) => y.refCount - x.refCount);
                 return {
@@ -160,6 +185,12 @@ namespace ts {
             return acquireOrUpdateDocument(fileName, path, compilationSettings, key, scriptSnapshot, version, /*acquiring*/ false, scriptKind);
         }
 
+        function getDocumentRegistryEntry(bucketEntry: BucketEntry, scriptKind: ScriptKind | undefined) {
+            const entry = isDocumentRegistryEntry(bucketEntry) ? bucketEntry : bucketEntry.get(Debug.checkDefined(scriptKind, "If there are more than one scriptKind's for same document the scriptKind should be provided"));
+            Debug.assert(scriptKind === undefined || !entry || entry.sourceFile.scriptKind === scriptKind, `Script kind should match provided ScriptKind:${scriptKind} and sourceFile.scriptKind: ${entry?.sourceFile.scriptKind}, !entry: ${!entry}`);
+            return entry;
+        }
+
         function acquireOrUpdateDocument(
             fileName: string,
             path: Path,
@@ -169,10 +200,11 @@ namespace ts {
             version: string,
             acquiring: boolean,
             scriptKind?: ScriptKind): SourceFile {
-
-            const bucket = getOrUpdate(buckets, key, () => new Map<Path, DocumentRegistryEntry>());
-            let entry = bucket.get(path);
+            scriptKind = ensureScriptKind(fileName, scriptKind);
             const scriptTarget = scriptKind === ScriptKind.JSON ? ScriptTarget.JSON : compilationSettings.target || ScriptTarget.ES5;
+            const bucket = getOrUpdate(buckets, key, () => new Map());
+            const bucketEntry = bucket.get(path);
+            let entry = bucketEntry && getDocumentRegistryEntry(bucketEntry, scriptKind);
             if (!entry && externalCache) {
                 const sourceFile = externalCache.getDocument(key, path);
                 if (sourceFile) {
@@ -181,7 +213,7 @@ namespace ts {
                         sourceFile,
                         languageServiceRefCount: 0
                     };
-                    bucket.set(path, entry);
+                    setBucketEntry();
                 }
             }
 
@@ -195,7 +227,7 @@ namespace ts {
                     sourceFile,
                     languageServiceRefCount: 1,
                 };
-                bucket.set(path, entry);
+                setBucketEntry();
             }
             else {
                 // We have an entry for this file.  However, it may be for a different version of
@@ -221,28 +253,53 @@ namespace ts {
             Debug.assert(entry.languageServiceRefCount !== 0);
 
             return entry.sourceFile;
+
+            function setBucketEntry() {
+                if (!bucketEntry) {
+                    bucket.set(path, entry!);
+                }
+                else if (isDocumentRegistryEntry(bucketEntry)) {
+                    const scriptKindMap = new Map<ScriptKind, DocumentRegistryEntry>();
+                    scriptKindMap.set(bucketEntry.sourceFile.scriptKind, bucketEntry);
+                    scriptKindMap.set(scriptKind!, entry!);
+                    bucket.set(path, scriptKindMap);
+                }
+                else {
+                    bucketEntry.set(scriptKind!, entry!);
+                }
+            }
         }
 
-        function releaseDocument(fileName: string, compilationSettings: CompilerOptions): void {
+        function releaseDocument(fileName: string, compilationSettings: CompilerOptions, scriptKind?: ScriptKind): void {
             const path = toPath(fileName, currentDirectory, getCanonicalFileName);
             const key = getKeyForCompilationSettings(compilationSettings);
-            return releaseDocumentWithKey(path, key);
+            return releaseDocumentWithKey(path, key, scriptKind);
         }
 
-        function releaseDocumentWithKey(path: Path, key: DocumentRegistryBucketKey): void {
+        function releaseDocumentWithKey(path: Path, key: DocumentRegistryBucketKey, scriptKind?: ScriptKind): void {
             const bucket = Debug.checkDefined(buckets.get(key));
-            const entry = bucket.get(path)!;
+            const bucketEntry = bucket.get(path)!;
+            const entry = getDocumentRegistryEntry(bucketEntry, scriptKind)!;
             entry.languageServiceRefCount--;
 
             Debug.assert(entry.languageServiceRefCount >= 0);
             if (entry.languageServiceRefCount === 0) {
-                bucket.delete(path);
+                if (isDocumentRegistryEntry(bucketEntry)) {
+                    bucket.delete(path);
+                }
+                else {
+                    bucketEntry.delete(scriptKind!);
+                    if (bucketEntry.size === 1) {
+                        bucket.set(path, firstDefinedIterator(bucketEntry.values(), identity)!);
+                    }
+                }
             }
         }
 
-        function getLanguageServiceRefCounts(path: Path) {
+        function getLanguageServiceRefCounts(path: Path, scriptKind: ScriptKind) {
             return arrayFrom(buckets.entries(), ([key, bucket]): [string, number | undefined] => {
-                const entry = bucket.get(path);
+                const bucketEntry = bucket.get(path);
+                const entry = bucketEntry && getDocumentRegistryEntry(bucketEntry, scriptKind);
                 return [key, entry && entry.languageServiceRefCount];
             });
         }
diff --git a/src/services/services.ts b/src/services/services.ts
@@ -1444,7 +1444,7 @@ namespace ts {
             // not part of the new program.
             function onReleaseOldSourceFile(oldSourceFile: SourceFile, oldOptions: CompilerOptions) {
                 const oldSettingsKey = documentRegistry.getKeyForCompilationSettings(oldOptions);
-                documentRegistry.releaseDocumentWithKey(oldSourceFile.resolvedPath, oldSettingsKey);
+                documentRegistry.releaseDocumentWithKey(oldSourceFile.resolvedPath, oldSettingsKey, oldSourceFile.scriptKind);
             }
 
             function getOrCreateSourceFile(fileName: string, languageVersion: ScriptTarget, onError?: (message: string) => void, shouldCreateNewSourceFile?: boolean): SourceFile | undefined {
@@ -1493,9 +1493,13 @@ namespace ts {
                         // We do not support the scenario where a host can modify a registered
                         // file's script kind, i.e. in one project some file is treated as ".ts"
                         // and in another as ".js"
-                        Debug.assertEqual(hostFileInformation.scriptKind, oldSourceFile.scriptKind, "Registered script kind should match new script kind.");
-
-                        return documentRegistry.updateDocumentWithKey(fileName, path, newSettings, documentRegistryBucketKey, hostFileInformation.scriptSnapshot, hostFileInformation.version, hostFileInformation.scriptKind);
+                        if (hostFileInformation.scriptKind === oldSourceFile.scriptKind) {
+                            return documentRegistry.updateDocumentWithKey(fileName, path, newSettings, documentRegistryBucketKey, hostFileInformation.scriptSnapshot, hostFileInformation.version, hostFileInformation.scriptKind);
+                        }
+                        else {
+                            // Release old source file and fall through to aquire new file with new script kind
+                            documentRegistry.releaseDocumentWithKey(oldSourceFile.resolvedPath, documentRegistry.getKeyForCompilationSettings(program.getCompilerOptions()), oldSourceFile.scriptKind);
+                        }
                     }
 
                     // We didn't already have the file.  Fall through and acquire it from the registry.
@@ -1531,7 +1535,7 @@ namespace ts {
                 // Use paths to ensure we are using correct key and paths as document registry could be created with different current directory than host
                 const key = documentRegistry.getKeyForCompilationSettings(program.getCompilerOptions());
                 forEach(program.getSourceFiles(), f =>
-                    documentRegistry.releaseDocumentWithKey(f.resolvedPath, key));
+                    documentRegistry.releaseDocumentWithKey(f.resolvedPath, key, f.scriptKind));
                 program = undefined!; // TODO: GH#18217
             }
             host = undefined!;
diff --git a/src/testRunner/unittests/tsserver/documentRegistry.ts b/src/testRunner/unittests/tsserver/documentRegistry.ts
@@ -27,7 +27,7 @@ namespace ts.projectSystem {
             assert.isDefined(moduleInfo);
             assert.equal(moduleInfo.isOrphan(), moduleIsOrphan);
             const key = service.documentRegistry.getKeyForCompilationSettings(project.getCompilationSettings());
-            assert.deepEqual(service.documentRegistry.getLanguageServiceRefCounts(moduleInfo.path), [[key, moduleIsOrphan ? undefined : 1]]);
+            assert.deepEqual(service.documentRegistry.getLanguageServiceRefCounts(moduleInfo.path, moduleInfo.scriptKind), [[key, moduleIsOrphan ? undefined : 1]]);
         }
 
         function createServiceAndHost() {
diff --git a/src/testRunner/unittests/tsserver/dynamicFiles.ts b/src/testRunner/unittests/tsserver/dynamicFiles.ts
@@ -121,6 +121,28 @@ var x = 10;`
             service.openClientFile(file.path);
             checkNumberOfProjects(service, { configuredProjects: 1 });
         });
+
+        it("when changing scriptKind of the untitled files", () => {
+            const host = createServerHost([libFile], { useCaseSensitiveFileNames: true });
+            const service = createProjectService(host, { useInferredProjectPerProjectRoot: true });
+            service.openClientFile(untitledFile, "const x = 10;", ScriptKind.TS, tscWatch.projectRoot);
+            checkNumberOfProjects(service, { inferredProjects: 1 });
+            checkProjectActualFiles(service.inferredProjects[0], [untitledFile, libFile.path]);
+            const program = service.inferredProjects[0].getCurrentProgram()!;
+            const sourceFile = program.getSourceFile(untitledFile)!;
+
+            // Close untitled file
+            service.closeClientFile(untitledFile);
+
+            // Open untitled file with different mode
+            service.openClientFile(untitledFile, "const x = 10;", ScriptKind.TSX, tscWatch.projectRoot);
+            checkNumberOfProjects(service, { inferredProjects: 1 });
+            checkProjectActualFiles(service.inferredProjects[0], [untitledFile, libFile.path]);
+            const newProgram = service.inferredProjects[0].getCurrentProgram()!;
+            const newSourceFile = newProgram.getSourceFile(untitledFile)!;
+            assert.notStrictEqual(newProgram, program);
+            assert.notStrictEqual(newSourceFile, sourceFile);
+        });
     });
 
     describe("unittests:: tsserver:: dynamicFiles:: ", () => {
diff --git a/tests/baselines/reference/api/tsserverlibrary.d.ts b/tests/baselines/reference/api/tsserverlibrary.d.ts
@@ -6516,8 +6516,23 @@ declare namespace ts {
          * @param fileName The name of the file to be released
          * @param compilationSettings The compilation settings used to acquire the file
          */
+        /**@deprecated pass scriptKind for correctness */
         releaseDocument(fileName: string, compilationSettings: CompilerOptions): void;
+        /**
+         * Informs the DocumentRegistry that a file is not needed any longer.
+         *
+         * Note: It is not allowed to call release on a SourceFile that was not acquired from
+         * this registry originally.
+         *
+         * @param fileName The name of the file to be released
+         * @param compilationSettings The compilation settings used to acquire the file
+         * @param scriptKind The script kind of the file to be released
+         */
+        releaseDocument(fileName: string, compilationSettings: CompilerOptions, scriptKind: ScriptKind): void;
+        /**
+         * @deprecated pass scriptKind for correctness */
         releaseDocumentWithKey(path: Path, key: DocumentRegistryBucketKey): void;
+        releaseDocumentWithKey(path: Path, key: DocumentRegistryBucketKey, scriptKind: ScriptKind): void;
         reportStats(): string;
     }
     type DocumentRegistryBucketKey = string & {
diff --git a/tests/baselines/reference/api/typescript.d.ts b/tests/baselines/reference/api/typescript.d.ts
@@ -6516,8 +6516,23 @@ declare namespace ts {
          * @param fileName The name of the file to be released
          * @param compilationSettings The compilation settings used to acquire the file
          */
+        /**@deprecated pass scriptKind for correctness */
         releaseDocument(fileName: string, compilationSettings: CompilerOptions): void;
+        /**
+         * Informs the DocumentRegistry that a file is not needed any longer.
+         *
+         * Note: It is not allowed to call release on a SourceFile that was not acquired from
+         * this registry originally.
+         *
+         * @param fileName The name of the file to be released
+         * @param compilationSettings The compilation settings used to acquire the file
+         * @param scriptKind The script kind of the file to be released
+         */
+        releaseDocument(fileName: string, compilationSettings: CompilerOptions, scriptKind: ScriptKind): void;
+        /**
+         * @deprecated pass scriptKind for correctness */
         releaseDocumentWithKey(path: Path, key: DocumentRegistryBucketKey): void;
+        releaseDocumentWithKey(path: Path, key: DocumentRegistryBucketKey, scriptKind: ScriptKind): void;
         reportStats(): string;
     }
     type DocumentRegistryBucketKey = string & {

Original file line number	Diff line number	Diff line change
`@@ -27,7 +27,7 @@ namespace ts.projectSystem {`
`27`	`27`	`assert.isDefined(moduleInfo);`
`28`	`28`	`assert.equal(moduleInfo.isOrphan(), moduleIsOrphan);`
`29`	`29`	`const key = service.documentRegistry.getKeyForCompilationSettings(project.getCompilationSettings());`
`30`		`- assert.deepEqual(service.documentRegistry.getLanguageServiceRefCounts(moduleInfo.path), [[key, moduleIsOrphan ? undefined : 1]]);`
	`30`	`+ assert.deepEqual(service.documentRegistry.getLanguageServiceRefCounts(moduleInfo.path, moduleInfo.scriptKind), [[key, moduleIsOrphan ? undefined : 1]]);`
`31`	`31`	`}`
`32`	`32`
`33`	`33`	`function createServiceAndHost() {`