From 6363323e89a6b6f8fe817de9fabb77a4c5803b22 Mon Sep 17 00:00:00 2001
From: Rich Chiodo <rchiodo@microsoft.com>
Date: Wed, 13 Nov 2019 16:17:03 -0800
Subject: [PATCH] Fix for 7999 - CTRL+Z should not undo at the top level

---
 news/2 Fixes/7999.md                          |  1 +
 .../interactiveWindowTypes.ts                 |  1 +
 .../interactive-common/mainStateController.ts | 48 ++++++++++---------
 .../native-editor/nativeCell.tsx              | 15 ++++--
 .../native-editor/nativeEditor.tsx            | 12 +++++
 .../nativeEditor.functional.test.tsx          | 42 +++++++++++++++-
 6 files changed, 91 insertions(+), 28 deletions(-)
 create mode 100644 news/2 Fixes/7999.md

diff --git a/news/2 Fixes/7999.md b/news/2 Fixes/7999.md
new file mode 100644
index 000000000000..216f1c461491
--- /dev/null
+++ b/news/2 Fixes/7999.md	
@@ -0,0 +1 @@
+CTRL+Z is deleting cells. It should only undo changes inside of the code for a cell. 'Z' and 'SHIFT+Z' are for undoing/redoing cell adds/moves.
diff --git a/src/client/datascience/interactive-common/interactiveWindowTypes.ts b/src/client/datascience/interactive-common/interactiveWindowTypes.ts
index 5b04aa885d4d..e1489b3b71ed 100644
--- a/src/client/datascience/interactive-common/interactiveWindowTypes.ts
+++ b/src/client/datascience/interactive-common/interactiveWindowTypes.ts
@@ -97,6 +97,7 @@ export enum NativeCommandType {
     InsertBelow,
     MoveCellDown,
     MoveCellUp,
+    Redo,
     Run,
     RunAbove,
     RunAll,
diff --git a/src/datascience-ui/interactive-common/mainStateController.ts b/src/datascience-ui/interactive-common/mainStateController.ts
index f59699093fb8..c95d56d38121 100644
--- a/src/datascience-ui/interactive-common/mainStateController.ts
+++ b/src/datascience-ui/interactive-common/mainStateController.ts
@@ -267,31 +267,35 @@ export class MainStateController implements IMessageHandler {
     }
 
     public redo = () => {
-        // Pop one off of our redo stack and update our undo
-        const cells = this.pendingState.redoStack[this.pendingState.redoStack.length - 1];
-        const redoStack = this.pendingState.redoStack.slice(0, this.pendingState.redoStack.length - 1);
-        const undoStack = this.pushStack(this.pendingState.undoStack, this.pendingState.cellVMs);
-        this.sendMessage(InteractiveWindowMessages.Redo);
-        this.setState({
-            cellVMs: cells,
-            undoStack: undoStack,
-            redoStack: redoStack,
-            skipNextScroll: true
-        });
+        if (this.pendingState.redoStack.length > 0) {
+            // Pop one off of our redo stack and update our undo
+            const cells = this.pendingState.redoStack[this.pendingState.redoStack.length - 1];
+            const redoStack = this.pendingState.redoStack.slice(0, this.pendingState.redoStack.length - 1);
+            const undoStack = this.pushStack(this.pendingState.undoStack, this.pendingState.cellVMs);
+            this.sendMessage(InteractiveWindowMessages.Redo);
+            this.setState({
+                cellVMs: cells,
+                undoStack: undoStack,
+                redoStack: redoStack,
+                skipNextScroll: true
+            });
+        }
     }
 
     public undo = () => {
-        // Pop one off of our undo stack and update our redo
-        const cells = this.pendingState.undoStack[this.pendingState.undoStack.length - 1];
-        const undoStack = this.pendingState.undoStack.slice(0, this.pendingState.undoStack.length - 1);
-        const redoStack = this.pushStack(this.pendingState.redoStack, this.pendingState.cellVMs);
-        this.sendMessage(InteractiveWindowMessages.Undo);
-        this.setState({
-            cellVMs: cells,
-            undoStack: undoStack,
-            redoStack: redoStack,
-            skipNextScroll: true
-        });
+        if (this.pendingState.undoStack.length > 0) {
+            // Pop one off of our undo stack and update our redo
+            const cells = this.pendingState.undoStack[this.pendingState.undoStack.length - 1];
+            const undoStack = this.pendingState.undoStack.slice(0, this.pendingState.undoStack.length - 1);
+            const redoStack = this.pushStack(this.pendingState.redoStack, this.pendingState.cellVMs);
+            this.sendMessage(InteractiveWindowMessages.Undo);
+            this.setState({
+                cellVMs: cells,
+                undoStack: undoStack,
+                redoStack: redoStack,
+                skipNextScroll: true
+            });
+        }
     }
 
     public deleteCell = (cellId: string) => {
diff --git a/src/datascience-ui/native-editor/nativeCell.tsx b/src/datascience-ui/native-editor/nativeCell.tsx
index 08500ae84dd0..567216abeb58 100644
--- a/src/datascience-ui/native-editor/nativeCell.tsx
+++ b/src/datascience-ui/native-editor/nativeCell.tsx
@@ -325,10 +325,17 @@ export class NativeCell extends React.Component<INativeCellProps> {
                 }
                 break;
             case 'z':
-                if (!this.isFocused() && this.props.stateController.canUndo()) {
-                    e.stopPropagation();
-                    this.props.stateController.undo();
-                    this.props.stateController.sendCommand(NativeCommandType.Undo, 'keyboard');
+            case 'Z':
+                if (!this.isFocused()) {
+                    if (e.shiftKey && !e.ctrlKey && !e.altKey && this.props.stateController.canRedo()) {
+                        e.stopPropagation();
+                        this.props.stateController.redo();
+                        this.props.stateController.sendCommand(NativeCommandType.Redo, 'keyboard');
+                    } else if (!e.shiftKey && !e.altKey && !e.ctrlKey && this.props.stateController.canUndo()) {
+                        e.stopPropagation();
+                        this.props.stateController.undo();
+                        this.props.stateController.sendCommand(NativeCommandType.Undo, 'keyboard');
+                    }
                 }
                 break;
 
diff --git a/src/datascience-ui/native-editor/nativeEditor.tsx b/src/datascience-ui/native-editor/nativeEditor.tsx
index 804382da9c1d..666e53377ac8 100644
--- a/src/datascience-ui/native-editor/nativeEditor.tsx
+++ b/src/datascience-ui/native-editor/nativeEditor.tsx
@@ -365,6 +365,18 @@ export class NativeEditor extends React.Component<INativeEditorProps, IMainState
                 }
                 break;
             }
+            case 'z':
+            case 'Z':
+                if (event.shiftKey && !event.ctrlKey && !event.altKey && this.stateController.canRedo()) {
+                    event.stopPropagation();
+                    this.stateController.redo();
+                    this.stateController.sendCommand(NativeCommandType.Redo, 'keyboard');
+                } else if (!event.shiftKey && !event.altKey && !event.ctrlKey && this.stateController.canUndo()) {
+                    event.stopPropagation();
+                    this.stateController.undo();
+                    this.stateController.sendCommand(NativeCommandType.Undo, 'keyboard');
+                }
+                break;
             default:
                 break;
         }
