Fix gj/gk so it maintains cursor position

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 '],
+  });
 });