PowerShell · springcomp · Aug 18, 2023 · Aug 18, 2023 · Dec 3, 2020 · Aug 18, 2023
diff --git a/PSReadLine/History.cs b/PSReadLine/History.cs
@@ -87,6 +87,8 @@ public class HistoryItem
             internal List<EditItem> _edits;
             internal int _undoEditIndex;
             internal int _editGroupStart;
+            internal GroupUndoHelper _groupUndoHelper;
+            internal Stack<GroupUndoState> _groupUndoStates = new();
         }
 
         // History state
@@ -138,6 +140,8 @@ private void ClearSavedCurrentLine()
             _savedCurrentLine._edits = null;
             _savedCurrentLine._undoEditIndex = 0;
             _savedCurrentLine._editGroupStart = -1;
+            _savedCurrentLine._groupUndoHelper.Clear();
+            _savedCurrentLine._groupUndoStates.Clear();
         }
 
         private AddToHistoryOption GetAddToHistoryOption(string line, bool fromHistoryFile)
@@ -787,13 +791,19 @@ private void UpdateFromHistory(HistoryMoveCursor moveCursor)
                 _edits = new List<EditItem>(_savedCurrentLine._edits);
                 _undoEditIndex = _savedCurrentLine._undoEditIndex;
                 _editGroupStart = _savedCurrentLine._editGroupStart;
+                _groupUndoHelper = _savedCurrentLine._groupUndoHelper;
+
+                _savedCurrentLine._groupUndoStates.CopyTo(_groupUndoStates);
             }
             else
             {
                 line = _history[_currentHistoryIndex].CommandLine;
                 _edits = new List<EditItem>(_history[_currentHistoryIndex]._edits);
                 _undoEditIndex = _history[_currentHistoryIndex]._undoEditIndex;
                 _editGroupStart = _history[_currentHistoryIndex]._editGroupStart;
+                _groupUndoHelper = _history[_currentHistoryIndex]._groupUndoHelper;
+
+                _history[_currentHistoryIndex]._groupUndoStates.CopyTo(_groupUndoStates);
             }
             _buffer.Clear();
             _buffer.Append(line);
@@ -831,6 +841,9 @@ private void SaveCurrentLine()
                 _savedCurrentLine._edits = _edits;
                 _savedCurrentLine._undoEditIndex = _undoEditIndex;
                 _savedCurrentLine._editGroupStart = _editGroupStart;
+                _savedCurrentLine._groupUndoHelper = _groupUndoHelper;
+
+                _groupUndoStates.CopyTo(_savedCurrentLine._groupUndoStates);
             }
         }
 

diff --git a/PSReadLine/ReadLine.cs b/PSReadLine/ReadLine.cs
@@ -36,7 +36,7 @@ public partial class PSConsoleReadLine : IPSConsoleReadLineMockableMethods
         // *must* be initialized in the static ctor
         // because the static member _clipboard depends upon it
         // for its own initialization
-        private static readonly PSConsoleReadLine _singleton;
+        internal static readonly PSConsoleReadLine _singleton;
 
         // This is used by PowerShellEditorServices (the backend of the PowerShell VSCode extension)
         // so that it can call PSReadLine from a delegate and not hit nested pipeline issues.

diff --git a/PSReadLine/ReadLine.vi.cs b/PSReadLine/ReadLine.vi.cs
@@ -573,6 +573,8 @@ public static void ViInsertAtEnd(ConsoleKeyInfo? key = null, object arg = null)
         /// </summary>
         public static void ViInsertWithAppend(ConsoleKeyInfo? key = null, object arg = null)
         {
+            _singleton._groupUndoHelper.StartGroup(ViInsertWithAppend, arg);
+
             ViInsertMode(key, arg);
             ForwardChar(key, arg);
         }

diff --git a/PSReadLine/Replace.vi.cs b/PSReadLine/Replace.vi.cs
@@ -277,6 +277,10 @@ private static void ViReplaceToChar(char keyChar, ConsoleKeyInfo? key = null, ob
                     ViInsertMode(key, arg);
                 }
             }
+            else
+            {
+                _singleton._groupUndoHelper.EndGroup();
+            }
         }
 
         /// <summary>
@@ -295,6 +299,10 @@ private static void ViReplaceToCharBack(char keyChar, ConsoleKeyInfo? key = null
             {
                 ViInsertMode(key, arg);
             }
+            else
+            {
+                _singleton._groupUndoHelper.EndGroup();
+            }
         }
 
         /// <summary>
@@ -314,6 +322,10 @@ private static void ViReplaceToBeforeChar(char keyChar, ConsoleKeyInfo? key = nu
             {
                 ViInsertMode(key, arg);
             }
+            else
+            {
+                _singleton._groupUndoHelper.EndGroup();
+            }
         }
 
         /// <summary>
@@ -332,6 +344,10 @@ private static void ViReplaceToBeforeCharBack(char keyChar, ConsoleKeyInfo? key
             {
                 ViInsertMode(key, arg);
             }
+            else
+            {
+                _singleton._groupUndoHelper.EndGroup();
+            }
         }
 
 

