Allow batch undo

madskristensen · Aug 6, 2023 · 95870e2 · 95870e2
1 parent fb4ec27
commit 95870e2
Show file tree

Hide file tree

Showing 4 changed files with 83 additions and 39 deletions.
diff --git a/src/Commands/IncrementalShiftDownCommand.cs b/src/Commands/IncrementalShiftDownCommand.cs
@@ -14,7 +14,7 @@ protected override async Task ExecuteAsync(OleMenuCmdEventArgs e)
                 return;
             }
 
-            ShiftDownCommand.Shift(docView, ShiftDirection.Down, true);
+            await ShiftDownCommand.ShiftAsync(docView, ShiftDirection.Down, true);
         }
     }
 }
diff --git a/src/Commands/IncrementalShiftUpCommand.cs b/src/Commands/IncrementalShiftUpCommand.cs
@@ -14,7 +14,7 @@ protected override async Task ExecuteAsync(OleMenuCmdEventArgs e)
                 return;
             }
 
-            ShiftDownCommand.Shift(docView, ShiftDirection.Up, true);
+            await ShiftDownCommand.ShiftAsync(docView, ShiftDirection.Up, true);
         }
     }
 }
diff --git a/src/Commands/ShiftDownCommand.cs b/src/Commands/ShiftDownCommand.cs
@@ -1,7 +1,9 @@
 using System.Collections.Generic;
 using System.Linq;
+using Microsoft.VisualStudio.OLE.Interop;
 using Microsoft.VisualStudio.Text;
 using Microsoft.VisualStudio.Text.Editor;
+using Microsoft.VisualStudio.Text.Operations;
 using Shifter.Providers;
 using Selection = Microsoft.VisualStudio.Text.Selection;
 
@@ -19,10 +21,10 @@ protected override async Task ExecuteAsync(OleMenuCmdEventArgs e)
                 return;
             }
 
-            Shift(docView, ShiftDirection.Down, false);
+            await ShiftAsync(docView, ShiftDirection.Down, false);
         }
 
