Fix #5284: Restore worksheet decorations

When switching tabs, VSCode discards the text decorations that were
added by worksheet evaluation. We now keep them in memory so that they
can be added again when coming back to the worksheet.

Author: Duhemm

vscode-dotty/src/worksheet.ts

@@ -5,10 +5,10 @@ import {
 } from 'vscode'
 
 import {
-  asWorksheetRunParams, WorksheetRunRequest, WorksheetRunParams, WorksheetRunResult,
+  asWorksheetRunParams, WorksheetRunRequest, WorksheetRunResult,
   WorksheetPublishOutputParams, WorksheetPublishOutputNotification
 } from './protocol'
-import { BaseLanguageClient, DocumentSelector } from 'vscode-languageclient'
+import { BaseLanguageClient } from 'vscode-languageclient'
 import { Disposable } from 'vscode-jsonrpc'
 
 /**
@@ -23,11 +23,26 @@ export const worksheetRunKey = "dotty.worksheet.run"
  */
 export const worksheetCancelKey = "dotty.worksheet.cancel"
 
+/**
+ * A wrapper around the information that VSCode needs to display text decorations.
+ *
+ * @param decorationType    The styling options of this decoration
+ * @param decorationOptions The options of this decoraiton.
+ */
+class Decoration {
+  constructor(readonly decorationType: vscode.TextEditorDecorationType,
+              readonly decorationOptions: vscode.DecorationOptions) {
+  }
+}
+
 /** A worksheet managed by vscode */
 class Worksheet implements Disposable {
 
+  /** The version of this document the last time it was run */
+  private runVersion: number = -1
+
   /** All decorations that have been added so far */
-  private decorationTypes: vscode.TextEditorDecorationType[] = []
+  private decorations: Decoration[] = []
 
   /** The number of blank lines that have been inserted to fit the output so far. */
   private insertedLines: number = 0
@@ -62,9 +77,11 @@ class Worksheet implements Disposable {
 
   /** Remove all decorations, and resets this worksheet. */
   private reset(): void {
-    this.decorationTypes.forEach(decoration => decoration.dispose())
+    this.decorations.forEach(decoration => decoration.decorationType.dispose())
+    this.decorations = []
     this.insertedLines = 0
     this.decoratedLines.clear()
+    this.runVersion = -1
     this.margin = this.longestLine() + 5
   }
 
@@ -112,6 +129,7 @@ class Worksheet implements Disposable {
       const edit = new vscode.WorkspaceEdit()
       edit.set(this.document.uri, textEdits)
       vscode.workspace.applyEdit(edit).then(editSucceeded => {
+        this.runVersion = this.document.version
         if (editSucceeded && !token.isCancellationRequested)
           resolve(vscode.window.withProgress({
             location: ProgressLocation.Window,
@@ -135,7 +153,8 @@ class Worksheet implements Disposable {
   }
 
   /**
-   * Parse and display the result of running part of this worksheet.
+   * Parse and display the result of running part of this worksheet. The result is saved so that it
+   * can be restored if this buffer is closed.
    *
    * @param lineNumber The number of the line in the source that produced the result.
    * @param runResult  The result itself.
@@ -144,7 +163,7 @@ class Worksheet implements Disposable {
    * @return A `Thenable` that will insert necessary lines to fit the output
    *         and display the decorations upon completion.
    */
-  public displayResult(lineNumber: number, runResult: string, editor: vscode.TextEditor) {
+  public displayAndSaveResult(lineNumber: number, runResult: string, editor: vscode.TextEditor) {
     const resultLines = runResult.trim().split(/\r\n|\r|\n/g)
 
     // The line where the next decoration should be put.
@@ -160,23 +179,41 @@ class Worksheet implements Disposable {
       const editPos = new vscode.Position(actualLine + 1, 0) // add after the line
       addNewLinesEdit.insert(editor.document.uri, editPos, "\n".repeat(linesToInsert))
       this.insertedLines += linesToInsert
+      // Increase the `runVersion`, because the text edit will increase the document's version
+      this.runVersion += 1
     }
 
     return vscode.workspace.applyEdit(addNewLinesEdit).then(_ => {
       for (let line of resultLines) {
         const decorationPosition = new vscode.Position(actualLine, 0)
         const decorationMargin = this.margin - editor.document.lineAt(actualLine).text.length
         const decorationType = this.createDecoration(decorationMargin, line)
-        this.decorationTypes.push(decorationType)
+        const decorationOptions = { range: new vscode.Range(decorationPosition, decorationPosition), hoverMessage: line }
+        const decoration = new Decoration(decorationType, decorationOptions)
+
         this.decoratedLines.add(actualLine)
+        this.decorations.push(decoration)
 
-        const decoration = { range: new vscode.Range(decorationPosition, decorationPosition), hoverMessage: line }
-        editor.setDecorations(decorationType, [decoration])
+        editor.setDecorations(decorationType, [decorationOptions])
         actualLine += 1
       }
     })
   }
 
+  /**
+   * Restore the decorations that belong to this worksheet in `editor`. If the document has been
+   * changed, since the last run of the worksheet, the decorations won't be added.
+   *
+   * @param editor The editor where to display the decorations.
+   */
+  public restoreDecorations(editor: vscode.TextEditor) {
+    if (editor.document.version == this.runVersion) {
+      this.decorations.forEach(decoration => {
+        editor.setDecorations(decoration.decorationType, [decoration.decorationOptions])
+      })
+    }
+  }
+
   /**
    * Create a new `TextEditorDecorationType` showing `text`. The decoration
    * will appear `margin` characters after the end of the line.
@@ -307,6 +344,14 @@ export class WorksheetProvider implements Disposable {
           this.worksheets.delete(document)
         }
       }),
+      vscode.window.onDidChangeActiveTextEditor(editor => {
+        if (editor) {
+          const worksheet = this.worksheetFor(editor.document)
+          if (worksheet) {
+            worksheet.restoreDecorations(editor)
+          }
+        }
+      }),
       vscode.commands.registerCommand(worksheetRunKey, () => {
         this.callOnActiveWorksheet(w => w.run())
       }),
@@ -375,7 +420,7 @@ export class WorksheetProvider implements Disposable {
     if (editor) {
       const worksheet = this.worksheetFor(editor.document)
       if (worksheet) {
-        worksheet.displayResult(output.line - 1, output.content, editor)
+        worksheet.displayAndSaveResult(output.line - 1, output.content, editor)
       }
     }
   }