diff --git a/PSReadLine/UndoRedo.cs b/PSReadLine/UndoRedo.cs
@@ -59,8 +59,8 @@ private void SaveEditItem(EditItem editItem)
             _undoEditIndex = _edits.Count;
         }
 
-        private void StartEditGroup()
-        {
+        internal void StartEditGroup()
+        {
             if (_editGroupStart != -1)
             {
                 // Nesting not supported.
@@ -71,7 +71,7 @@ private void StartEditGroup()
             _editGroupStart = _edits.Count;
         }
 
-        private void EndEditGroup(Action<ConsoleKeyInfo?, object> instigator = null, object instigatorArg = null)
+        internal void EndEditGroup(Action<ConsoleKeyInfo?, object> instigator = null, object instigatorArg = null)
         {
             // Remove the undone edits when closing an edit group, so the generated group
             // doesn't contain those edits that were already undone.
@@ -94,7 +94,8 @@ private void EndEditGroup(Action<ConsoleKeyInfo?, object> instigator = null, obj
                 _edits.RemoveRange(_editGroupStart, groupEditCount);
                 SaveEditItem(GroupedEdit.Create(groupedEditItems, instigator, instigatorArg));
             }
-            _editGroupStart = -1;
+
+            _editGroupStart = -1;
         }
 
         private bool IsLastEditItemReplaceable

diff --git a/PSReadLine/UndoRedo.vi.cs b/PSReadLine/UndoRedo.vi.cs
@@ -3,24 +3,36 @@
 --********************************************************************/
 
 using System;
+using System.Collections.Generic;
 
 namespace Microsoft.PowerShell
 {
     public partial class PSConsoleReadLine
     {
-        private class GroupUndoHelper
+        /// <summary>
+        /// Encapsulates state and behaviour for an edit group.
+        ///
+        /// An edit group is a sequence of distinct internal changes
+        /// to the buffer that are considered to be a single user-facing
+        /// change for the purpose of the undo/redo actions.
+        /// </summary>
+        internal struct GroupUndoHelper
         {
             public Action<ConsoleKeyInfo?, object> _instigator;
             public object _instigatorArg;
 
-            public GroupUndoHelper()
-            {
-                _instigator = null;
-                _instigatorArg = null;
-            }
-
             public void StartGroup(Action<ConsoleKeyInfo?, object> instigator, object instigatorArg)
             {
+                if (_singleton._editGroupStart != -1)
+                {
+                    // A nested "start" of a group is being made, so we
+                    // need to record the state of the preceding start.
+                    _singleton._groupUndoStates.Push(
+                        new GroupUndoState(_singleton._groupUndoHelper, _singleton._editGroupStart));
+
+                    _singleton._editGroupStart = -1;
+                }
+
                 _instigator = instigator;
                 _instigatorArg = instigatorArg;
                 _singleton.StartEditGroup();
@@ -34,14 +46,60 @@ public void Clear()
 
             public void EndGroup()
             {
-                if (_singleton._editGroupStart >= 0)
-                {
-                    _singleton.EndEditGroup(_instigator, _instigatorArg);
+                if (_singleton._editGroupStart >= 0)
+                {
+                    _singleton.EndEditGroup(
+                        _singleton._groupUndoHelper._instigator,
+                        _singleton._groupUndoHelper._instigatorArg);
+
+                    while (_singleton._groupUndoStates.Count > 0)
+                    {
+                        var groupUndoState = _singleton._groupUndoStates.Pop();
+                        _singleton._editGroupStart = groupUndoState.EditGroupStart;
+                        _singleton._groupUndoHelper = groupUndoState.GroupUndoHelper;
+
+                        _singleton.EndEditGroup(
+                            _singleton._groupUndoHelper._instigator,
+                            _singleton._groupUndoHelper._instigatorArg);
+                    }
                 }
-                Clear();
+
+                _singleton._groupUndoHelper.Clear();
+            }
+        }
+
+        internal GroupUndoHelper _groupUndoHelper = new();
+
+        /// <summary>
+        /// Records states of changes made as part of an edit group.
+        ///
+        /// In an ideal situation, edit groups would be started and ended
+        /// in balanced pairs of calls. However, we expose public methods
+        /// that may start an edit group and rely on future actions to
+        /// properly end the group.
+        ///
+        /// To improve robustness of the code, we allow starting "nested"
+        /// edit groups and end the whole sequence of groups once. That is,
+        /// a single call to _singleton.EndEditGroup() will properly record the
+        /// changes made up to that point from calls to _singleton.StartEditGroup()
+        /// that have been made at different points in the overall sequence of changes.
+        /// </summary>
+        internal class GroupUndoState
+        {
+            public GroupUndoState(GroupUndoHelper undoHelper, int editGroupStart)
+            {
+                GroupUndoHelper = undoHelper;
+                EditGroupStart = editGroupStart;
             }
+
+            public GroupUndoHelper GroupUndoHelper { get; }
+            public int EditGroupStart { get; }
         }
-        private readonly GroupUndoHelper _groupUndoHelper = new GroupUndoHelper();
+
+        /// <summary>
+        /// Records the sequence of "nested" starts of a edit group.
+        /// </summary>
+        private readonly Stack<GroupUndoState> _groupUndoStates = new();
 
         /// <summary>
         /// Undo all previous edits for line.
@@ -63,4 +121,23 @@ public static void UndoAll(ConsoleKeyInfo? key = null, object arg = null)
             }
         }
     }
+
+    internal static class StackExtensions
+    {
+        /// <summary>
+        /// This helper method copies the contents of the target stack
+        /// with items from the supplied stack.
+        /// </summary>
+        /// <typeparam name="T"></typeparam>
+        /// <param name="origin"></param>
+        /// <param name="target"></param>
+        public static void CopyTo<T>(this Stack<T> origin, Stack<T> target)
+        {
+            target.Clear();
+            foreach (var item in origin)
+            {
+                target.Push(item);
+            }
+        }
+    }
 }
diff --git a/test/BasicEditingTest.VI.cs b/test/BasicEditingTest.VI.cs
@@ -1121,5 +1121,45 @@ public void ViInsertModeMoveCursor()
                 _.RightArrow, // 'RightArrow' again does nothing, but doesn't crash
                 "c"));
         }