-        public static void Shift(DocumentView docView, ShiftDirection direction, bool incremental)
+        public static async Task ShiftAsync(DocumentView docView, ShiftDirection direction, bool incremental)
         {
             IWpfTextView textView = docView.TextView!;
             IMultiSelectionBroker multiSelectionBroker = textView.GetMultiSelectionBroker();
@@ -75,45 +77,87 @@ public static void Shift(DocumentView docView, ShiftDirection direction, bool in
                     (Selection, ShiftResult)[] shifts = shiftResults
                         .Select<ShiftResult, (Selection Selection, ShiftResult ShiftResult)>((shiftResult, i) => (selections[i], shiftResult))
                         .Where(static pair => pair.ShiftResult != null)
-                        // We need descending order to preserve positions after shift transformations
-                        .OrderByDescending(static pair => pair.Selection.Start.Position.Position)
+                        // We need ascending order to preserve selections after shift transformations, we'll calculate cumulative difference
+                        .OrderBy(static pair => pair.Selection.Start.Position.Position)
                         .ToArray();
 
-                    ITextSnapshot updatedSnapshot = snapshot;
-                    Selection[] newSelections = shifts
-                        .Select(
-                            pair =>
-                            {
-                                (Selection selection, ShiftResult shiftResult) = pair;
-                                // Update text buffer
-                                Span span = new(shiftResult.Start + selection.Start.Position, shiftResult.Length);
-                                updatedSnapshot = updatedSnapshot.TextBuffer.Replace(span, shiftResult.ShiftedText);
-
-                                // Save new caret position and selection
-                                SnapshotSpan newSpan = new(updatedSnapshot, shiftResult.Start + selection.Start.Position, shiftResult.ShiftedText.Length);
-                                SnapshotPoint newEndPosition = new(updatedSnapshot, newSpan.Contains(selection.InsertionPoint.Position.Position) ? selection.InsertionPoint.Position.Position : newSpan.End.Position);
-                                Selection newSelection = new(newEndPosition, newSpan.Start, newSpan.End);
-                                return newSelection;
-                            })
-                        .ToArray();
+                    ITextBuffer buffer = textView.TextBuffer;
+                    string undoText = (direction, incremental) switch
+                    {
+                        (ShiftDirection.Up, false)   => "Shift up",
+                        (ShiftDirection.Down, false) => "Shift down",
+                        (ShiftDirection.Up, true)    => "Incremental Shift up",
+                        (ShiftDirection.Down, true)  => "Incremental Shift down",
+                        _                            => "Shifter"
+                    };
+
+                    using ITextUndoTransaction undo = await buffer.OpenUndoContextAsync(undoText);
+                    Selection[] newSelections;
+                    using (ITextEdit edit = buffer.CreateEdit())
+                    {
+                        int diff = 0;
+
+                        // Collect positions, we have to create new snapshot after all changes are applied with single edit
+                        List<(int InsertionPoint, Span Span)> newSelectionPositions = shifts
+                            .Select(
+                                pair =>
+                                {
+                                    (Selection selection, ShiftResult shiftResult) = pair;
+
+                                    // Update text buffer
+                                    Span span = new(shiftResult.Start + selection.Start.Position, shiftResult.Length);
+                                    edit.Replace(span, shiftResult.ShiftedText);
+
+                                    Span newSpan = new(diff + shiftResult.Start + selection.Start.Position, shiftResult.ShiftedText.Length);
+
+                                    // Cumulate total difference so far so spans will match the new ones
+                                    diff += shiftResult.ShiftedText.Length - span.Length;
+
+                                    // Try to remain insertion point
+                                    int insertionPoint = newSpan.Contains(selection.InsertionPoint.Position.Position + diff)
+                                        ? selection.InsertionPoint.Position.Position
+                                        : newSpan.End;
+                                    return (insertionPoint, newSpan);
+                                })
+                            .ToList();
+
+                        ITextSnapshot updatedSnapshot = edit.Apply();
+
+                        // Create new selections from calculated new positions
+                        newSelections = new Selection[newSelectionPositions.Count];
+                        for (int i = 0; i < newSelectionPositions.Count; i++)
+                        {
+                            (int insertionPoint, Span span) = newSelectionPositions[i];
+                            SnapshotSpan newSpan = new(updatedSnapshot, span);
+                            SnapshotPoint newEndPosition = new(updatedSnapshot, insertionPoint);
+                            Selection newSelection = new(newEndPosition, newSpan.Start, newSpan.End);
+                            newSelections[i] = newSelection;
+                        }
+                    }
 
-                    // Refresh selections
-                    selections = multiSelectionBroker.AllSelections
-                        // New selections are also ordered by position in descending order
-                        .OrderByDescending(static selection => selection.Start.Position.Position)
-                        .ToArray();
-                    for (int i = 0; i < selections.Count; i++)
+                    using (multiSelectionBroker.BeginBatchOperation())
                     {
-                        Selection selection = selections[i];
-                        Selection newSelection = newSelections[i];
-                        multiSelectionBroker.TryPerformActionOnSelection(
-                            selection,
-                            transformer =>
-                            {
-                                // Move selection to new position
-                                transformer.MoveTo(newSelection.AnchorPoint, newSelection.ActivePoint, newSelection.InsertionPoint, newSelection.InsertionPointAffinity);
-                            }, out _);
+                        // Refresh selections
+                        selections = multiSelectionBroker.AllSelections
+                        // New selections are also ordered by position in ascending order
+                        .OrderBy(static selection => selection.Start.Position.Position)
+                        .ToArray();
+                        for (int i = 0; i < selections.Count; i++)
+                        {
+                            Selection selection = selections[i];
+                            Selection newSelection = newSelections[i];
+                            multiSelectionBroker.TryPerformActionOnSelection(
+                                selection,
+                                transformer =>
+                                {
+                                    // Move selection to new position
+                                    transformer.MoveTo(newSelection.AnchorPoint, newSelection.ActivePoint, newSelection.InsertionPoint, newSelection.InsertionPointAffinity);
+                                }, out _);
+                        }
                     }
+
+                    // Safe to history as single named edit
+                    undo.Complete();
                 }
             }
             else

diff --git a/src/Commands/ShiftUpCommand.cs b/src/Commands/ShiftUpCommand.cs
@@ -14,7 +14,7 @@ protected override async Task ExecuteAsync(OleMenuCmdEventArgs e)
                 return;
             }
 
-            ShiftDownCommand.Shift(docView, ShiftDirection.Up, false);
+            await ShiftDownCommand.ShiftAsync(docView, ShiftDirection.Up, false);
         }
     }
 }