Improve undo stack titles for suggestions applied via light bulb

This change introduces a workspace service that can be used to register the display strings that should be displayed in VS's undo stack for each change applied using light bulb. The string displayed in the undo stack is currently identical to the one displayed in the light bulb menu.

Fixes #273

src/EditorFeatures/Core/EditorFeatures.csproj

@@ -653,6 +653,10 @@
     <Compile Include="Implementation\Workspaces\WorkspaceTaskSchedulerFactoryFactory.cs" />
     <Compile Include="InternalsVisibleTo.cs" />
     <Compile Include="IRefactorNotifyService.cs" />
+    <Compile Include="Undo\DefaultSourceTextUndoService.cs" />
+    <Compile Include="Undo\EditorSourceTextUndoService.cs" />
+    <Compile Include="Undo\ISourceTextUndoTransaction.cs" />
+    <Compile Include="Undo\ISourceTextUndoService.cs" />
     <Compile Include="MetadataAsSourceFile.cs" />
     <Compile Include="Implementation\Intellisense\IController.cs" />
     <Compile Include="Options\BraceCompletionOptions.cs" />

src/EditorFeatures/Core/Implementation/CodeActions/CodeActionEditHandlerService.cs

@@ -119,14 +119,41 @@ public void Apply(Workspace workspace, Document fromDocument, IEnumerable<CodeAc
                 if (updatedSolution == oldSolution)
                 {
                     updatedSolution = applyChanges.ChangedSolution;
+                    var projectChanges = updatedSolution.GetChanges(oldSolution).GetProjectChanges();
+                    var changedDocuments = projectChanges.SelectMany(pd => pd.GetChangedDocuments());
+                    var changedAddiionalDocuments = projectChanges.SelectMany(pd => pd.GetChangedAdditionalDocuments());
+                    var changedFiles = changedDocuments.Concat(changedAddiionalDocuments);
 
-                    // check whether it contains only 1 or 0 changed documents
-                    if (!updatedSolution.GetChanges(oldSolution).GetProjectChanges().SelectMany(pd => pd.GetChangedDocuments()).Skip(1).Any())
+                    // 0 file changes
+                    if (!changedFiles.Any())
                     {
                         operation.Apply(workspace, cancellationToken);
                         continue;
                     }
 
+                    // 1 file change
+                    SourceText text = null;
+                    if (!changedFiles.Skip(1).Any())
+                    {
+                        if (changedDocuments.Any())
+                        {
+                            text = oldSolution.GetDocument(changedDocuments.Single()).GetTextAsync(cancellationToken).WaitAndGetResult(cancellationToken);
+                        }
+                        else if (changedAddiionalDocuments.Any())
+                        {
+                            text = oldSolution.GetAdditionalDocument(changedAddiionalDocuments.Single()).GetTextAsync(cancellationToken).WaitAndGetResult(cancellationToken);
+                        }
+                    }
+
+                    if (text != null)
+                    {
+                        using (workspace.Services.GetService<ISourceTextUndoService>().RegisterUndoTransaction(text, title))
+                        {
+                            operation.Apply(workspace, cancellationToken);
+                            continue;
+                        }
+                    }
+
                     // multiple file changes
                     using (var undoTransaction = workspace.OpenGlobalUndoTransaction(title))
                     {

src/EditorFeatures/Core/Undo/DefaultSourceTextUndoService.cs

+// Copyright (c) Microsoft.  All Rights Reserved.  Licensed under the Apache License, Version 2.0.  See License.txt in the project root for license information.
+
+using System.Composition;
+using Microsoft.CodeAnalysis.Host.Mef;
+using Microsoft.CodeAnalysis.Text;
+using Microsoft.VisualStudio.Text;
+
+namespace Microsoft.CodeAnalysis.Editor.Undo
+{
+    [ExportWorkspaceService(typeof(ISourceTextUndoService), ServiceLayer.Default), Shared]
+    internal sealed class DefaultSourceTextUndoService : ISourceTextUndoService
+    {
+        public ISourceTextUndoTransaction RegisterUndoTransaction(SourceText sourceText, string description)
+        {
+            return null;
+        }
+
+        public bool BeginUndoTransaction(ITextSnapshot snapshot)
+        {
+            return false;
+        }
+
+        public bool EndUndoTransaction(ISourceTextUndoTransaction transaction)
+        {
+            return false;
+        }
+    }
+}

src/EditorFeatures/Core/Undo/EditorSourceTextUndoService.cs

+// Copyright (c) Microsoft.  All Rights Reserved.  Licensed under the Apache License, Version 2.0.  See License.txt in the project root for license information.
+
+using System.Collections.Generic;
+using System.Composition;
+using Microsoft.CodeAnalysis.Editor.Shared.Utilities;
+using Microsoft.CodeAnalysis.Host.Mef;
+using Microsoft.CodeAnalysis.Text;
+using Microsoft.VisualStudio.Text.Operations;
+using Microsoft.VisualStudio.Text;
+
+namespace Microsoft.CodeAnalysis.Editor.Undo
+{
+    [ExportWorkspaceService(typeof(ISourceTextUndoService), ServiceLayer.Editor), Shared]
+    internal sealed class EditorSourceTextUndoService : ISourceTextUndoService
+    {
+        private Dictionary<SourceText, SourceTextUndoTransaction> transactions = new Dictionary<SourceText, SourceTextUndoTransaction>();
+
+        private readonly ITextUndoHistoryRegistry _undoHistoryRegistry;
+
+        [ImportingConstructor]
+        public EditorSourceTextUndoService(ITextUndoHistoryRegistry undoHistoryRegistry)
+        {
+            _undoHistoryRegistry = undoHistoryRegistry;
+        }
+
+        public ISourceTextUndoTransaction RegisterUndoTransaction(SourceText sourceText, string description)
+        {
+            if (sourceText != null && !string.IsNullOrWhiteSpace(description))
+            {
+                var transaction = new SourceTextUndoTransaction(this, sourceText, description);
+                transactions.Add(sourceText, transaction);
+                return transaction;
+            }
+
+            return null;
+        }
+
+        public bool BeginUndoTransaction(ITextSnapshot snapshot)
+        {
+            SourceTextUndoTransaction transaction = null;
+            var sourceText = snapshot?.AsText();
+            if (sourceText != null)
+            {
+                transactions.TryGetValue(sourceText, out transaction);
+                if (transaction != null)
+                {
+                    return transaction.Begin(_undoHistoryRegistry?.GetHistory(snapshot.TextBuffer));
+                }
+            }
+
+            return false;
+        }
+
+        public bool EndUndoTransaction(ISourceTextUndoTransaction transaction)
+        {
+            if (transaction != null && transactions.ContainsKey(transaction.SourceText))
+            {
+                transactions.Remove(transaction.SourceText);
+                return true;
+            }
+
+            return false;
+        }
+
+        private sealed class SourceTextUndoTransaction : ISourceTextUndoTransaction
+        {
+            private readonly ISourceTextUndoService _service;
+            public SourceText SourceText { get; }
+            public string Description { get; }
+
+            private ITextUndoTransaction _transaction;
+
+            public SourceTextUndoTransaction(ISourceTextUndoService service, SourceText sourceText, string description)
+            {
+                _service = service;
+                SourceText = sourceText;
+                Description = description;
+            }
+
+            internal bool Begin(ITextUndoHistory undoHistory)
+            {
+                if (undoHistory != null)
+                {
+                    _transaction = new HACK_TextUndoTransactionThatRollsBackProperly(undoHistory.CreateTransaction(Description));
+                    return true;
+                }
+
+                return false;
+            }
+
+            public void Dispose()
+            {
+                if (_transaction != null)
+                {
+                    _transaction.Complete();
+                }
+
+                _service.EndUndoTransaction(this);
+            }
+        }
+    }
+}

src/EditorFeatures/Core/Undo/ISourceTextUndoService.cs

+// Copyright (c) Microsoft.  All Rights Reserved.  Licensed under the Apache License, Version 2.0.  See License.txt in the project root for license information.
+
+using Microsoft.CodeAnalysis.Host;
+using Microsoft.CodeAnalysis.Text;
+using Microsoft.VisualStudio.Text;
+
+namespace Microsoft.CodeAnalysis.Editor.Undo
+{
+    /// <summary>
+    /// A service that allows consumers to register undo transactions for a supplied
+    /// <see cref="SourceText"/> with a supplied description. The description is the
+    /// display string by which the IDE's undo stack UI will subsequently refer to the transaction.
+    /// </summary>
+    internal interface ISourceTextUndoService : IWorkspaceService
+    {
+        /// <summary>
+        /// Registers undo transaction for the supplied <see cref="SourceText"/>.
+        /// </summary>
+        /// <param name="sourceText">The <see cref="SourceText"/> for which undo transaction is being registered.</param>
+        /// <param name="description">The display string by which the IDE's undo stack UI will subsequently refer to the transaction.</param>
+        ISourceTextUndoTransaction RegisterUndoTransaction(SourceText sourceText, string description);
+
+        /// <summary>
+        /// Starts previously registered undo transaction for the supplied <see cref="ITextSnapshot"/> (if any).
+        /// </summary>
+        /// <param name="snapshot">The <see cref="ITextSnapshot"/> for the <see cref="SourceText"/> for undo transaction being started.</param>
+        /// <remarks>
+        /// This method will handle the translation from <see cref="ITextSnapshot"/> to <see cref="SourceText"/>
+        /// and update the IDE's undo stack UI with the transaction's previously registered description string.
+        /// </remarks>
+        bool BeginUndoTransaction(ITextSnapshot snapshot);
+
+        /// <summary>
+        /// Completes and deletes the supplied undo transaction.
+        /// </summary>
+        /// <param name="transaction">The undo transaction that is being ended.</param>
+        bool EndUndoTransaction(ISourceTextUndoTransaction transaction);
+    }
+}

src/EditorFeatures/Core/Undo/ISourceTextUndoTransaction.cs

+// Copyright (c) Microsoft.  All Rights Reserved.  Licensed under the Apache License, Version 2.0.  See License.txt in the project root for license information.
+
+using System;
+using Microsoft.CodeAnalysis.Text;
+
+namespace Microsoft.CodeAnalysis.Editor.Undo
+{
+    /// <summary>
+    /// Represents undo transaction for a <see cref="Microsoft.CodeAnalysis.Text.SourceText"/>
+    /// with a display string by which the IDE's undo stack UI refers to the transaction.
+    /// </summary>
+    internal interface ISourceTextUndoTransaction : IDisposable
+    {
+        /// <summary>
+        /// The <see cref="Microsoft.CodeAnalysis.Text.SourceText"/> for this undo transaction.
+        /// </summary>
+        SourceText SourceText { get; }
+
+        /// <summary>
+        /// The display string by which the IDE's undo stack UI refers to the transaction.
+        /// </summary>
+        string Description { get; }
+    }
+}

src/VisualStudio/Core/Def/Implementation/ProjectSystem/DocumentProvider.StandardTextDocument.cs

@@ -4,9 +4,9 @@
 using System.Collections.Generic;
 using System.Diagnostics;
 using System.IO;
-using System.Text;
 using Microsoft.CodeAnalysis;
 using Microsoft.CodeAnalysis.Editor.Shared.Utilities;
+using Microsoft.CodeAnalysis.Editor.Undo;
 using Microsoft.CodeAnalysis.Text;
 using Microsoft.VisualStudio.Shell.Interop;
 using Microsoft.VisualStudio.Text;
@@ -217,9 +217,17 @@ private static void UpdateText(SourceText newText, ITextBuffer buffer, EditOptio
             {
                 using (var edit = buffer.CreateEdit(options, reiteratedVersionNumber: null, editTag: null))
                 {
-                    var oldText = buffer.CurrentSnapshot.AsText();
+                    var oldSnapshot = buffer.CurrentSnapshot;
+                    var oldText = oldSnapshot.AsText();
                     var changes = newText.GetTextChanges(oldText);
 
+                    Workspace workspace = null;
+                    if (Workspace.TryGetWorkspace(oldText.Container, out workspace))
+                    {
+                        var undoService = workspace.Services.GetService<ISourceTextUndoService>();
+                        undoService.BeginUndoTransaction(oldSnapshot);
+                    }
+
                     foreach (var change in changes)
                     {
                         edit.Replace(change.Span.Start, change.Span.Length, change.NewText);