Set parents of augmented module exports (#59609)

Co-authored-by: Isabel Duan <isabelduan@microsoft.com>

src/compiler/checker.ts

@@ -2683,7 +2683,7 @@ export function createTypeChecker(host: TypeCheckerHost): TypeChecker {
             }
             if (source.exports) {
                 if (!target.exports) target.exports = createSymbolTable();
-                mergeSymbolTable(target.exports, source.exports, unidirectional);
+                mergeSymbolTable(target.exports, source.exports, unidirectional, target);
             }
             if (!unidirectional) {
                 recordMergedSymbol(target, source);
@@ -2772,10 +2772,29 @@ export function createTypeChecker(host: TypeCheckerHost): TypeChecker {
         return combined;
     }
 
-    function mergeSymbolTable(target: SymbolTable, source: SymbolTable, unidirectional = false) {
+    function mergeSymbolTable(target: SymbolTable, source: SymbolTable, unidirectional = false, mergedParent?: Symbol) {
         source.forEach((sourceSymbol, id) => {
             const targetSymbol = target.get(id);
-            target.set(id, targetSymbol ? mergeSymbol(targetSymbol, sourceSymbol, unidirectional) : getMergedSymbol(sourceSymbol));
+            const merged = targetSymbol ? mergeSymbol(targetSymbol, sourceSymbol, unidirectional) : getMergedSymbol(sourceSymbol);
+            if (mergedParent && targetSymbol) {
+                // If a merge was performed on the target symbol, set its parent to the merged parent that initiated the merge
+                // of its exports. Otherwise, `merged` came only from `sourceSymbol` and can keep its parent:
+                //
+                // // a.ts
+                // export interface A { x: number; }
+                //
+                // // b.ts
+                // declare module "./a" {
+                //   interface A { y: number; }
+                //   interface B {}
+                // }
+                //
+                // When merging the module augmentation into a.ts, the symbol for `A` will itself be merged, so its parent
+                // should be the merged module symbol. But the symbol for `B` has only one declaration, so its parent should
+                // be the module augmentation symbol, which contains its only declaration.
+                merged.parent = mergedParent;
+            }
+            target.set(id, merged);
         });
     }
 

tests/cases/fourslash/autoImportCompletionExportListAugmentation1.ts

@@ -0,0 +1,58 @@
+/// <reference path="fourslash.ts" />
+// @module: nodenext
+
+// @Filename: /node_modules/@sapphire/pieces/index.d.ts
+//// interface Container {
+////   stores: unknown;
+//// }
+////
+//// declare class Piece {
+////   container: Container;
+//// }
+////
+//// export { Piece, type Container };
+
+// @FileName: /augmentation.ts
+//// declare module "@sapphire/pieces" {
+////   interface Container {
+////     client: unknown;
+////   }
+////   export { Container };
+//// }
+
+// @Filename: /index.ts
+//// import { Piece } from "@sapphire/pieces";
+//// class FullPiece extends Piece {
+////   /*1*/
+//// }
+
+const preferences = {
+  includeCompletionsWithClassMemberSnippets: true,
+  includeCompletionsWithInsertText: true,
+};
+
+verify.completions({
+  marker: "1",
+  includes: [
+    {
+      name: "container",
+      insertText: "container: Container;",
+      filterText: "container",
+      hasAction: true,
+      source: "ClassMemberSnippet/",
+    },
+  ],
+  preferences,
+  isNewIdentifierLocation: true,
+});
+
+verify.applyCodeActionFromCompletion("1", {
+  name: "container",
+  source: "ClassMemberSnippet/",
+  description: `Includes imports of types referenced by 'container'`,
+  newFileContent: `import { Container, Piece } from "@sapphire/pieces";
+class FullPiece extends Piece {
+  
+}`,
+  preferences,
+});

tests/cases/fourslash/autoImportCompletionExportListAugmentation2.ts

@@ -0,0 +1,68 @@
+/// <reference path="fourslash.ts" />
+// @module: nodenext
+
+// @Filename: /node_modules/@sapphire/pieces/index.d.ts
+//// interface Container {
+////   stores: unknown;
+//// }
+////
+//// declare class Piece {
+////   get container(): Container;
+//// }
+////
+//// declare class AliasPiece extends Piece {}
+////
+//// export { AliasPiece, type Container };
+
+// @Filename: /node_modules/@sapphire/framework/index.d.ts
+//// import { AliasPiece } from "@sapphire/pieces";
+////
+//// declare class Command extends AliasPiece {}
+////
+//// declare module "@sapphire/pieces" {
+////   interface Container {
+////     client: unknown;
+////   }
+//// }
+////
+//// export { Command };
+
+// @Filename: /index.ts
+//// import "@sapphire/pieces";
+//// import { Command } from "@sapphire/framework";
+//// class PingCommand extends Command {
+////   /*1*/
+//// }
+
+const preferences = {
+  includeCompletionsWithClassMemberSnippets: true,
+  includeCompletionsWithInsertText: true,
+};
+
+verify.completions({
+  marker: "1",
+  includes: [
+    {
+      name: "container",
+      insertText: "get container(): Container {\n}",
+      filterText: "container",
+      hasAction: true,
+      source: "ClassMemberSnippet/",
+    },
+  ],
+  preferences,
+  isNewIdentifierLocation: true,
+});
+
+verify.applyCodeActionFromCompletion("1", {
+  name: "container",
+  source: "ClassMemberSnippet/",
+  description: `Includes imports of types referenced by 'container'`,
+  newFileContent: `import "@sapphire/pieces";
+import { Command } from "@sapphire/framework";
+import { Container } from "@sapphire/pieces";
+class PingCommand extends Command {
+  
+}`,
+  preferences,
+});

tests/cases/fourslash/autoImportCompletionExportListAugmentation3.ts

@@ -0,0 +1,57 @@
+/// <reference path="fourslash.ts" />
+// @module: nodenext
+
+// @Filename: /node_modules/@sapphire/pieces/index.d.ts
+//// export interface Container {
+////   stores: unknown;
+//// }
+////
+//// declare class Piece {
+////   container: Container;
+//// }
+////
+//// export { Piece };
+
+// @FileName: /augmentation.ts
+//// declare module "@sapphire/pieces" {
+////   interface Container {
+////     client: unknown;
+////   }
+//// }
+
+// @Filename: /index.ts
+//// import { Piece } from "@sapphire/pieces";
+//// class FullPiece extends Piece {
+////   /*1*/
+//// }
+
+const preferences = {
+  includeCompletionsWithClassMemberSnippets: true,
+  includeCompletionsWithInsertText: true,
+};
+
+verify.completions({
+  marker: "1",
+  includes: [
+    {
+      name: "container",
+      insertText: "container: Container;",
+      filterText: "container",
+      hasAction: true,
+      source: "ClassMemberSnippet/",
+    },
+  ],
+  preferences,
+  isNewIdentifierLocation: true,
+});
+
+verify.applyCodeActionFromCompletion("1", {
+  name: "container",
+  source: "ClassMemberSnippet/",
+  description: `Includes imports of types referenced by 'container'`,
+  newFileContent: `import { Container, Piece } from "@sapphire/pieces";
+class FullPiece extends Piece {
+  
+}`,
+  preferences,
+});

tests/cases/fourslash/autoImportCompletionExportListAugmentation4.ts

@@ -0,0 +1,67 @@
+/// <reference path="fourslash.ts" />
+
+// @module: nodenext
+
+// @Filename: /node_modules/@sapphire/pieces/index.d.ts
+//// interface Container {
+////   stores: unknown;
+//// }
+////
+//// declare class Piece {
+////   get container(): Container;
+//// }
+////
+//// export { Piece as Alias, type Container };
+
+// @Filename: /node_modules/@sapphire/framework/index.d.ts
+//// import { Alias } from "@sapphire/pieces";
+////
+//// declare class Command extends Alias {}
+////
+//// declare module "@sapphire/pieces" {
+////   interface Container {
+////     client: unknown;
+////   }
+//// }
+////
+//// export { Command as CommandAlias };
+
+// @Filename: /index.ts
+//// import "@sapphire/pieces";
+//// import { CommandAlias } from "@sapphire/framework";
+//// class PingCommand extends CommandAlias {
+////   /*1*/
+//// }
+
+const preferences = {
+  includeCompletionsWithClassMemberSnippets: true,
+  includeCompletionsWithInsertText: true,
+};
+
+verify.completions({
+  marker: "1",
+  includes: [
+    {
+      name: "container",
+      insertText: "get container(): Container {\n}",
+      filterText: "container",
+      hasAction: true,
+      source: "ClassMemberSnippet/",
+    },
+  ],
+  preferences,
+  isNewIdentifierLocation: true,
+});
+
+verify.applyCodeActionFromCompletion("1", {
+  name: "container",
+  source: "ClassMemberSnippet/",
+  description: `Includes imports of types referenced by 'container'`,
+  newFileContent: `import "@sapphire/pieces";
+import { CommandAlias } from "@sapphire/framework";
+import { Container } from "@sapphire/pieces";
+class PingCommand extends CommandAlias {
+  
+}`,
+  preferences,
+});