Fix zz updating desired column

src/actions/base.ts

@@ -17,6 +17,12 @@ export class BaseAction {
 
   public canBeRepeatedWithDot = false;
 
+  /**
+   * Whether we should change desiredColumn in VimState.
+   */
+  public doesntChangeDesiredColumn(): boolean {
+    return false;
+  }
   /**
    * Modes that this action can be run in.
    */

src/actions/baseMotion.ts

@@ -53,11 +53,6 @@ export abstract class BaseMovement extends BaseAction {
    */
   isRepeat = false;
 
-  /**
-   * Whether we should change desiredColumn in VimState.
-   */
-  public doesntChangeDesiredColumn = false;
-
   /**
    * This is for commands like $ which force the desired column to be at
    * the end of even the longest line.

src/actions/commands/actions.ts

@@ -711,13 +711,19 @@ abstract class CommandEditorScroll extends BaseCommand {
 @RegisterAction
 class CommandCtrlE extends CommandEditorScroll {
   keys = ['<C-e>'];
+  doesntChangeDesiredColumn() {
+    return true;
+  }
   to: EditorScrollDirection = 'down';
   by: EditorScrollByUnit = 'line';
 }
 
 @RegisterAction
 class CommandCtrlY extends CommandEditorScroll {
   keys = ['<C-y>'];
+  doesntChangeDesiredColumn() {
+    return true;
+  }
   to: EditorScrollDirection = 'up';
   by: EditorScrollByUnit = 'line';
 }
@@ -2440,6 +2446,10 @@ class CommandCenterScroll extends BaseCommand {
   modes = [Mode.Normal, Mode.Visual, Mode.VisualLine, Mode.VisualBlock];
   keys = ['z', 'z'];
 
+  doesntChangeDesiredColumn() {
+    return true;
+  }
+
   public doesActionApply(vimState: VimState, keysPressed: string[]): boolean {
     // Don't run if there's an operator because the Sneak plugin uses <operator>z
     return (
@@ -2490,6 +2500,10 @@ class CommandTopScroll extends BaseCommand {
   modes = [Mode.Normal];
   keys = ['z', 't'];
 
+  doesntChangeDesiredColumn() {
+    return true;
+  }
+
   public doesActionApply(vimState: VimState, keysPressed: string[]): boolean {
     // Don't run if there's an operator because the Sneak plugin uses <operator>z
     return (
@@ -2544,6 +2558,10 @@ class CommandBottomScroll extends BaseCommand {
   modes = [Mode.Normal];
   keys = ['z', 'b'];
 
+  doesntChangeDesiredColumn() {
+    return true;
+  }
+
   public doesActionApply(vimState: VimState, keysPressed: string[]): boolean {
     // Don't run if there's an operator because the Sneak plugin uses <operator>z
     return (

src/actions/motion.ts

@@ -99,7 +99,9 @@ class MoveDownByScreenLine extends MoveByScreenLine {
 }
 
 abstract class MoveByScreenLineMaintainDesiredColumn extends MoveByScreenLine {
-  doesntChangeDesiredColumn = true;
+  doesntChangeDesiredColumn() {
+    return true;
+  }
   public async execAction(position: Position, vimState: VimState): Promise<Position | IMovement> {
     let prevDesiredColumn = vimState.desiredColumn;
     let prevLine = vimState.editor.selection.active.line;
@@ -214,7 +216,9 @@ class MoveDownFoldFix extends MoveByScreenLineMaintainDesiredColumn {
 @RegisterAction
 class MoveDown extends BaseMovement {
   keys = ['j'];
-  doesntChangeDesiredColumn = true;
+  doesntChangeDesiredColumn() {
+    return true;
+  }
 
   public async execAction(position: Position, vimState: VimState): Promise<Position | IMovement> {
     if (configuration.foldfix && vimState.currentMode !== Mode.VisualBlock) {
@@ -237,7 +241,9 @@ class MoveDownArrow extends MoveDown {
 @RegisterAction
 class MoveUp extends BaseMovement {
   keys = ['k'];
-  doesntChangeDesiredColumn = true;
+  doesntChangeDesiredColumn() {
+    return true;
+  }
 
   public async execAction(position: Position, vimState: VimState): Promise<Position | IMovement> {
     if (configuration.foldfix && vimState.currentMode !== Mode.VisualBlock) {
@@ -857,7 +863,9 @@ class MoveUpByScreenLineVisualBlock extends BaseMovement {
     ['g', 'k'],
     ['g', '<up>'],
   ];
-  doesntChangeDesiredColumn = true;
+  doesntChangeDesiredColumn() {
+    return true;
+  }
 
   public async execAction(position: Position, vimState: VimState): Promise<Position | IMovement> {
     return position.getUp(vimState.desiredColumn);
@@ -876,7 +884,9 @@ class MoveDownByScreenLineVisualBlock extends BaseMovement {
     ['g', 'j'],
     ['g', '<down>'],
   ];
-  doesntChangeDesiredColumn = true;
+  doesntChangeDesiredColumn() {
+    return true;
+  }
 
   public async execAction(position: Position, vimState: VimState): Promise<Position | IMovement> {
     return position.getDown(vimState.desiredColumn);

src/mode/modeHandler.ts

@@ -564,8 +564,10 @@ export class ModeHandler implements vscode.Disposable {
     const movement = action instanceof BaseMovement ? action : undefined;
 
     if (
-      (movement && !movement.doesntChangeDesiredColumn) ||
-      (!movement && vimState.currentMode !== Mode.VisualBlock)
+      (movement && !movement.doesntChangeDesiredColumn()) ||
+      (!movement &&
+        vimState.currentMode !== Mode.VisualBlock &&
+        !action.doesntChangeDesiredColumn())
     ) {
       // We check !operator here because e.g. d$ should NOT set the desired column to EOL.
 

test/mode/modeNormal.test.ts

@@ -1093,6 +1093,34 @@ suite('Mode Normal', () => {
     end: ['text te|x'],
   });
 
+  newTest({
+    title: 'A should update desiredColumn',
+    start: ['|longer line', 'short'],
+    keysPressed: 'A<Esc>jk',
+    end: ['longer lin|e', 'short'],
+  });
+
+  newTest({
+    title: 'I should updated desiredColumn',
+    start: ['     hell|o', 'a'],
+    keysPressed: 'I<Esc>jk',
+    end: ['    | hello', 'a'],
+  });
+
+  newTest({
+    title: 'leaving insert mode should update desired column when entered with a',
+    start: ['a long line of text', 'shor|t'],
+    keysPressed: 'amore<Esc>kj',
+    end: ['a long line of text', 'shortmor|e'],
+  });
+
+  newTest({
+    title: 'leaving insert mode should update desired column when entered with i',
+    start: ['a long line of text', 'shor|t'],
+    keysPressed: 'imore<Esc>kj',
+    end: ['a long line of text', 'shormor|et'],
+  });
+
   newTest({
     title: "Can handle 'yy' without changing cursor position",
     start: ['one', 'tw|o'],

test/mode/modeVisualLine.test.ts

@@ -470,5 +470,12 @@ suite('Mode Visual Line', () => {
       keysPressed: 'VjjA_',
       end: ['111', '222_|', ' ', '444_', '555'],
     });
+
+    newTest({
+      title: 'updates desired column correctly',
+      start: ['|111111', '222', '333'],
+      keysPressed: 'VjjA<Esc>jk',
+      end: ['11111|1', '222', '333'],
+    });
   });
 });

test/mode/normalModeTests/motions.test.ts

@@ -803,4 +803,41 @@ suite('Motions in Normal Mode', () => {
     keysPressed: 'gjgj',
     end: ['blah', 'duh', 'a', 'hu|r '],
   });
+
+  suite.only("doesn't update desiredColumn when it shouldn't", () => {
+    newTest({
+      title: 'Preserves desired cursor position when pressing zz',
+      start: ['very long line of text....|.', 'short line'],
+      keysPressed: 'jzzk',
+      end: ['very long line of text....|.', 'short line'],
+    });
+
+    newTest({
+      title: 'Preserves desired cursor position when pressing zt',
+      start: ['very long line of text....|.', 'short line'],
+      keysPressed: 'jztk',
+      end: ['very long line of text....|.', 'short line'],
+    });
+
+    newTest({
+      title: 'Preserves desired cursor position when pressing zb',
+      start: ['very long line of text....|.', 'short line'],
+      keysPressed: 'jzbk',
+      end: ['very long line of text....|.', 'short line'],
+    });
+
+    newTest({
+      title: 'Preserves desired cursor position when pressing <C-e>',
+      start: ['very long line of text....|.', 'short line'],
+      keysPressed: 'j<C-e>k',
+      end: ['very long line of text....|.', 'short line'],
+    });
+
+    newTest({
+      title: 'Preserves desired cursor position when pressing <C-y>',
+      start: ['short line', 'very long line of text....|.'],
+      keysPressed: 'k<C-y>j',
+      end: ['short line', 'very long line of text....|.'],
+    });
+  });
 });