Handle buffers throughout script snapshots

src/server/scriptInfo.ts

@@ -20,6 +20,7 @@ namespace ts.server {
          * Only on edits to the script version cache, the text will be set to undefined
          */
         private text: string | undefined;
+        private buffer: Uint8Array | undefined;
         /**
          * Line map for the text when there is no script version cache present
          */
@@ -164,7 +165,7 @@ namespace ts.server {
         public getSnapshot(): IScriptSnapshot {
             return this.useScriptVersionCacheIfValidOrOpen()
                 ? this.svc!.getSnapshot()
-                : ScriptSnapshot.fromString(this.getOrLoadText());
+                : ScriptSnapshot.fromString(this.getOrLoadText(), this.buffer);
         }
 
         public getAbsolutePositionAndLineText(line: number): AbsolutePositionAndLineText {
@@ -208,6 +209,9 @@ namespace ts.server {
             let text: string;
             const fileName = tempFileName || this.info.fileName;
             const getText = () => text === undefined ? (text = this.host.readFile(fileName) || "") : text;
+            if (endsWith(fileName, ".wasm")) {
+                this.buffer = this.host.readFileBuffer(fileName);
+            }
             // Only non typescript files have size limitation
             if (!hasTSFileExtension(this.info.fileName)) {
                 const fileSize = this.host.getFileSize ? this.host.getFileSize(fileName) : getText().length;

src/server/scriptVersionCache.ts

@@ -370,6 +370,10 @@ namespace ts.server {
             return this.index.getText(rangeStart, rangeEnd - rangeStart);
         }
 
+        getBuffer() {
+            return ts.sys.bufferFrom?.(this.getText(0, this.index.getLength()), "utf8") as Uint8Array; // TODO: Very unsafe, discarding `undefined`
+        }
+
         getLength() {
             return this.index.getLength();
         }

src/services/services.ts

@@ -1066,7 +1066,7 @@ namespace ts {
     }
 
     export function createLanguageServiceSourceFile(fileName: string, scriptSnapshot: IScriptSnapshot, scriptTarget: ScriptTarget, version: string, setNodeParents: boolean, scriptKind?: ScriptKind): SourceFile {
-        const sourceFile = createSourceFile(fileName, getSnapshotText(scriptSnapshot), scriptTarget, setNodeParents, scriptKind);
+        const sourceFile = createSourceFile(fileName, endsWith(fileName, ".wasm") ? getSnapshotBuffer(scriptSnapshot) : getSnapshotText(scriptSnapshot), scriptTarget, setNodeParents, scriptKind);
         setSourceFileFields(sourceFile, scriptSnapshot, version);
         return sourceFile;
     }

src/services/shims.ts

@@ -312,6 +312,8 @@ namespace ts {
             return this.scriptSnapshotShim.getText(start, end);
         }
 
+        getBuffer = notImplemented;
+
         public getLength(): number {
             return this.scriptSnapshotShim.getLength();
         }

src/services/types.ts

@@ -121,6 +121,11 @@ namespace ts {
         /** Gets a portion of the script snapshot specified by [start, end). */
         getText(start: number, end: number): string;
 
+        /**
+         * Gets the udnerlying buffer for the snapshot
+         */
+        getBuffer(): Uint8Array;
+
         /** Gets the length of this script snapshot. */
         getLength(): number;
 
@@ -140,7 +145,7 @@ namespace ts {
     export namespace ScriptSnapshot {
         class StringScriptSnapshot implements IScriptSnapshot {
 
-            constructor(private text: string) {
+            constructor(private text: string, private buffer: Uint8Array | undefined) {
             }
 
             public getText(start: number, end: number): string {
@@ -149,6 +154,10 @@ namespace ts {
                     : this.text.substring(start, end);
             }
 
+            public getBuffer() {
+                return this.buffer || (ts.sys.bufferFrom?.(this.text, "utf8") as Uint8Array);
+            }
+
             public getLength(): number {
                 return this.text.length;
             }
@@ -160,8 +169,8 @@ namespace ts {
             }
         }
 
-        export function fromString(text: string): IScriptSnapshot {
-            return new StringScriptSnapshot(text);
+        export function fromString(text: string, buffer?: Uint8Array): IScriptSnapshot {
+            return new StringScriptSnapshot(text, buffer);
         }
     }
 

src/services/utilities.ts

@@ -1723,6 +1723,10 @@ namespace ts {
         return snap.getText(0, snap.getLength());
     }
 
+    export function getSnapshotBuffer(snap: IScriptSnapshot): Uint8Array {
+        return snap.getBuffer();
+    }
+
     export function repeatString(str: string, count: number): string {
         let result = "";
         for (let i = 0; i < count; i++) {