+
+        [SkippableFact]
+        public void ViDefect1281()
+        {
+            TestSetup(KeyMode.Vi);
+
+            Test("bcd", Keys(
+                "abcdabcd", _.Escape,
+
+                // return to the [B]eginning of the word,
+                // then [c]hange text un[t]il just before the [2]nd [b] character
+                // this leaves the cursor at the current position (0) but erases
+                // the "abcda" text portion, / switches to edit mode and
+                // positions the cursor just before the "bcd" text portion.
+
+                "Bc2tb",
+
+                // going back to normal mode again without having modified the buffer further
+                // even though the [c] command started an edit group, going back to normal
+                // mode closes the pending edit group.
+
+                _.Escape, CheckThat(() => AssertCursorLeftIs(0)),
+
+                // attempt to [c]hange text un[t]il just before the [2]nd [b] character again
+                // because the [b] character only appears once further down in the buffer
+                // relative to where the cursor position is – currently set to 0 - the command
+                // fails. Therefore, we are still in normal mode.
+                //
+                // however, even though the command failed, it still started an edit group
+                // which is now pending further edit actions
+
+                "c2tb", CheckThat(() => AssertLineIs("bcd")),
+
+                // attempt to [c]hange text un[t]il just before the [2]nd [b] character a third time.
+                // this exercises a code path where starting a edit group while another
+                // pending edit group was previously started crashed PSRL.
+
+                "c2tb" // should not crash
+                ));
+        }
     }
 }
diff --git a/test/UndoRedoTest.cs b/test/UndoRedoTest.cs
@@ -2,6 +2,8 @@
 
 namespace Test
 {
+    using PSConsoleReadLine = Microsoft.PowerShell.PSConsoleReadLine;
+
     public partial class ReadLine
     {
         [SkippableFact]
@@ -56,5 +58,49 @@ public void ProperlySetEditGroupStartMarkWhenLoopInHistory()
                 _.Escape, _.u, CheckThat(() => AssertLineIs("yuiogh"))
             ));
         }
+
+        [SkippableFact]
+        public void SupportNestedEditGroups()
+        {
+            TestSetup(KeyMode.Vi);
+
+            Test(" ", Keys(
+                " ", _.Escape,
+                CheckThat(() =>
+                {
+                    PSConsoleReadLine._singleton._groupUndoHelper.StartGroup((k, o) => { }, null);
+                    PSConsoleReadLine._singleton._groupUndoHelper.EndGroup();
+                })));
+
+            Test("012 45", Keys(
+                "012 45", _.Escape,
+                "b", // back one word
+                CheckThat(() =>
+                {
+                    PSConsoleReadLine._singleton._groupUndoHelper.StartGroup((k, o) => { }, null);
+
+                    // In an ideal situation, edit groups would be started and ended
+                    // in balanced pairs of calls. However, we expose public methods
+                    // that may start an edit group and rely on future actions to
+                    // properly end the group.
+                    //
+                    // To improve robustness of the code, we allow starting "nested"
+                    // edit groups and end the whole sequence of groups once. That is,
+                    // a single call to _singleton.EndEditGroup() will properly record the
+                    // changes made up to that point from calls to _singleton.StartEditGroup()
+                    // that have been made at different points in the overall sequence of changes.
+
+                    // Thus, the following, unneeded call, should not crash.
+
+                    PSConsoleReadLine._singleton._groupUndoHelper.StartGroup((k, o) => { }, null);
+
+                    PSConsoleReadLine._singleton._groupUndoHelper.EndGroup();
+
+                    // Likewise, closing an edit group that is no longer open
+                    // is now a no-op. That is, the following line should not crash either.
+
+                    PSConsoleReadLine._singleton._groupUndoHelper.EndGroup();
+                })));
+        }
     }
 }