Ensure the desired column number is used while moving up or down in VI mode (#1122)

Author: springcomp

PSReadLine/Movement.cs

@@ -10,6 +10,7 @@ namespace Microsoft.PowerShell
 {
     public partial class PSConsoleReadLine
     {
+        private int _moveToEndOfLineCommandCount;
         private int _moveToLineCommandCount;
         private int _moveToLineDesiredColumn;
 
@@ -125,6 +126,18 @@ private static void ViOffsetCursorPosition(int count)
         }
 
         private void MoveToLine(int lineOffset)
+        {
+            if (InViCommandMode())
+            {
+                ViMoveToLine(lineOffset);
+            }
+            else
+            {
+                MoveToLineImpl(lineOffset);
+            }
+        }
+
+        private void MoveToLineImpl(int lineOffset)
         {
             // Behavior description:
             //  - If the cursor is at the end of a logical line, then 'UpArrow' (or 'DownArrow') moves the cursor up (or down)

PSReadLine/Movement.vi.cs

@@ -156,7 +156,7 @@ public static void ViEndOfPreviousGlob(ConsoleKeyInfo? key = null, object arg =
         private static int ViEndOfLineFactor => InViCommandMode() ? -1 : 0;
 
         /// <summary>
-        /// Move the cursor to the end of the input.
+        /// Move the cursor to the end of the current logical line.
         /// </summary>
         public static void MoveToEndOfLine(ConsoleKeyInfo? key = null, object arg = null)
         {
@@ -166,7 +166,9 @@ public static void MoveToEndOfLine(ConsoleKeyInfo? key = null, object arg = null
                 if (eol != _singleton._current)
                 {
                     _singleton.MoveCursor(eol);
-                }
+                }
+                _singleton._moveToEndOfLineCommandCount++;
+                _singleton._moveToLineDesiredColumn = int.MaxValue;
             }
             else
             {

PSReadLine/Movement.vi.multiline.cs

@@ -56,5 +56,69 @@ public void MoveToLastLine(ConsoleKeyInfo? key = null, object arg = null)
 
             _singleton.MoveCursor(position);
         }
+
+        private void ViMoveToLine(int lineOffset)
+        {
+            // When moving up or down in a buffer in VI mode
+            // the cursor wants to be positioned at a desired column number, which is:
+            // - either a specified column number, the 0-based offset from the start of the logical line.
+            // - or the end of the line
+            //
+            // Only one of those desired position is available at any given time.
+            //
+            // If the desired column number is specified, the cursor will be positioned at
+            // the specified offset in the target logical line, or at the end of the line as appropriate.
+            // The fact that a logical line is shorter than the desired column number *does not*
+            // change its value. If a subsequent move to another logical line is performed, the
+            // desired column number will take effect.
+            //
+            // If the desired column number is the end of the line, the cursor will be positioned at
+            // the end of the target logical line.
+
+            const int endOfLine = int.MaxValue;
+
+            _moveToLineCommandCount += 1;
+
+            // if this is the first "move to line" command
+            // record the desired column number from the current position
+            // on the logical line
+
+            if (_moveToLineCommandCount == 1 && _moveToLineDesiredColumn == -1)
+            {
+                var startOfLine = GetBeginningOfLinePos(_current);
+                _moveToLineDesiredColumn = _current - startOfLine;
+            }
+
+            // Nothing needs to be done when:
+            //  - actually not moving the line, or
+            //  - moving the line down when it's at the end of the last logical line.
+            if (lineOffset == 0 || (lineOffset > 0 && _current == _buffer.Length))
+            {
+                return;
+            }
+
+            int targetLineOffset;
+
+            var currentLineIndex = _singleton.GetLogicalLineNumber() - 1;
+
+            if (lineOffset < 0)
+            {
+                targetLineOffset = Math.Max(0, currentLineIndex + lineOffset);
+            }
+            else
+            {
+                var lastLineIndex = _singleton.GetLogicalLineCount() - 1;
+                targetLineOffset = Math.Min(lastLineIndex, currentLineIndex + lineOffset);
+            }
+
+            var startOfTargetLinePos = GetBeginningOfNthLinePos(targetLineOffset);
+            var endOfTargetLinePos = GetEndOfLogicalLinePos(startOfTargetLinePos);
+
+            var newCurrent = _moveToLineDesiredColumn == endOfLine
+                ? endOfTargetLinePos
+                : Math.Min(startOfTargetLinePos + _moveToLineDesiredColumn, endOfTargetLinePos);
+
+            MoveCursor(newCurrent);
+        }
     }
 }

PSReadLine/Position.cs

@@ -35,6 +35,41 @@ private static int GetBeginningOfLinePos(int current)
             return newCurrent;
         }
 
+        /// <summary>
+        /// Returns the position of the beginning of line
+        /// for the 0-based specified line number.
+        /// </summary>
+        private static int GetBeginningOfNthLinePos(int lineIndex)
+        {
+            System.Diagnostics.Debug.Assert(lineIndex >= 0 || lineIndex < _singleton.GetLogicalLineCount());
+
+            var nth = 0;
+            var index = 0;
+            var result = 0;
+
+            for (; index < _singleton._buffer.Length; index++)
+            {
+                if (nth == lineIndex)
+                {
+                    result = index;
+                    break;
+                }
+
+                if (_singleton._buffer[index] == '\n')
+                {
+                    nth++;
+                }
+            }
+
+            if (nth == lineIndex)
+            {
+                result = index;
+            }
+
+
+            return result;
+        }
+
         /// <summary>
         /// Returns the position of the end of the logical line
         /// as specified by the "current" position.

PSReadLine/ReadLine.cs

@@ -476,7 +476,8 @@ private string InputLoop()
                 var recallHistoryCommandCount = _recallHistoryCommandCount;
                 var yankLastArgCommandCount = _yankLastArgCommandCount;
                 var visualSelectionCommandCount = _visualSelectionCommandCount;
-                var movingAtEndOfLineCount = _moveToLineCommandCount;
+                var moveToLineCommandCount = _moveToLineCommandCount;
+                var moveToEndOfLineCommandCount = _moveToEndOfLineCommandCount;
 
                 var key = ReadKey();
                 ProcessOneKey(key, _dispatchTable, ignoreIfNoAction: false, arg: null);
@@ -537,9 +538,18 @@ private string InputLoop()
                     _visualSelectionCommandCount = 0;
                     Render();  // Clears the visual selection
                 }
