CodeMapper support

Author: zkat

src/harness/client.ts

@@ -76,7 +76,7 @@ import {
     UserPreferences,
 } from "./_namespaces/ts";
 import {
-    protocol,
+    protocol
 } from "./_namespaces/ts.server";
 
 export interface SessionClientHost extends LanguageServiceHost {
@@ -781,6 +781,8 @@ export class SessionClient implements LanguageService {
         });
     }
 
+    mapCode = notImplemented;
+
     private createFileLocationOrRangeRequestArgs(positionOrRange: number | TextRange, fileName: string): protocol.FileLocationOrRangeRequestArgs {
         return typeof positionOrRange === "number"
             ? this.createFileLocationRequestArgs(fileName, positionOrRange)

src/harness/harnessLanguageService.ts

@@ -698,6 +698,9 @@ class LanguageServiceShimProxy implements ts.LanguageService {
     provideInlayHints(fileName: string, span: ts.TextSpan, preference: ts.UserPreferences) {
         return unwrapJSONCallResult(this.shim.provideInlayHints(fileName, span, preference));
     }
+    mapCode(fileName: string, contents: string[], formatOptions: ts.FormatCodeSettings, preferences: ts.UserPreferences = ts.emptyOptions, focusLocations?: ts.PrioritizedSpan[], updates?: ts.FileTextChanges[]): ts.FileTextChanges[] {
+        return unwrapJSONCallResult(this.shim.mapCode(fileName, contents, formatOptions, preferences, focusLocations, updates));
+    }
     getEmitOutput(fileName: string): ts.EmitOutput {
         return unwrapJSONCallResult(this.shim.getEmitOutput(fileName));
     }

src/server/protocol.ts

@@ -173,6 +173,7 @@ export const enum CommandTypes {
     ProvideCallHierarchyIncomingCalls = "provideCallHierarchyIncomingCalls",
     ProvideCallHierarchyOutgoingCalls = "provideCallHierarchyOutgoingCalls",
     ProvideInlayHints = "provideInlayHints",
+    MapCode = "mapCode",
 }
 
 /**
@@ -2687,6 +2688,35 @@ export interface InlayHintsResponse extends Response {
     body?: InlayHintItem[];
 }
 
+export interface MapCodeRequestArgs extends FileRequestArgs {
+    /// The specific code to map/insert/replace in the file.
+    contents: string[];
+
+    /// Areas of "focus" to inform the code mapper with. For example, cursor
+    /// location, current selection, viewport, etc.
+    focusLocations?: MapCodeFocusLocation[];
+
+    /// Edits to apply to the current workspace before performing the mapping.
+    updates?: FileCodeEdits[]
+}
+
+/// Area of "focus" to inform the code mapper with. For example, cursor
+/// location, current selection, viewport, etc.
+export interface MapCodeFocusLocation extends FileSpan {
+    /// The priority to give this focus location. Multiple FocusLocations can
+    /// share the same priority value.
+    priority: number;
+}
+
+export interface MapCodeRequest extends Request {
+    command: CommandTypes.MapCode,
+    arguments: MapCodeRequestArgs;
+}
+
+export interface MapCodeResponse extends Response {
+    body: FileCodeEdits[]
+}
+
 /**
  * Synchronous request for semantic diagnostics of one file.
  */

src/server/session.ts

@@ -1887,6 +1887,51 @@ export class Session<TMessage = string> implements EventSender {
         });
     }
 