diff --git a/src/test/datascience/nativeEditor.functional.test.tsx b/src/test/datascience/nativeEditor.functional.test.tsx
index 9723729bc294..f8458191c94f 100644
--- a/src/test/datascience/nativeEditor.functional.test.tsx
+++ b/src/test/datascience/nativeEditor.functional.test.tsx
@@ -14,7 +14,7 @@ import { Disposable, TextDocument, TextEditor, Uri, WindowState } from 'vscode';
 
 import { IApplicationShell, IDocumentManager } from '../../client/common/application/types';
 import { IFileSystem } from '../../client/common/platform/types';
-import { createDeferred } from '../../client/common/utils/async';
+import { createDeferred, sleep, waitForPromise } from '../../client/common/utils/async';
 import { createTemporaryFile } from '../../client/common/utils/fs';
 import { noop } from '../../client/common/utils/misc';
 import { Identifiers } from '../../client/datascience/constants';
@@ -939,10 +939,14 @@ for _ in range(50):
                 // Add, then undo, keep doing at least 3 times and confirm it works as expected.
                 for (let i = 0; i < 3; i += 1) {
                     // Add a new cell
-                    let update = waitForUpdate(wrapper, NativeEditor, 1);
+                    let update = waitForMessage(ioc, InteractiveWindowMessages.FocusedCellEditor);
                     simulateKeyPressOnCell(0, { code: 'a' });
                     await update;
 
+                    // Wait a bit for the time out to try and set focus a second time (this will be
+                    // fixed when we switch to redux)
+                    await sleep(100);
+
                     // There should be 4 cells and first cell is focused.
                     assert.equal(isCellSelected(wrapper, 'NativeCell', 0), false);
                     assert.equal(isCellSelected(wrapper, 'NativeCell', 1), false);
@@ -956,6 +960,40 @@ for _ in range(50):
                     await update;
                     assert.equal(isCellSelected(wrapper, 'NativeCell', 0), true);
 
+                    // Press 'ctrl+z'. This should do nothing
+                    update = waitForUpdate(wrapper, NativeEditor, 1);
+                    simulateKeyPressOnCell(0, { code: 'z', ctrlKey: true });
+                    await waitForPromise(update, 100);
+
+                    // There should be 4 cells and first cell is selected.
+                    assert.equal(isCellSelected(wrapper, 'NativeCell', 0), true);
+                    assert.equal(isCellSelected(wrapper, 'NativeCell', 1), false);
+                    assert.equal(isCellFocused(wrapper, 'NativeCell', 0), false);
+                    assert.equal(isCellFocused(wrapper, 'NativeCell', 1), false);
+                    assert.equal(wrapper.find('NativeCell').length, 4);
+
+                    // Press 'z' to undo.
+                    update = waitForUpdate(wrapper, NativeEditor, 1);
+                    simulateKeyPressOnCell(0, { code: 'z' });
+                    await update;
+
+                    // There should be 3 cells and first cell is selected & nothing focused.
+                    assert.equal(isCellSelected(wrapper, 'NativeCell', 0), true);
+                    assert.equal(isCellSelected(wrapper, 'NativeCell', 1), false);
+                    assert.equal(wrapper.find('NativeCell').length, 3);
+
+                    // Press 'shift+z' to redo
+                    update = waitForUpdate(wrapper, NativeEditor, 1);
+                    simulateKeyPressOnCell(0, { code: 'z', shiftKey: true });
+                    await update;
+
+                    // There should be 4 cells and first cell is selected.
+                    assert.equal(isCellSelected(wrapper, 'NativeCell', 0), true);
+                    assert.equal(isCellSelected(wrapper, 'NativeCell', 1), false);
+                    assert.equal(isCellFocused(wrapper, 'NativeCell', 0), false);
+                    assert.equal(isCellFocused(wrapper, 'NativeCell', 1), false);
+                    assert.equal(wrapper.find('NativeCell').length, 4);
+
                     // Press 'z' to undo.
                     update = waitForUpdate(wrapper, NativeEditor, 1);
                     simulateKeyPressOnCell(0, { code: 'z' });