Restore caret position when undoing a Complete Statement operation

Fixes #43400

src/EditorFeatures/CSharp/CompleteStatement/CompleteStatementCommandHandler.cs

@@ -13,6 +13,7 @@
 using Microsoft.CodeAnalysis.CSharp.Syntax;
 using Microsoft.CodeAnalysis.Editor.Implementation.AutomaticCompletion;
 using Microsoft.CodeAnalysis.Editor.Shared.Extensions;
+using Microsoft.CodeAnalysis.Editor.Shared.Utilities;
 using Microsoft.CodeAnalysis.Host.Mef;
 using Microsoft.CodeAnalysis.Internal.Log;
 using Microsoft.CodeAnalysis.LanguageServices;
@@ -23,6 +24,7 @@
 using Microsoft.VisualStudio.Language.Intellisense.AsyncCompletion;
 using Microsoft.VisualStudio.Text;
 using Microsoft.VisualStudio.Text.Editor.Commanding.Commands;
+using Microsoft.VisualStudio.Text.Operations;
 using Microsoft.VisualStudio.Utilities;
 
 namespace Microsoft.CodeAnalysis.Editor.CSharp.CompleteStatement
@@ -36,43 +38,67 @@ namespace Microsoft.CodeAnalysis.Editor.CSharp.CompleteStatement
     [Order(After = PredefinedCompletionNames.CompletionCommandHandler)]
     internal sealed class CompleteStatementCommandHandler : IChainedCommandHandler<TypeCharCommandArgs>
     {
+        private readonly ITextUndoHistoryRegistry _textUndoHistoryRegistry;
+        private readonly IEditorOperationsFactoryService _editorOperationsFactoryService;
+
         public CommandState GetCommandState(TypeCharCommandArgs args, Func<CommandState> nextCommandHandler) => nextCommandHandler();
 
         [ImportingConstructor]
         [Obsolete(MefConstruction.ImportingConstructorMessage, error: true)]
-        public CompleteStatementCommandHandler()
+        public CompleteStatementCommandHandler(ITextUndoHistoryRegistry textUndoHistoryRegistry, IEditorOperationsFactoryService editorOperationsFactoryService)
         {
+            _textUndoHistoryRegistry = textUndoHistoryRegistry;
+            _editorOperationsFactoryService = editorOperationsFactoryService;
         }
 
         public string DisplayName => CSharpEditorResources.Complete_statement_on_semicolon;
 
         public void ExecuteCommand(TypeCharCommandArgs args, Action nextCommandHandler, CommandExecutionContext executionContext)
         {
+            var history = _textUndoHistoryRegistry.GetHistory(args.SubjectBuffer);
+            var operations = _editorOperationsFactoryService.GetEditorOperations(args.TextView);
+
+            var willMoveSemicolon = BeforeExecuteCommand(speculative: true, args: args, executionContext: executionContext);
+            if (!willMoveSemicolon)
+            {
+                // Pass this on without altering the undo stack
+                nextCommandHandler();
+                return;
+            }
+
+            using var transaction = new HACK_TextUndoTransactionThatRollsBackProperly(history.CreateTransaction(nameof(CompleteStatementCommandHandler)));
+
+            operations.AddBeforeTextBufferChangePrimitive();
+
             // Determine where semicolon should be placed and move caret to location
-            BeforeExecuteCommand(args, executionContext);
+            BeforeExecuteCommand(speculative: false, args: args, executionContext: executionContext);
 
             // Insert the semicolon using next command handler
             nextCommandHandler();
+
+            operations.AddAfterTextBufferChangePrimitive();
+
+            transaction.Complete();
         }
 
-        private static void BeforeExecuteCommand(TypeCharCommandArgs args, CommandExecutionContext executionContext)
+        private static bool BeforeExecuteCommand(bool speculative, TypeCharCommandArgs args, CommandExecutionContext executionContext)
         {
             if (args.TypedChar != ';' || !args.TextView.Selection.IsEmpty)
             {
-                return;
+                return false;
             }
 
             var caretOpt = args.TextView.GetCaretPoint(args.SubjectBuffer);
             if (!caretOpt.HasValue)
             {
-                return;
+                return false;
             }
 
             var caret = caretOpt.Value;
             var document = caret.Snapshot.GetOpenDocumentInCurrentContextWithChanges();
             if (document == null)
             {
-                return;
+                return false;
             }
 
             var syntaxFacts = document.GetLanguageService<ISyntaxFactsService>();
@@ -81,10 +107,10 @@ private static void BeforeExecuteCommand(TypeCharCommandArgs args, CommandExecut
             var cancellationToken = executionContext.OperationContext.UserCancellationToken;
             if (!TryGetStartingNode(root, caret, out var currentNode, cancellationToken))
             {
-                return;
+                return false;
             }
 
-            MoveCaretToSemicolonPosition(args, document, root, caret, syntaxFacts, currentNode,
+            return MoveCaretToSemicolonPosition(speculative, args, document, root, originalCaret: caret, caret, syntaxFacts, currentNode,
                 isInsideDelimiters: false, cancellationToken);
         }
 
@@ -131,10 +157,12 @@ private static void BeforeExecuteCommand(TypeCharCommandArgs args, CommandExecut
             return true;
         }
 
-        private static void MoveCaretToSemicolonPosition(
+        private static bool MoveCaretToSemicolonPosition(
+            bool speculative,
             TypeCharCommandArgs args,
             Document document,
             SyntaxNode root,
+            SnapshotPoint originalCaret,
             SnapshotPoint caret,
             ISyntaxFactsService syntaxFacts,
             SyntaxNode currentNode,
@@ -145,7 +173,7 @@ private static void BeforeExecuteCommand(TypeCharCommandArgs args, CommandExecut
                 IsInAStringOrCharacter(currentNode, caret))
             {
                 // Don't complete statement.  Return without moving the caret.
-                return;
+                return false;
             }
 
             if (currentNode.IsKind(
@@ -161,7 +189,7 @@ private static void BeforeExecuteCommand(TypeCharCommandArgs args, CommandExecut
                 // make sure the closing delimiter exists
                 if (RequiredDelimiterIsMissing(currentNode))
                 {
-                    return;
+                    return false;
                 }
 
                 // set caret to just outside the delimited span and analyze again
@@ -169,39 +197,37 @@ private static void BeforeExecuteCommand(TypeCharCommandArgs args, CommandExecut
                 var newCaretPosition = currentNode.Span.End;
                 if (newCaretPosition == caret.Position)
                 {
-                    return;
+                    return false;
                 }
 
                 var newCaret = args.SubjectBuffer.CurrentSnapshot.GetPoint(newCaretPosition);
                 if (!TryGetStartingNode(root, newCaret, out currentNode, cancellationToken))
                 {
-                    return;
+                    return false;
                 }
 
-                MoveCaretToSemicolonPosition(args, document, root, newCaret, syntaxFacts, currentNode,
+                return MoveCaretToSemicolonPosition(speculative, args, document, root, originalCaret, newCaret, syntaxFacts, currentNode,
                     isInsideDelimiters: true, cancellationToken);
             }
             else if (currentNode.IsKind(SyntaxKind.DoStatement))
             {
                 if (IsInConditionOfDoStatement(currentNode, caret))
                 {
-                    MoveCaretToFinalPositionInStatement(currentNode, args, caret, true);
+                    return MoveCaretToFinalPositionInStatement(speculative, currentNode, args, originalCaret, caret, true);
                 }
-                return;
+                return false;
             }
             else if (syntaxFacts.IsStatement(currentNode)
                 || currentNode.IsKind(SyntaxKind.FieldDeclaration, SyntaxKind.DelegateDeclaration, SyntaxKind.ArrowExpressionClause))
             {
-                MoveCaretToFinalPositionInStatement(currentNode, args, caret, isInsideDelimiters);
-                return;
+                return MoveCaretToFinalPositionInStatement(speculative, currentNode, args, originalCaret, caret, isInsideDelimiters);
             }
             else
             {
                 // keep caret the same, but continue analyzing with the parent of the current node
                 currentNode = currentNode.Parent;
-                MoveCaretToSemicolonPosition(args, document, root, caret, syntaxFacts, currentNode,
+                return MoveCaretToSemicolonPosition(speculative, args, document, root, originalCaret, caret, syntaxFacts, currentNode,
                     isInsideDelimiters, cancellationToken);
-                return;
             }
         }
 
@@ -216,24 +242,33 @@ private static bool IsInConditionOfDoStatement(SyntaxNode currentNode, SnapshotP
             return (caret >= condition.Span.Start && caret <= condition.Span.End);
         }
 
-        private static void MoveCaretToFinalPositionInStatement(SyntaxNode statementNode, TypeCharCommandArgs args, SnapshotPoint caret, bool isInsideDelimiters)
+        private static bool MoveCaretToFinalPositionInStatement(bool speculative, SyntaxNode statementNode, TypeCharCommandArgs args, SnapshotPoint originalCaret, SnapshotPoint caret, bool isInsideDelimiters)
         {
             if (StatementClosingDelimiterIsMissing(statementNode))
             {
                 // Don't complete statement.  Return without moving the caret.
-                return;
+                return false;
             }
 
-            if (TryGetCaretPositionToMove(statementNode, caret, isInsideDelimiters, out var targetPosition))
+            if (TryGetCaretPositionToMove(statementNode, caret, isInsideDelimiters, out var targetPosition)
+                && targetPosition != originalCaret)
             {
+                if (speculative)
+                {
+                    // Return an indication that moving the caret is required, but don't actually move it
+                    return true;
+                }
+
                 Logger.Log(FunctionId.CommandHandler_CompleteStatement, KeyValueLogMessage.Create(LogType.UserAction, m =>
                 {
                     m[nameof(isInsideDelimiters)] = isInsideDelimiters;
                     m[nameof(statementNode)] = statementNode.Kind();
                 }));
 
-                args.TextView.TryMoveCaretToAndEnsureVisible(targetPosition);
+                return args.TextView.TryMoveCaretToAndEnsureVisible(targetPosition);
             }
+
+            return false;
         }
 
         private static bool TryGetCaretPositionToMove(SyntaxNode statementNode, SnapshotPoint caret, bool isInsideDelimiters, out SnapshotPoint targetPosition)

src/VisualStudio/IntegrationTest/IntegrationTests/CSharp/CSharpCompleteStatement.cs

+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using Microsoft.CodeAnalysis;
+using Microsoft.VisualStudio.IntegrationTest.Utilities;
+using Roslyn.Test.Utilities;
+using Xunit;
+
+namespace Roslyn.VisualStudio.IntegrationTests.CSharp
+{
+    [Collection(nameof(SharedIntegrationHostFixture))]
+    public class CSharpCompleteStatement : AbstractEditorTest
+    {
+        protected override string LanguageName => LanguageNames.CSharp;
+
+        public CSharpCompleteStatement(VisualStudioInstanceFactory instanceFactory)
+            : base(instanceFactory, nameof(CSharpCompleteStatement))
+        {
+        }
+
+        [WpfFact]
+        public void UndoRestoresCaretPosition1()
+        {
+            SetUpEditor(@"
+public class Test
+{
+    private object f;
+
+    public void Method()
+    {
+        f.ToString($$)
+    }
+}
+");
+
+            VisualStudio.Editor.SendKeys(';');
+            VisualStudio.Editor.Verify.CurrentLineText("f.ToString();$$", assertCaretPosition: true);
+
+            VisualStudio.Editor.Undo();
+            VisualStudio.Editor.Verify.CurrentLineText("f.ToString($$)", assertCaretPosition: true);
+        }
+
+        [WpfFact]
+        [WorkItem(43400, "https://github.com/dotnet/roslyn/issues/43400")]
+        public void UndoRestoresCaretPosition2()
+        {
+            SetUpEditor(@"
+public class Test
+{
+    private object f;
+
+    public void Method()
+    {
+        Method(condition ? whenTrue $$)
+    }
+}
+");
+
+            VisualStudio.Editor.SendKeys(';');
+            VisualStudio.Editor.Verify.CurrentLineText("Method(condition ? whenTrue );$$", assertCaretPosition: true);
+
+            VisualStudio.Editor.Undo();
+            VisualStudio.Editor.Verify.CurrentLineText("Method(condition ? whenTrue $$)", assertCaretPosition: true);
+        }
+
+        [WpfFact]
+        public void UndoRestoresFormatBeforeRestoringCaretPosition()
+        {
+            SetUpEditor(@"
+public class Test
+{
+    private object f;
+
+    public void Method()
+    {
+        f.ToString($$ )
+    }
+}
+");
+
+            VisualStudio.Editor.SendKeys(';');
+            VisualStudio.Editor.Verify.CurrentLineText("f.ToString();$$", assertCaretPosition: true);
+
+            VisualStudio.Editor.Undo();
+            VisualStudio.Editor.Verify.CurrentLineText("f.ToString( );$$", assertCaretPosition: true);
+
+            VisualStudio.Editor.Undo();
+            VisualStudio.Editor.Verify.CurrentLineText("f.ToString($$ )", assertCaretPosition: true);
+        }
+    }
+}