From 253cc1c7b89497801666083173d78a99e379922c Mon Sep 17 00:00:00 2001
From: Matthew Bastien <matthew_bastien@apple.com>
Date: Wed, 18 Dec 2024 11:33:09 -0500
Subject: [PATCH] clean up live preview VS Code event handlers

---
 .../DocumentationPreviewEditor.ts             | 76 ++++++++++++-------
 1 file changed, 50 insertions(+), 26 deletions(-)

diff --git a/src/documentation/DocumentationPreviewEditor.ts b/src/documentation/DocumentationPreviewEditor.ts
index 156104df..71bb964d 100644
--- a/src/documentation/DocumentationPreviewEditor.ts
+++ b/src/documentation/DocumentationPreviewEditor.ts
@@ -22,10 +22,12 @@ import { ConvertDocumentationRequest } from "../sourcekit-lsp/extensions/Convert
 export enum PreviewEditorConstant {
     VIEW_TYPE = "swift.previewDocumentationEditor",
     TITLE = "Preview Swift Documentation",
+    UNSUPPORTED_EDITOR_ERROR_MESSAGE = "The active text editor does not support Swift Documentation Live Preview",
 }
 
 export class DocumentationPreviewEditor implements vscode.Disposable {
     private readonly webviewPanel: vscode.WebviewPanel;
+    private activeTextEditor?: vscode.TextEditor;
     private subscriptions: vscode.Disposable[] = [];
 
     private disposeEmitter = new vscode.EventEmitter<void>();
@@ -36,6 +38,7 @@ export class DocumentationPreviewEditor implements vscode.Disposable {
         private readonly extension: vscode.ExtensionContext,
         private readonly context: WorkspaceContext
     ) {
+        this.activeTextEditor = vscode.window.activeTextEditor;
         const swiftDoccRenderPath = this.extension.asAbsolutePath(
             path.join("assets", "swift-docc-render")
         );
@@ -65,26 +68,20 @@ export class DocumentationPreviewEditor implements vscode.Disposable {
                 )
             )
         );
-        fs.readFile(path.join(swiftDoccRenderPath, "index.html"), "utf-8").then(
-            documentationHTML => {
-                documentationHTML = documentationHTML
-                    .replaceAll("{{BASE_PATH}}", webviewBaseURI.toString())
-                    .replace("</body>", `<script src="${scriptURI.toString()}"></script></body>`);
-                this.webviewPanel.webview.html = documentationHTML;
-                this.subscriptions.push(
-                    this.webviewPanel.webview.onDidReceiveMessage(this.receiveMessage.bind(this)),
-                    vscode.window.onDidChangeActiveTextEditor(editor => {
-                        this.convertDocumentation(editor);
-                    }),
-                    vscode.window.onDidChangeTextEditorSelection(event => {
-                        this.convertDocumentation(event.textEditor);
-                    }),
-                    this.webviewPanel.onDidDispose(this.dispose.bind(this))
-                );
-                // Reveal the editor, but don't change the focus of the active text editor
-                this.webviewPanel.reveal(undefined, true);
-            }
-        );
+        fs.readFile(path.join(swiftDoccRenderPath, "index.html"), "utf-8").then(doccRenderHTML => {
+            doccRenderHTML = doccRenderHTML
+                .replaceAll("{{BASE_PATH}}", webviewBaseURI.toString())
+                .replace("</body>", `<script src="${scriptURI.toString()}"></script></body>`);
+            this.webviewPanel.webview.html = doccRenderHTML;
+            this.subscriptions.push(
+                this.webviewPanel.webview.onDidReceiveMessage(this.receiveMessage, this),
+                vscode.window.onDidChangeActiveTextEditor(this.handleActiveTextEditorChange, this),
+                vscode.workspace.onDidChangeTextDocument(this.handleDocumentChange, this),
+                this.webviewPanel.onDidDispose(this.dispose, this)
+            );
+            // Reveal the editor, but don't change the focus of the active text editor
+            this.webviewPanel.reveal(undefined, true);
+        });
     }
 
     /** An event that is fired when the Documentation Preview Editor is disposed */
@@ -117,7 +114,10 @@ export class DocumentationPreviewEditor implements vscode.Disposable {
     private receiveMessage(message: WebviewMessage) {
         switch (message.type) {
             case "loaded":
-                this.convertDocumentation(vscode.window.activeTextEditor);
+                if (!this.activeTextEditor) {
+                    break;
+                }
+                this.convertDocumentation(this.activeTextEditor);
                 break;
             case "rendered":
                 this.renderEmitter.fire();
@@ -125,10 +125,34 @@ export class DocumentationPreviewEditor implements vscode.Disposable {
         }
     }
 
-    private async convertDocumentation(editor: vscode.TextEditor | undefined): Promise<void> {
-        const document = editor?.document;
-        if (!document || document.uri.scheme !== "file") {
-            return undefined;
+    private handleActiveTextEditorChange(activeTextEditor: vscode.TextEditor | undefined) {
+        if (this.activeTextEditor === activeTextEditor || activeTextEditor === undefined) {
+            return;
+        }
+        this.activeTextEditor = activeTextEditor;
+        this.convertDocumentation(activeTextEditor);
+    }
+
+    private handleDocumentChange(event: vscode.TextDocumentChangeEvent) {
+        if (this.activeTextEditor?.document === event.document) {
+            this.convertDocumentation(this.activeTextEditor);
+        }
+    }
+
+    private async convertDocumentation(textEditor: vscode.TextEditor): Promise<void> {
+        const document = textEditor.document;
+        if (
+            document.uri.scheme !== "file" ||
+            !["markdown", "tutorial", "swift"].includes(document.languageId)
+        ) {
+            this.postMessage({
+                type: "update-content",
+                content: {
+                    type: "error",
+                    errorMessage: PreviewEditorConstant.UNSUPPORTED_EDITOR_ERROR_MESSAGE,
+                },
+            });
+            return;
         }
 
         const response = await this.context.languageClientManager.useLanguageClient(
@@ -137,7 +161,7 @@ export class DocumentationPreviewEditor implements vscode.Disposable {
                     textDocument: {
                         uri: document.uri.toString(),
                     },
-                    position: editor.selection.start,
+                    position: textEditor.selection.start,
                 });
             }
         );