-                if (movingAtEndOfLineCount == _moveToLineCommandCount)
+                if (moveToLineCommandCount == _moveToLineCommandCount)
                 {
-                    _moveToLineCommandCount = 0;
+                    _moveToLineCommandCount = 0;
+
+                    if (InViCommandMode() && moveToEndOfLineCommandCount == _moveToEndOfLineCommandCount)
+                    {
+                        // the previous command was neither a "move to end of line" command
+                        // nor a "move to line" command. In that case, the desired column
+                        // number will be computed from the current position on the logical line.
+
+                        _moveToLineDesiredColumn = -1;
+                    }
                 }
             }
         }

test/MovementTest.VI.Multiline.cs

@@ -5,6 +5,46 @@ namespace Test
 {
     public partial class ReadLine
     {
+        [SkippableFact]
+        public void ViMoveToLine_DesiredColumn()
+        {
+            TestSetup(KeyMode.Vi);
+
+            const string buffer = "\"\n12345\n1234\n123\n12\n1\n\"";
+
+            var continuationPrefixLength = PSConsoleReadLineOptions.DefaultContinuationPrompt.Length;
+
+            Test(buffer, Keys(
+                _.DQuote, _.Enter,
+                "12345", _.Enter,
+                "1234", _.Enter,
+                "123", _.Enter,
+                "12", _.Enter,
+                "1", _.Enter,
+                _.DQuote,
+                _.Escape,
+
+                // move to second line at column 4
+                "ggj3l", CheckThat(() => AssertCursorLeftIs(continuationPrefixLength + 3)),
+                // moving down on shorter lines will position the cursor at the end of each logical line
+                _.j, CheckThat(() => AssertCursorLeftIs(continuationPrefixLength + 3)),
+                _.j, CheckThat(() => AssertCursorLeftIs(continuationPrefixLength + 2)),
+                // moving back up will position the cursor at the end of shorter lines or at the desired column number
+                _.k, CheckThat(() => AssertCursorLeftIs(continuationPrefixLength + 3)),
+                _.k, CheckThat(() => AssertCursorLeftIs(continuationPrefixLength + 3)),
+
+                // move at end of line (column 5)
+                _.Dollar, CheckThat(() => AssertCursorLeftIs(continuationPrefixLength + 4)),
+                // moving down on shorter lines will position the cursor at the end of each logical line
+                _.j, CheckThat(() => AssertCursorLeftIs(continuationPrefixLength + 3)),
+                _.j, CheckThat(() => AssertCursorLeftIs(continuationPrefixLength + 2)),
+                // moving back up will position the cursor at the end of each logical line
+                _.k, CheckThat(() => AssertCursorLeftIs(continuationPrefixLength + 3)),
+                _.k, CheckThat(() => AssertCursorLeftIs(continuationPrefixLength + 4))
+            ));
+        }
+
+
         [SkippableFact]
         public void ViBackwardChar()
         {