+    private mapCode(args: protocol.MapCodeRequestArgs): protocol.FileCodeEdits[] {
+        const { file, project } = this.getFileAndProject(args);
+        const scriptInfo = this.projectService.getScriptInfoForNormalizedPath(file)!;
+        const formatOptions = this.getHostFormatOptions();
+        const preferences = this.getHostPreferences();
+        const selected = args.focusLocations?.map(loc => {
+            const start = scriptInfo.lineOffsetToPosition(loc.start.line, loc.start.offset);
+            const end = scriptInfo.lineOffsetToPosition(loc.end.line, loc.end.offset);
+            return {
+                fileName: loc.file,
+                priority: loc.priority,
+                textSpan: {
+                    start,
+                    length: end - start,
+                }
+            };
+        });
+        const updates = args.updates?.map(edit => {
+            return {
+                fileName: edit.fileName,
+                textChanges: edit.textChanges.map(({ start, end, newText }) => {
+                    const newStart = scriptInfo.lineOffsetToPosition(start.line, start.offset);
+                    const newEnd = scriptInfo.lineOffsetToPosition(end.line, end.offset);
+                    return {
+                        span: { start: newStart, length: newEnd - newStart },
+                        newText,
+                    };
+                }),
+            };
+        });
+        return project.getLanguageService().mapCode(file, args.contents, formatOptions, preferences, selected, updates)?.map(change => {
+            return {
+                fileName: change.fileName,
+                textChanges: change.textChanges.map(({ span, newText }) => {
+                    const newSpan = toProtocolTextSpan(span, scriptInfo);
+                    return {
+                        start: newSpan.start,
+                        end: newSpan.end,
+                        newText
+                    };
+                }),
+            };
+        });
+    }
+
     private setCompilerOptionsForInferredProjects(args: protocol.SetCompilerOptionsForInferredProjectsArgs): void {
         this.projectService.setCompilerOptionsForInferredProjects(args.options, args.projectRootPath);
     }
