Fix highlights, jump to and hover in foreign documents

Prior to this commit the positions would be always resolved with respect
to the root document because virtual document was being taken from
adapter using `virtualPositionToRootPosition` func which is error
prone as the assumption of virtual document == root document is not
obvious; this function was removed. To make the fixes work, a
workaround for foreign document created in inherited classes always having
the usptream class was added; the issue was reported in
jupyterlab/jupyterlab#15126

Author: krassowski

packages/jupyterlab-lsp/src/converter.ts

@@ -10,7 +10,7 @@ import type {
 } from '@jupyterlab/lsp';
 import type * as lsProtocol from 'vscode-languageserver-protocol';
 
-import type { VirtualDocument } from './virtual/document';
+import { VirtualDocument } from './virtual/document';
 
 export class PositionConverter {
   static lsp_to_cm(position: lsProtocol.Position): IPosition {
@@ -86,18 +86,6 @@ export function rootPositionToVirtualPosition(
   return adapter.virtualDocument!.virtualPositionAtDocument(rootAsSource);
 }
 
-export function virtualPositionToRootPosition(
-  adapter: WidgetLSPAdapter<any>,
-  position: IVirtualPosition
-): IRootPosition | null {
-  if (!adapter.virtualDocument) {
-    throw Error('Virtual document of adapter disposed!');
-  }
-  return (adapter.virtualDocument as VirtualDocument).transformVirtualToRoot(
-    position
-  );
-}
-
 export function rootPositionToEditorPosition(
   adapter: WidgetLSPAdapter<any>,
   position: IRootPosition
@@ -126,12 +114,20 @@ export function editorPositionToRootPosition(
 export function rangeToEditorRange(
   adapter: WidgetLSPAdapter<any>,
   range: lsProtocol.Range,
-  editor: CodeEditor.IEditor | null
+  editor: CodeEditor.IEditor | null,
+  virtualDocument: VirtualDocument
 ): IEditorRange {
   let start = PositionConverter.lsp_to_cm(range.start) as IVirtualPosition;
   let end = PositionConverter.lsp_to_cm(range.end) as IVirtualPosition;
 
-  let startInRoot = virtualPositionToRootPosition(adapter, start);
+  // because `openForeign()` does not use new this.constructor, we need to workaround it for now:
+  // const startInRoot = virtualDocument.transformVirtualToRoot(start);
+  // https://github.com/jupyterlab/jupyterlab/issues/15126
+  const startInRoot = VirtualDocument.prototype.transformVirtualToRoot.call(
+    virtualDocument,
+    start
+  );
+
   if (!startInRoot) {
     throw Error('Could not determine position in root');
   }

packages/jupyterlab-lsp/src/features/highlights.ts

@@ -14,8 +14,7 @@ import {
   ILSPFeatureManager,
   IEditorPosition,
   ILSPDocumentConnectionManager,
-  WidgetLSPAdapter,
-  VirtualDocument
+  WidgetLSPAdapter
 } from '@jupyterlab/lsp';
 import { ISettingRegistry } from '@jupyterlab/settingregistry';
 import { LabIcon } from '@jupyterlab/ui-components';
@@ -37,6 +36,7 @@ import { DocumentHighlightKind } from '../lsp';
 import { createMarkManager, ISimpleMarkManager } from '../marks';
 import { PLUGIN_ID } from '../tokens';
 import { BrowserConsole } from '../virtual/console';
+import { VirtualDocument } from '../virtual/document';
 
 export const highlightIcon = new LabIcon({
   name: 'lsp:highlight',
@@ -160,7 +160,8 @@ export class HighlightsFeature extends Feature {
 
   protected handleHighlight(
     items: lsProtocol.DocumentHighlight[] | null,
-    adapter: WidgetLSPAdapter<any>
+    adapter: WidgetLSPAdapter<any>,
+    document: VirtualDocument
   ) {
     this.markManager.clearAllMarks();
 
@@ -174,7 +175,7 @@ export class HighlightsFeature extends Feature {
     >();
 
     for (let item of items) {
-      let range = rangeToEditorRange(adapter, item.range, null);
+      let range = rangeToEditorRange(adapter, item.range, null, document);
       const editor = range.editor;
 
       let optionsList = highlightsByEditor.get(editor);
@@ -350,7 +351,7 @@ export class HighlightsFeature extends Feature {
         return;
       }
 
-      this.handleHighlight(highlights, adapter);
+      this.handleHighlight(highlights, adapter, document);
       this._lastToken = token;
     } catch (e) {
       this.console.warn('Could not get highlights:', e);

packages/jupyterlab-lsp/src/features/hover.ts

@@ -18,7 +18,6 @@ import {
   isEqual,
   ILSPDocumentConnectionManager,
   WidgetLSPAdapter,
-  VirtualDocument,
   Document
 } from '@jupyterlab/lsp';
 import { IRenderMimeRegistry } from '@jupyterlab/rendermime';
@@ -45,6 +44,7 @@ import { createMarkManager, ISimpleMarkManager } from '../marks';
 import { PLUGIN_ID } from '../tokens';
 import { getModifierState } from '../utils';
 import { BrowserConsole } from '../virtual/console';
+import { VirtualDocument } from '../virtual/document';
 
 export const hoverIcon = new LabIcon({
   name: 'lsp:hover',
@@ -365,7 +365,8 @@ export class HoverFeature extends Feature {
       context.adapter,
       response,
       context.token,
-      context.editor
+      context.editor,
+      virtualDocument
     );
     return this._addRange(
       context.adapter,
@@ -619,7 +620,8 @@ export class HoverFeature extends Feature {
             adapter,
             response,
             token,
-            editor
+            editor,
+            document
           );
           responseData = {
             response: response,
@@ -689,10 +691,11 @@ export class HoverFeature extends Feature {
     adapter: WidgetLSPAdapter<any>,
     response: lsProtocol.Hover,
     token: CodeEditor.IToken,
-    editor: CodeEditor.IEditor
+    editor: CodeEditor.IEditor,
+    document: VirtualDocument
   ): IEditorRange {
     if (typeof response.range !== 'undefined') {
-      return rangeToEditorRange(adapter, response.range, editor);
+      return rangeToEditorRange(adapter, response.range, editor, document);
     }
 
     const startInEditor = editor.getPositionAt(token.offset);

packages/jupyterlab-lsp/src/features/jump_to.ts

@@ -48,13 +48,13 @@ import {
   documentAtRootPosition,
   editorAtRootPosition,
   rootPositionToVirtualPosition,
-  rootPositionToEditorPosition,
-  virtualPositionToRootPosition
+  rootPositionToEditorPosition
 } from '../converter';
 import { FeatureSettings, Feature } from '../feature';
 import { PLUGIN_ID } from '../tokens';
 import { getModifierState, uriToContentsPath, urisEqual } from '../utils';
 import { BrowserConsole } from '../virtual/console';
+import { VirtualDocument } from '../virtual/document';
 
 export const jumpToIcon = new LabIcon({
   name: 'lsp:jump-to',
@@ -245,7 +245,7 @@ export class NavigationFeature extends Feature {
     connection.clientRequests['textDocument/definition']
       .request(positionParams)
       .then(targets => {
-        this.handleJump(targets, positionParams, adapter)
+        this.handleJump(targets, positionParams, adapter, document)
           .then((result: JumpResult | undefined) => {
             if (
               result === JumpResult.NoTargetsFound ||
@@ -259,7 +259,7 @@ export class NavigationFeature extends Feature {
                 })
                 .then(targets =>
                   // TODO: explain that we are now presenting references?
-                  this.handleJump(targets, positionParams, adapter)
+                  this.handleJump(targets, positionParams, adapter, document)
                 )
                 .catch(this.console.warn);
             }
@@ -361,7 +361,8 @@ export class NavigationFeature extends Feature {
   async handleJump(
     locationData: AnyLocation,
     positionParams: lsp.TextDocumentPositionParams,
-    adapter: WidgetLSPAdapter<any>
+    adapter: WidgetLSPAdapter<any>,
+    document: VirtualDocument
   ) {
     const locations = this._harmonizeLocations(locationData);
     const targetInfo = await this._chooseTarget(locations);
@@ -382,10 +383,16 @@ export class NavigationFeature extends Feature {
 
     if (urisEqual(uri, positionParams.textDocument.uri)) {
       // if in current file, transform from the position within virtual document to the editor position:
-      const rootPosition = virtualPositionToRootPosition(
-        adapter,
-        virtualPosition
-      );
+
+      // because `openForeign()` does not use new this.constructor, we need to workaround it for now:
+      // const rootPosition = document.transformVirtualToRoot(virtualPosition);
+      // https://github.com/jupyterlab/jupyterlab/issues/15126
+      const rootPosition =
+        VirtualDocument.prototype.transformVirtualToRoot.call(
+          document,
+          virtualPosition
+        );
+
       if (rootPosition === null) {
         this.console.warn(
           'Could not jump: conversion from virtual position to editor position failed',
@@ -581,7 +588,7 @@ export const JUMP_PLUGIN: JupyterFrontEndPlugin<void> = {
         const targets = await connection.clientRequests[
           'textDocument/definition'
         ].request(positionParams);
-        await feature.handleJump(targets, positionParams, adapter);
+        await feature.handleJump(targets, positionParams, adapter, document);
       },
       label: trans.__('Jump to definition'),
       icon: jumpToIcon,
@@ -627,7 +634,7 @@ export const JUMP_PLUGIN: JupyterFrontEndPlugin<void> = {
           ...positionParams,
           context: { includeDeclaration: false }
         });
-        await feature.handleJump(targets, positionParams, adapter);
+        await feature.handleJump(targets, positionParams, adapter, document);
       },
       label: trans.__('Jump to references'),
       icon: jumpToIcon,

packages/jupyterlab-lsp/src/virtual/document.ts

@@ -184,11 +184,10 @@ export class VirtualDocument extends VirtualDocumentBase {
    * @experimental
    */
   transformVirtualToRoot(position: IVirtualPosition): IRootPosition | null {
-    // a method which was previously implemented in CodeMirrorIntegration
-    // but probably should have been in VirtualDocument all along
     let editor = this.virtualLines.get(position.line)!.editor;
     let editorPosition = this.transformVirtualToEditor(position);
-    return this.transformFromEditorToRoot(editor, editorPosition!);
+    // only root holds the full editor mapping
+    return this.root.transformFromEditorToRoot(editor, editorPosition!);
   }
 
   /**