Fix gj/gk so it maintains cursor position (VSCodeVim#3890)

Fix gj & gk motion commands so that they maintain column position when jumping across short lines.
Fixes VSCodeVim#1323 .

The previous fix (commit e14d23d) was somewhat naive as it conflated
what should have been entirely separate classes. One set for moving the
cursor up and down, and the other for moving it up and down while
maintaining cursors position in wrapped lines. This conflation would
probably be fine except that the former set of classes was in turn being
used by the MoveUpFoldFix/MoveDownFoldFix classes. Digging through the
commit history, the original confusion might have been introduced by
commit 5a9f2f0.

We now correctly separate Move(Up|Down)ByScreenLine classes from the
Move(Up|Down)ByScreenLineMaintainDesiredColumn classes.

src/actions/motion.ts

@@ -86,12 +86,42 @@ abstract class MoveByScreenLine extends BaseMovement {
   }
 }
 
+export class MoveUpByScreenLine extends MoveByScreenLine {
+  movementType: CursorMovePosition = 'up';
+  by: CursorMoveByUnit = 'wrappedLine';
+  value = 1;
+}
+
+class MoveDownByScreenLine extends MoveByScreenLine {
+  movementType: CursorMovePosition = 'down';
+  by: CursorMoveByUnit = 'wrappedLine';
+  value = 1;
+}
+
 abstract class MoveByScreenLineMaintainDesiredColumn extends MoveByScreenLine {
   doesntChangeDesiredColumn = true;
   public async execAction(position: Position, vimState: VimState): Promise<Position | IMovement> {
     let prevDesiredColumn = vimState.desiredColumn;
     let prevLine = vimState.editor.selection.active.line;
 
+    if (vimState.currentMode !== ModeName.Normal) {
+      /**
+       * As VIM and VSCode handle the end of selection index a little
+       * differently we need to sometimes move the cursor at the end
+       * of the selection back by a character.
+       */
+      let start = Position.FromVSCodePosition(vimState.editor.selection.start);
+      if ((this.movementType === 'down' && position.line > start.line) ||
+          (this.movementType === 'up'   && position.line < prevLine)) {
+          await vscode.commands.executeCommand('cursorMove', {
+            to: 'left',
+            select: true,
+            by: 'character',
+            value: 1,
+          });
+      }
+    }
+
     await vscode.commands.executeCommand('cursorMove', {
       to: this.movementType,
       select: vimState.currentMode !== ModeName.Normal,
@@ -115,23 +145,22 @@ abstract class MoveByScreenLineMaintainDesiredColumn extends MoveByScreenLine {
       let curPos = Position.FromVSCodePosition(vimState.editor.selection.active);
 
       // We want to swap the cursor start stop positions based on which direction we are moving, up or down
-      if (start.isEqual(curPos)) {
-        position = start;
+      if (start.isEqual(curPos) && !start.isEqual(stop)) {
         [start, stop] = [stop, start];
-        start = start.getLeft();
+        if (prevLine !== start.line) {
+          start = start.getLeft();
+        }
+      }
+
+      if (position.line !== stop.line) {
+        stop = stop.withColumn(prevDesiredColumn);
       }
 
       return { start, stop };
     }
   }
 }
 
-class MoveDownByScreenLineMaintainDesiredColumn extends MoveByScreenLineMaintainDesiredColumn {
-  movementType: CursorMovePosition = 'down';
-  by: CursorMoveByUnit = 'wrappedLine';
-  value = 1;
-}
-
 class MoveDownFoldFix extends MoveByScreenLineMaintainDesiredColumn {
   movementType: CursorMovePosition = 'down';
   by: CursorMoveByUnit = 'line';
@@ -191,12 +220,6 @@ class MoveDownArrow extends MoveDown {
   keys = ['<down>'];
 }
 
-class MoveUpByScreenLineMaintainDesiredColumn extends MoveByScreenLineMaintainDesiredColumn {
-  movementType: CursorMovePosition = 'up';
-  by: CursorMoveByUnit = 'wrappedLine';
-  value = 1;
-}
-
 @RegisterAction
 class MoveUp extends BaseMovement {
   keys = ['k'];
@@ -760,7 +783,7 @@ class MoveScreenLineCenter extends MoveByScreenLine {
 }
 
 @RegisterAction
-export class MoveUpByScreenLine extends MoveByScreenLine {
+export class MoveUpByScreenLineMaintainDesiredColumn extends MoveByScreenLineMaintainDesiredColumn {
   modes = [ModeName.Insert, ModeName.Normal, ModeName.Visual];
   keys = [['g', 'k'], ['g', '<up>']];
   movementType: CursorMovePosition = 'up';
@@ -769,7 +792,7 @@ export class MoveUpByScreenLine extends MoveByScreenLine {
 }
 
 @RegisterAction
-class MoveDownByScreenLine extends MoveByScreenLine {
+class MoveDownByScreenLineMaintainDesiredColumn extends MoveByScreenLineMaintainDesiredColumn {
   modes = [ModeName.Insert, ModeName.Normal, ModeName.Visual];
   keys = [['g', 'j'], ['g', '<down>']];
   movementType: CursorMovePosition = 'down';

test/mode/modeVisual.test.ts

@@ -226,6 +226,36 @@ suite('Mode Visual', () => {
     });
   });
 
+  suite('Screen line motions in Visual Mode', () => {
+    newTest({
+      title: "Can handle 'gk'",
+      start: ['blah', 'duh', '|dur', 'hur'],
+      keysPressed: 'vgkx',
+      end: ['blah', '|ur', 'hur'],
+    });
+
+    newTest({
+      title: "Can handle 'gj'",
+      start: ['blah', 'duh', '|dur', 'hur'],
+      keysPressed: 'vgjx',
+      end: ['blah', 'duh', '|ur'],
+    });
+
+    newTest({
+      title: "Preserves cursor position when handling 'gk'",
+      start: ['blah', 'word', 'a', 'la|st'],
+      keysPressed: 'vgkgkx',
+      end: ['blah', 'wo|t'],
+    });
+
+    newTest({
+      title: "Preserves cursor position when handling 'gj'",
+      start: ['blah', 'wo|rd', 'a', 'last'],
+      keysPressed: 'vgjgjx',
+      end: ['blah', 'wo|t'],
+    });
+  });
+
   suite('handles aw in visual mode', () => {
     newTest({
       title: "Can handle 'vawd' on word with cursor inside spaces",

test/mode/modeVisualLine.test.ts

@@ -184,6 +184,22 @@ suite('Mode Visual Line', () => {
     });
   });
 
+  suite('Screen line motions in Visual Line Mode', () => {
+    newTest({
+      title: "Can handle 'gk'",
+      start: ['blah', 'duh', '|dur', 'hur'],
+      keysPressed: 'Vgkx',
+      end: ['blah', '|hur'],
+    });
+
+    newTest({
+      title: "Can handle 'gj'",
+      start: ['blah', 'duh', '|dur', 'hur'],
+      keysPressed: 'Vgjx',
+      end: ['blah', '|duh'],
+    });
+  });
+
   suite('Can handle d/c correctly in Visual Line Mode', () => {
     newTest({
       title: 'Can handle d key',

test/mode/normalModeTests/motions.test.ts

@@ -725,4 +725,32 @@ suite('Motions in Normal Mode', () => {
     keysPressed: '<right>',
     end: ['blah', 'duh', 'd|ur', 'hur'],
   });
+
+  newTest({
+    title: "Can handle 'gk'",
+    start: ['blah', 'duh', '|dur', 'hur'],
+    keysPressed: 'gk',
+    end: ['blah', '|duh', 'dur', 'hur'],
+  });
+
+  newTest({
+    title: "Can handle 'gj'",
+    start: ['blah', 'duh', '|dur', 'hur'],
+    keysPressed: 'gj',
+    end: ['blah', 'duh', 'dur', '|hur'],
+  });
+
+  newTest({
+    title: "Preserves cursor position when handling 'gk'",
+    start: ['blah', 'duh', 'a', 'hu|r '],
+    keysPressed: 'gkgk',
+    end: ['blah', 'du|h', 'a', 'hur '],
+  });
+
+  newTest({
+    title: "Preserves cursor position when handling 'gj'",
+    start: ['blah', 'du|h', 'a', 'hur '],
+    keysPressed: 'gjgj',
+    end: ['blah', 'duh', 'a', 'hu|r '],
+  });
 });