@@ -3556,6 +3601,9 @@ export class Session<TMessage = string> implements EventSender {
         [protocol.CommandTypes.ProvideInlayHints]: (request: protocol.InlayHintsRequest) => {
             return this.requiredResponse(this.provideInlayHints(request.arguments));
         },
+        [protocol.CommandTypes.MapCode]: (request: protocol.MapCodeRequest) => {
+            return this.requiredResponse(this.mapCode(request.arguments));
+        },
     }));
 
     public addProtocolHandler(command: string, handler: (request: protocol.Request) => HandlerResponse) {

src/services/_namespaces/ts.MapCode.ts

@@ -0,0 +1,4 @@
+/* Generated file to emulate the ts.MapCode namespace. */
+
+export * from "../mapCode";
+

src/services/_namespaces/ts.ts

@@ -33,6 +33,8 @@ import * as GoToDefinition from "./ts.GoToDefinition";
 export { GoToDefinition };
 import * as InlayHints from "./ts.InlayHints";
 export { InlayHints };
+import * as MapCode from "./ts.MapCode";
+export { MapCode };
 import * as JsDoc from "./ts.JsDoc";
 export { JsDoc };
 import * as NavigateTo from "./ts.NavigateTo";

src/services/mapCode.ts

@@ -0,0 +1,33 @@
+import {
+    FileTextChanges,
+    formatting,
+    LanguageServiceHost,
+    PrioritizedSpan,
+    Program,
+    SourceFile,
+    SourceMapper,
+    textChanges,
+    UserPreferences,
+} from "./_namespaces/ts";
+
+// useful stuff:
+// textChanges.ChangeTracker - lots of methods for inserting nodes in the right place.
+// textChanges.ChangeTracker.pushRaw() - to push `updates` before running our analysis.
+
+/** @internal */
+export function mapCode(
+    _file: SourceFile,
+    _program: Program,
+    host: LanguageServiceHost,
+    formatContext: formatting.FormatContext,
+    preferences: UserPreferences,
+    _sourceMapper: SourceMapper,
+    _contents: string[],
+    _focusLocations?: PrioritizedSpan[],
+    _updates?: FileTextChanges[]
+): FileTextChanges[] {
+    // const _checker = program.getTypeChecker();
+    return textChanges.ChangeTracker.with({ host, formatContext, preferences }, _changeTracker => {
+        throw new Error("IT WORKED!!");
+    });
+}

src/services/services.ts

@@ -207,6 +207,7 @@ import {
     LinkedEditingInfo,
     LiteralType,
     map,
+    MapCode,
     mapDefined,
     MapLike,
     mapOneOrMany,
@@ -242,6 +243,7 @@ import {
     positionIsSynthesized,
     PossibleProgramFileInfo,
     PragmaMap,
+    PrioritizedSpan,
     PrivateIdentifier,
     Program,
     PropertyName,
@@ -3070,6 +3072,12 @@ export function createLanguageService(
         return InlayHints.provideInlayHints(getInlayHintsContext(sourceFile, span, preferences));
     }
 
+    function mapCode(fileName: string, contents: string[], formatOptions: FormatCodeSettings, preferences: UserPreferences = emptyOptions, focusLocations?: PrioritizedSpan[], updates?: FileTextChanges[]): FileTextChanges[] {
+        synchronizeHostData();
+        const sourceFile = getValidSourceFile(fileName);
+        return MapCode.mapCode(sourceFile, program, host, formatting.getFormatContext(formatOptions, host), preferences, sourceMapper, contents, focusLocations, updates);
+    }
+
     const ls: LanguageService = {
         dispose,
         cleanupSemanticCache,
@@ -3140,6 +3148,7 @@ export function createLanguageService(
         uncommentSelection,
         provideInlayHints,
         getSupportedCodeFixes,
+        mapCode,
     };
 
     switch (languageServiceMode) {

src/services/shims.ts

@@ -21,6 +21,7 @@ import {
     Extension,
     extensionFromPath,
     FileReference,
+    FileTextChanges,
     filter,
     flattenDiagnosticMessageText,
     FormatCodeOptions,
@@ -48,6 +49,7 @@ import {
     parseJsonSourceFileConfigFileContent,
     parseJsonText,
     preProcessFile,
+    PrioritizedSpan,
     ResolvedModuleFull,
     ResolvedTypeReferenceDirective,
     resolveModuleName,
@@ -361,6 +363,7 @@ export interface LanguageServiceShim extends Shim {
     provideCallHierarchyIncomingCalls(fileName: string, position: number): string;
     provideCallHierarchyOutgoingCalls(fileName: string, position: number): string;
     provideInlayHints(fileName: string, span: TextSpan, preference: UserPreferences | undefined): string;
+    mapCode(fileName: string, contents: string[], formatOptions: FormatCodeSettings, preferences: UserPreferences, focusLocations?: PrioritizedSpan[], updates?: FileTextChanges[]): string;
     getEmitOutput(fileName: string): string;
     getEmitOutputObject(fileName: string): EmitOutput;
 
@@ -1157,6 +1160,13 @@ class LanguageServiceShimObject extends ShimBase implements LanguageServiceShim
         );
     }
 
+    public mapCode(fileName: string, contents: string[], formatOptions: FormatCodeSettings, preferences: UserPreferences = emptyOptions, focusLocations?: PrioritizedSpan[], updates?: FileTextChanges[]): string {
+        return this.forwardJSONCall(
+            `mapCode('${fileName}', '${JSON.stringify(contents)}', '${JSON.stringify(formatOptions)}', '${JSON.stringify(preferences)}', '${JSON.stringify(focusLocations)}', '${JSON.stringify(updates)}')`,
+            () => this.languageService.mapCode(fileName, contents, formatOptions, preferences, focusLocations, updates),
+        )
+    }
+
     /// Emit
     public getEmitOutput(fileName: string): string {
         return this.forwardJSONCall(

src/services/types.ts

@@ -676,6 +676,8 @@ export interface LanguageService {
 
     getSupportedCodeFixes(fileName?: string): readonly string[];
 
+    mapCode(fileName: string, contents: string[], formatOptions: FormatCodeSettings, preferences: UserPreferences, focusLocations?: PrioritizedSpan[], updates?: FileTextChanges[]): FileTextChanges[];
+
     dispose(): void;
 }
 
@@ -1117,6 +1119,12 @@ export interface EditorSettings {
     trimTrailingWhitespace?: boolean;
 }
 
+export interface PrioritizedSpan {
+    fileName: string,
+    textSpan: TextSpan,
+    priority: number,
+}
+
 /** @deprecated - consider using FormatCodeSettings instead */
 export interface FormatCodeOptions extends EditorOptions {
     InsertSpaceAfterCommaDelimiter: boolean;