Fix undo/redo in Jupyter notebooks

Two things going on here:

- Notebook cells have the same file name, but different Documents/Uris
  (fragment is different). We now use Uri (rather than file name) as the
  ModeHandler key, so cells have separate VimStates and therefore undo
  histories. VS Code lets you customize that behavior, but we don't
  currently respect that setting.

- TextDocuments associated with notebook cells don't ever seem to have
  `isDirty` set (not sure if this is a bug/limitation or the intended
  design). This caused a "file changed on disk" handler to be triggered
  on every change, which wipes that file's undo history. I've added
  notebook cells as a special case disabling that, as a stopgap.

Fixes #6960

extensionBase.ts

@@ -2,7 +2,6 @@ import * as vscode from 'vscode';
 import * as path from 'path';
 
 import { CompositionState } from './src/state/compositionState';
-import { EditorIdentity } from './src/editorIdentity';
 import { Globals } from './src/globals';
 import { Jump } from './src/jumps/jump';
 import { ModeHandler } from './src/mode/modeHandler';
@@ -21,7 +20,7 @@ import { SpecialKeys } from './src/util/specialKeys';
 import { HistoryTracker } from './src/history/historyTracker';
 
 let extensionContext: vscode.ExtensionContext;
-let previousActiveEditorId: EditorIdentity | undefined;
+let previousActiveEditorUri: vscode.Uri | undefined;
 let lastClosedModeHandler: ModeHandler | null = null;
 
 interface ICodeKeybinding {
@@ -37,34 +36,30 @@ export async function getAndUpdateModeHandler(
     return undefined;
   }
 
-  const activeEditorId = EditorIdentity.fromEditor(activeTextEditor);
+  const uri = activeTextEditor.document.uri;
 
-  const [curHandler, isNew] = await ModeHandlerMap.getOrCreate(activeEditorId);
+  const [curHandler, isNew] = await ModeHandlerMap.getOrCreate(uri);
   if (isNew) {
     extensionContext.subscriptions.push(curHandler);
   }
 
   curHandler.vimState.editor = activeTextEditor;
 
-  if (
-    forceSyncAndUpdate ||
-    !previousActiveEditorId ||
-    !previousActiveEditorId.isEqual(activeEditorId)
-  ) {
+  if (forceSyncAndUpdate || !previousActiveEditorUri || previousActiveEditorUri !== uri) {
     // We sync the cursors here because ModeHandler is specific to a document, not an editor, so we
     // need to update our representation of the cursors when switching between editors for the same document.
     // This will be unnecessary once #4889 is fixed.
     curHandler.syncCursors();
     await curHandler.updateView({ drawSelection: false, revealRange: false });
   }
 
-  previousActiveEditorId = activeEditorId;
+  previousActiveEditorUri = uri;
 
   if (curHandler.focusChanged) {
     curHandler.focusChanged = false;
 
-    if (previousActiveEditorId) {
-      const prevHandler = ModeHandlerMap.get(previousActiveEditorId);
+    if (previousActiveEditorUri) {
+      const prevHandler = ModeHandlerMap.get(previousActiveEditorUri);
       prevHandler!.focusChanged = true;
     }
   }
@@ -168,14 +163,19 @@ export async function activate(context: vscode.ExtensionContext, handleLocal: bo
     };
 
     ModeHandlerMap.getAll()
-      .filter((modeHandler) => modeHandler.vimState.identity.fileName === event.document.fileName)
+      .filter((modeHandler) => modeHandler.vimState.documentUri === event.document.uri)
       .forEach((modeHandler) => {
         contentChangeHandler(modeHandler);
       });
 
     if (handleLocal) {
       setTimeout(() => {
-        if (!event.document.isDirty && !event.document.isUntitled && event.contentChanges.length) {
+        if (
+          !event.document.isDirty &&
+          !event.document.isUntitled &&
+          event.document.uri.scheme !== 'vscode-notebook-cell' && // TODO: Notebooks never seem to be marked dirty...
+          event.contentChanges.length
+        ) {
           handleContentChangedFromDisk(event.document);
         }
       }, 0);
@@ -189,8 +189,8 @@ export async function activate(context: vscode.ExtensionContext, handleLocal: bo
       const documents = vscode.workspace.textDocuments;
 
       // Delete modehandler once all tabs of this document have been closed
-      for (const editorIdentity of ModeHandlerMap.getKeys()) {
-        const modeHandler = ModeHandlerMap.get(editorIdentity);
+      for (const uri of ModeHandlerMap.keys()) {
+        const modeHandler = ModeHandlerMap.get(uri);
 
         let shouldDelete = false;
         if (modeHandler == null || modeHandler.vimState.editor === undefined) {
@@ -206,7 +206,7 @@ export async function activate(context: vscode.ExtensionContext, handleLocal: bo
         }
 
         if (shouldDelete) {
-          ModeHandlerMap.delete(editorIdentity);
+          ModeHandlerMap.delete(uri);
         }
       }
     },
@@ -228,8 +228,8 @@ export async function activate(context: vscode.ExtensionContext, handleLocal: bo
     context,
     vscode.window.onDidChangeActiveTextEditor,
     async () => {
-      const mhPrevious: ModeHandler | undefined = previousActiveEditorId
-        ? ModeHandlerMap.get(previousActiveEditorId)
+      const mhPrevious: ModeHandler | undefined = previousActiveEditorUri
+        ? ModeHandlerMap.get(previousActiveEditorUri)
         : undefined;
       // Track the closed editor so we can use it the next time an open event occurs.
       // When vscode changes away from a temporary file, onDidChangeActiveTextEditor first twice.
@@ -284,7 +284,7 @@ export async function activate(context: vscode.ExtensionContext, handleLocal: bo
         return;
       }
 
-      const mh = ModeHandlerMap.get(EditorIdentity.fromEditor(vscode.window.activeTextEditor));
+      const mh = ModeHandlerMap.get(vscode.window.activeTextEditor.document.uri);
       if (mh === undefined) {
         // We don't care if there is no active editor
         return;
@@ -649,7 +649,7 @@ async function forceStopRecursiveRemap(mh: ModeHandler): Promise<boolean> {
 
 function handleContentChangedFromDisk(document: vscode.TextDocument): void {
   ModeHandlerMap.getAll()
-    .filter((modeHandler) => modeHandler.vimState.identity.fileName === document.fileName)
+    .filter((modeHandler) => modeHandler.vimState.documentUri === document.uri)
     .forEach((modeHandler) => {
       modeHandler.vimState.historyTracker = new HistoryTracker(modeHandler.vimState);
     });

src/mode/modeHandler.ts

@@ -24,7 +24,6 @@ import { TextEditor } from './../textEditor';
 import { VimError, ForceStopRemappingError } from './../error';
 import { VimState } from './../state/vimState';
 import { VSCodeContext } from '../util/vscodeContext';
-import { ExCommandLine } from '../cmd_line/commandLine';
 import { configuration } from '../configuration/configuration';
 import { decoration } from '../configuration/decoration';
 import { scrollView } from '../util/util';
@@ -38,17 +37,16 @@ import { isTextTransformation } from '../transformations/transformations';
 import { executeTransformations, IModeHandler } from '../transformations/execute';
 import { globalState } from '../state/globalState';
 import { Notation } from '../configuration/notation';
-import { EditorIdentity } from '../editorIdentity';
 import { SpecialKeys } from '../util/specialKeys';
 import { BaseOperator } from '../actions/operator';
 import { SearchByNCharCommand } from '../actions/plugins/easymotion/easymotion.cmd';
-import { Position } from 'vscode';
+import { Position, Uri } from 'vscode';
 import { RemapState } from '../state/remapState';
 import * as process from 'process';
 import { EasyMotion } from '../actions/plugins/easymotion/easymotion';
 
 interface IModeHandlerMap {
-  get(editorId: EditorIdentity): ModeHandler | undefined;
+  get(editorId: Uri): ModeHandler | undefined;
 }
 
 /**
@@ -1533,9 +1531,7 @@ export class ModeHandler implements vscode.Disposable, IModeHandler {
       (configuration.incsearch && this.currentMode === Mode.SearchInProgressMode) ||
       (configuration.hlsearch && globalState.hl);
     for (const editor of vscode.window.visibleTextEditors) {
-      this.handlerMap
-        .get(EditorIdentity.fromEditor(editor))
-        ?.updateSearchHighlights(showHighlights);
+      this.handlerMap.get(editor.document.uri)?.updateSearchHighlights(showHighlights);
     }
 
     const easyMotionDimRanges =

src/mode/modeHandlerMap.ts

@@ -1,13 +1,13 @@
+import { Uri } from 'vscode';
 import { ModeHandler } from './modeHandler';
-import { EditorIdentity } from '../editorIdentity';
 
 /**
- * Stores one ModeHandler (and therefore VimState) per editor.
+ * Stores one ModeHandler (and therefore VimState) per TextDocument.
  */
 class ModeHandlerMapImpl {
-  private modeHandlerMap = new Map<EditorIdentity, ModeHandler>();
+  private modeHandlerMap = new Map<Uri, ModeHandler>();
 
-  public async getOrCreate(editorId: EditorIdentity): Promise<[ModeHandler, boolean]> {
+  public async getOrCreate(editorId: Uri): Promise<[ModeHandler, boolean]> {
     let isNew = false;
     let modeHandler: ModeHandler | undefined = this.get(editorId);
 
@@ -19,24 +19,19 @@ class ModeHandlerMapImpl {
     return [modeHandler, isNew];
   }
 
-  public get(editorId: EditorIdentity): ModeHandler | undefined {
-    for (const [key, value] of this.modeHandlerMap.entries()) {
-      if (key.isEqual(editorId)) {
-        return value;
-      }
-    }
-    return undefined;
+  public get(uri: Uri): ModeHandler | undefined {
+    return this.modeHandlerMap.get(uri);
   }
 
-  public getKeys(): EditorIdentity[] {
-    return [...this.modeHandlerMap.keys()];
+  public keys(): IterableIterator<Uri> {
+    return this.modeHandlerMap.keys();
   }
 
   public getAll(): ModeHandler[] {
     return [...this.modeHandlerMap.values()];
   }
 
-  public delete(editorId: EditorIdentity) {
+  public delete(editorId: Uri) {
     const modeHandler = this.modeHandlerMap.get(editorId);
     if (modeHandler) {
       modeHandler.dispose();

src/state/vimState.ts

@@ -3,7 +3,6 @@ import * as vscode from 'vscode';
 import { IMovement } from '../actions/baseMotion';
 import { configuration } from '../configuration/configuration';
 import { IEasyMotion } from '../actions/plugins/easymotion/types';
-import { EditorIdentity } from './../editorIdentity';
 import { HistoryTracker } from './../history/historyTracker';
 import { Logger } from '../util/logger';
 import { Mode } from '../mode/mode';
@@ -61,7 +60,7 @@ export class VimState implements vscode.Disposable {
 
   public easyMotion: IEasyMotion;
 
-  public readonly identity: EditorIdentity;
+  public readonly documentUri: vscode.Uri;
 
   public editor: vscode.TextEditor;
 
@@ -307,7 +306,7 @@ export class VimState implements vscode.Disposable {
 
   public constructor(editor: vscode.TextEditor, easyMotion: IEasyMotion) {
     this.editor = editor;
-    this.identity = EditorIdentity.fromEditor(editor);
+    this.documentUri = editor?.document.uri ?? vscode.Uri.file(''); // TODO: this is needed for some badly written tests
     this.historyTracker = new HistoryTracker(this);
     this.easyMotion = easyMotion;
   }

test/mode/modeHandlerMap.test.ts

@@ -1,13 +1,13 @@
 import * as assert from 'assert';
+import { Uri } from 'vscode';
 
 import { ModeHandlerMap } from '../../src/mode/modeHandlerMap';
-import { EditorIdentity } from '../../src/editorIdentity';
 
-function createRandomEditorIdentity(): EditorIdentity {
-  return new EditorIdentity(Math.random().toString(36).substring(7));
+function createRandomUri(): Uri {
+  return Uri.file(Math.random().toString(36).substring(7));
 }
 
-suite('Mode Handler Map', () => {
+suite.only('Mode Handler Map', () => {
   setup(() => {
     ModeHandlerMap.clear();
   });
@@ -17,22 +17,18 @@ suite('Mode Handler Map', () => {
   });
 
   test('getOrCreate', async () => {
-    const key = createRandomEditorIdentity();
+    const key = createRandomUri();
     let [modeHandler, isNew] = await ModeHandlerMap.getOrCreate(key);
     assert.strictEqual(isNew, true);
     assert.notStrictEqual(modeHandler, undefined);
 
     [, isNew] = await ModeHandlerMap.getOrCreate(key);
     assert.strictEqual(isNew, false);
 
-    // getKeys
-    const keys = ModeHandlerMap.getKeys();
-    assert.strictEqual(keys.length, 1, 'getKeys().length');
-    assert.strictEqual(keys[0], key, 'key');
+    assert.deepStrictEqual([...ModeHandlerMap.keys()], [key], 'keys');
 
     // getAll
-    const modeHandlerList = ModeHandlerMap.getAll();
-    assert.strictEqual(modeHandlerList.length, 1, 'getAll() should have only returned one');
+    assert.strictEqual(ModeHandlerMap.getAll().length, 1, 'getAll() should have only returned one');
 
     // delete
     ModeHandlerMap.delete(key);
@@ -41,14 +37,14 @@ suite('Mode Handler Map', () => {
 
   test('get', async () => {
     // same file name should return the same modehandler
-    const identity = createRandomEditorIdentity();
+    const identity = createRandomUri();
 
     let [modeHandler, isNew] = await ModeHandlerMap.getOrCreate(identity);
     assert.strictEqual(isNew, true);
     assert.notStrictEqual(modeHandler, undefined);
 
     const prevModeHandler = modeHandler;
-    [modeHandler, isNew] = await ModeHandlerMap.getOrCreate(new EditorIdentity(identity.fileName));
+    [modeHandler, isNew] = await ModeHandlerMap.getOrCreate(identity);
     assert.strictEqual(isNew, false);
     assert.strictEqual(prevModeHandler, modeHandler);
   });

test/testSimplifier.ts

@@ -13,7 +13,6 @@ import { vimrcKeyRemappingBuilder } from '../src/configuration/vimrcKeyRemapping
 import { IConfiguration } from '../src/configuration/iconfiguration';
 import { Position } from 'vscode';
 import { ModeHandlerMap } from '../src/mode/modeHandlerMap';
-import { EditorIdentity } from '../src/editorIdentity';
 import { StatusBar } from '../src/statusBar';
 import { Register } from '../src/register/register';
 import { ModeHandler } from '../src/mode/modeHandler';
@@ -347,7 +346,7 @@ async function testIt(testObj: ITestObject): Promise<ModeHandler> {
 
   // Generate a brand new ModeHandler for this editor
   ModeHandlerMap.clear();
-  const [modeHandler, _] = await ModeHandlerMap.getOrCreate(EditorIdentity.fromEditor(editor));
+  const [modeHandler, _] = await ModeHandlerMap.getOrCreate(editor.document.uri);
 
   let keysPressed = testObj.keysPressed;
   if (process.platform === 'win32') {
@@ -439,7 +438,7 @@ async function testItWithRemaps(testObj: ITestWithRemapsObject): Promise<ModeHan
 
   // Generate a brand new ModeHandler for this editor
   ModeHandlerMap.clear();
-  const [modeHandler, _] = await ModeHandlerMap.getOrCreate(EditorIdentity.fromEditor(editor));
+  const [modeHandler, _] = await ModeHandlerMap.getOrCreate(editor.document.uri);
 
   // Change remappings
   if (testObj.remaps) {