[WIP] Improve foldFix further by making more motion prefixed actions fold-aware

src/actions/motion.ts

@@ -136,7 +136,13 @@ export abstract class BaseMovement extends BaseAction {
     for (let i = 0; i < count; i++) {
       const firstIteration = i === 0;
       const lastIteration = i === count - 1;
-      result = await this.createMovementResult(position, vimState, recordedState, lastIteration);
+      result = await this.createMovementResult(
+        position,
+        vimState,
+        recordedState,
+        firstIteration,
+        lastIteration
+      );
 
       if (result instanceof Position) {
         position = result;
@@ -171,6 +177,7 @@ export abstract class BaseMovement extends BaseAction {
     position: Position,
     vimState: VimState,
     recordedState: RecordedState,
+    firstIteration: boolean,
     lastIteration: boolean
   ): Promise<Position | IMovement> {
     const result =
@@ -179,6 +186,7 @@ export abstract class BaseMovement extends BaseAction {
         : await this.execAction(position, vimState);
     return result;
   }
+
   protected adjustPosition(position: Position, result: IMovement, lastIteration: boolean) {
     if (!lastIteration) {
       position = result.stop.getRightThroughLineBreaks();
@@ -295,56 +303,118 @@ class MoveDownByScreenLineMaintainDesiredColumn extends MoveByScreenLineMaintain
   value = 1;
 }
 
-class MoveDownFoldFix extends MoveByScreenLineMaintainDesiredColumn {
+export abstract class FoldFixMovement extends BaseMovement {
+  protected async createMovementResult(
+    position: Position,
+    vimState: VimState,
+    recordedState: RecordedState,
+    firstIteration: boolean,
+    lastIteration: boolean
+  ): Promise<Position | IMovement> {
+    const isFirstUp = firstIteration && this instanceof MoveUp;
+    const includeFoldedRegionBelowBeforeUp =
+      configuration.foldfix && recordedState.operator && isFirstUp;
+
+    let result = includeFoldedRegionBelowBeforeUp
+      ? <IMovement>await new MoveDownFoldFix().execAction(position, vimState)
+      : <IMovement>(
+          await super.createMovementResult(
+            position,
+            vimState,
+            recordedState,
+            firstIteration,
+            lastIteration
+          )
+        );
+
+    // exclude line just below folded region added before first up
+    if (includeFoldedRegionBelowBeforeUp) {
+      result.start = result.stop.getUp(vimState.desiredColumn);
+      result.stop = (<IMovement>await new MoveUpFoldFix().execAction(position, vimState)).stop;
+    }
+
+    const lastDownBeforeAction =
+      lastIteration && configuration.foldfix && recordedState.operator && this instanceof MoveDown;
+    if (lastDownBeforeAction) {
+      // include folded region of last line after action prefixed down motion
+      result.stop = (<IMovement>(
+        await new MoveDownFoldFix().execAction(position, vimState)
+      )).stop.getUp(vimState.desiredColumn);
+    }
+    return result;
+  }
+}
+export class MoveDownFoldFix extends MoveByScreenLineMaintainDesiredColumn {
   movementType: CursorMovePosition = 'down';
   by: CursorMoveByUnit = 'line';
   value = 1;
 
   public async execAction(position: Position, vimState: VimState): Promise<Position | IMovement> {
     if (position.line >= TextEditor.getLineCount() - 1) {
-      return position;
+      return { start: position, stop: position };
     }
     let t: Position | IMovement;
+    let tStop: Position;
     let prevLine: number = position.line;
     let prevChar: number = position.character;
     const prevDesiredColumn = vimState.desiredColumn;
     const moveDownByScreenLine = new MoveDownByScreenLine();
     do {
       t = <Position | IMovement>await moveDownByScreenLine.execAction(position, vimState);
-      t = t instanceof Position ? t : t.stop;
-      const lineChanged = prevLine !== t.line;
+      tStop = t instanceof Position ? t : t.stop;
+      const lineChanged = prevLine !== tStop.line;
       // wrappedLine movement goes to eol character only when at the last line
       // thus a column change on wrappedLine movement represents a visual last line
-      const colChanged = prevChar !== t.character;
+      const colChanged = prevChar !== tStop.character;
       if (lineChanged || !colChanged) {
         break;
       }
-      prevChar = t.character;
-      prevLine = t.line;
-    } while (t.line === position.line);
+      prevChar = tStop.character;
+      prevLine = tStop.line;
+    } while (tStop.line === position.line);
+
     // fix column change at last line caused by wrappedLine movement
     // causes cursor lag and flicker if a large repeat prefix is given to movement
-    if (t.character !== prevDesiredColumn) {
-      t = new Position(t.line, prevDesiredColumn);
+    if (tStop.character !== prevDesiredColumn) {
+      tStop = new Position(tStop.line, prevDesiredColumn);
     }
-    return t;
+
+    // create and return an `IMovement` to support `execActionForOperator`
+    const tMovement: IMovement = { start: position, stop: tStop };
+    return tMovement;
   }
 }
 
 @RegisterAction
-class MoveDown extends BaseMovement {
+class MoveDown extends FoldFixMovement {
   keys = ['j'];
   doesntChangeDesiredColumn = true;
 
   public async execAction(position: Position, vimState: VimState): Promise<Position | IMovement> {
     if (configuration.foldfix && vimState.currentMode !== ModeName.VisualBlock) {
-      return new MoveDownFoldFix().execAction(position, vimState);
+      if (vimState.currentMode === ModeName.Normal) {
+        return await new MoveDownFoldFix().execAction(position, vimState);
+      } else {
+        // insert and visual modes break given an IMovement
+        return (<IMovement>await new MoveDownFoldFix().execAction(position, vimState)).stop;
+      }
     }
     return position.getDown(vimState.desiredColumn);
   }
 
-  public async execActionForOperator(position: Position, vimState: VimState): Promise<Position> {
+  public async execActionForOperator(
+    position: Position,
+    vimState: VimState
+  ): Promise<Position | IMovement> {
     vimState.currentRegisterMode = RegisterMode.LineWise;
+
+    if (configuration.foldfix && vimState.currentMode !== ModeName.VisualBlock) {
+      let actionRange = <IMovement>await new MoveDownFoldFix().execAction(position, vimState);
+      // adjust range to exclude line after wrapped and/or folded block
+      actionRange.stop = actionRange.stop.getUp(vimState.desiredColumn);
+      return actionRange;
+    }
+
     return position.getDown(position.getLineEnd().character);
   }
 }
@@ -361,46 +431,63 @@ class MoveUpByScreenLineMaintainDesiredColumn extends MoveByScreenLineMaintainDe
 }
 
 @RegisterAction
-class MoveUp extends BaseMovement {
+class MoveUp extends FoldFixMovement {
   keys = ['k'];
   doesntChangeDesiredColumn = true;
 
   public async execAction(position: Position, vimState: VimState): Promise<Position | IMovement> {
     if (configuration.foldfix && vimState.currentMode !== ModeName.VisualBlock) {
-      return new MoveUpFoldFix().execAction(position, vimState);
+      if (vimState.currentMode === ModeName.Normal) {
+        return await new MoveUpFoldFix().execAction(position, vimState);
+      } else {
+        return (<IMovement>await new MoveUpFoldFix().execAction(position, vimState)).stop;
+      }
     }
     return position.getUp(vimState.desiredColumn);
   }
 
-  public async execActionForOperator(position: Position, vimState: VimState): Promise<Position> {
+  public async execActionForOperator(
+    position: Position,
+    vimState: VimState
+  ): Promise<Position | IMovement> {
     vimState.currentRegisterMode = RegisterMode.LineWise;
+    if (configuration.foldfix && vimState.currentMode !== ModeName.VisualBlock) {
+      const stop = (<IMovement>await new MoveUpFoldFix().execAction(position, vimState)).stop;
+      const actionRange: IMovement = { start: position, stop };
+      return actionRange;
+    }
     return position.getUp(position.getLineEnd().character);
   }
 }
 
 @RegisterAction
-class MoveUpFoldFix extends MoveByScreenLineMaintainDesiredColumn {
+export class MoveUpFoldFix extends MoveByScreenLineMaintainDesiredColumn {
   movementType: CursorMovePosition = 'up';
   by: CursorMoveByUnit = 'line';
   value = 1;
 
   public async execAction(position: Position, vimState: VimState): Promise<Position | IMovement> {
     if (position.line === 0) {
-      return position;
+      return { start: position, stop: position };
     }
     let t: Position | IMovement;
+    let tStop: Position;
     const prevDesiredColumn = vimState.desiredColumn;
     const moveUpByScreenLine = new MoveUpByScreenLine();
     do {
       t = <Position | IMovement>await moveUpByScreenLine.execAction(position, vimState);
-      t = t instanceof Position ? t : t.stop;
-    } while (t.line === position.line);
+      tStop = t instanceof Position ? t : t.stop;
+    } while (tStop.line === position.line);
+
     // fix column change at last line caused by wrappedLine movement
     // causes cursor lag and flicker if a large repeat prefix is given to movement
-    if (t.character !== prevDesiredColumn) {
-      t = new Position(t.line, prevDesiredColumn);
+    if (tStop.character !== prevDesiredColumn) {
+      tStop = new Position(tStop.line, prevDesiredColumn);
     }
-    return t;
+
+    // create and return an `IMovement` to support `execActionForOperator`
+    const tMovement: IMovement = { start: position, stop: tStop };
+    return tMovement;
   }
 }
 

src/actions/operator.ts

@@ -10,6 +10,7 @@ import { TextEditor } from './../textEditor';
 import { BaseAction, RegisterAction } from './base';
 import { CommandNumber } from './commands/actions';
 import { TextObjectMovement } from './textobject';
+import { IMovement, MoveDownFoldFix, MoveUpFoldFix } from './motion';
 
 export class BaseOperator extends BaseAction {
   constructor(multicursorIndex?: number) {
@@ -121,8 +122,11 @@ export class DeleteOperator extends BaseOperator {
     if (registerMode === RegisterMode.LineWise) {
       start = start.getLineBegin();
       end = end.getLineEnd();
-    }
 
+      if (configuration.foldfix) {
+        end = await includeFoldedLines(start, end, vimState);
+      }
+    }
     end = new Position(end.line, end.character + 1);
 
     const isOnLastLine = end.line === TextEditor.getLineCount() - 1;
@@ -263,6 +267,10 @@ export class YankOperator extends BaseOperator {
     if (vimState.currentRegisterMode === RegisterMode.LineWise) {
       start = start.getLineBegin();
       end = end.getLineEnd();
+
+      if (configuration.foldfix) {
+        end = await includeFoldedLines(start, end, vimState);
+      }
     }
 
     let text = TextEditor.getText(new vscode.Range(start, end));
@@ -300,6 +308,25 @@ export class YankOperator extends BaseOperator {
   }
 }
 
+async function includeFoldedLines(
+  start: Position,
+  end: Position,
+  vimState: VimState
+): Promise<Position> {
+  const isddCommand = vimState.recordedState.commandWithoutCountPrefix === 'dd';
+  const isyyCommand = vimState.recordedState.commandWithoutCountPrefix === 'yy';
+
+  if (isddCommand || isyyCommand) {
+    const startEndLineDistance = Math.abs(end.line - start.line);
+    end = new Position(start.line, end.character);
+    for (let len = startEndLineDistance + 1; len > 0; len--) {
+      end = (<IMovement>await new MoveDownFoldFix().execAction(end, vimState)).stop;
+    }
+    end = end.getUp(0).getLineEnd();
+  }
+  return end;
+}
+
 @RegisterAction
 export class ShiftYankOperatorVisual extends BaseOperator {
   public keys = ['Y'];