Get the appropriate text for rolling back.

src/Features/Core/Portable/AddImport/CodeActions/PackageReference.InstallPackageAndAddImportCodeAction.cs

@@ -6,6 +6,7 @@
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.CodeActions;
 using Microsoft.CodeAnalysis.Shared.Utilities;
+using Microsoft.CodeAnalysis.Text;
 using Roslyn.Utilities;
 
 namespace Microsoft.CodeAnalysis.CodeFixes.AddImport
@@ -18,7 +19,7 @@ private class InstallPackageAndAddImportCodeAction : CodeAction
             {
                 private readonly string _title;
                 private readonly CodeActionPriority _priority;
-                private readonly AsyncLazy<ValueTuple<ApplyChangesOperation, InstallNugetPackageOperation>> _getOperations;
+                private readonly AsyncLazy<ValueTuple<Document, Document, InstallNugetPackageOperation>> _documentsAndInstallOperation;
 
                 public override string Title => _title;
                 public override string EquivalenceKey => _title;
@@ -26,11 +27,11 @@ private class InstallPackageAndAddImportCodeAction : CodeAction
 
                 public InstallPackageAndAddImportCodeAction(
                     string title, CodeActionPriority priority,
-                    AsyncLazy<ValueTuple<ApplyChangesOperation, InstallNugetPackageOperation>> getOperations)
+                    AsyncLazy<ValueTuple<Document, Document, InstallNugetPackageOperation>> documentsAndInstallOperation)
                 {
                     _title = title;
                     _priority = priority;
-                    _getOperations = getOperations;
+                    _documentsAndInstallOperation = documentsAndInstallOperation;
                 }
 
                 /// <summary>
@@ -41,8 +42,14 @@ private class InstallPackageAndAddImportCodeAction : CodeAction
                 /// </summary>
                 protected override async Task<IEnumerable<CodeActionOperation>> ComputePreviewOperationsAsync(CancellationToken cancellationToken)
                 {
-                    var operations = await _getOperations.GetValueAsync(cancellationToken).ConfigureAwait(false);
-                    return ImmutableArray.Create<CodeActionOperation>(operations.Item1, operations.Item2);
+                    var newDocumentAndInstallOperation = await _documentsAndInstallOperation.GetValueAsync(cancellationToken).ConfigureAwait(false);
+                    var solutionChangeAction = new SolutionChangeAction(
+                        "", c => Task.FromResult(newDocumentAndInstallOperation.Item2.Project.Solution));
+
+                    var result = ArrayBuilder<CodeActionOperation>.GetInstance();
+                    result.AddRange(await solutionChangeAction.GetPreviewOperationsAsync(cancellationToken).ConfigureAwait(false));
+                    result.Add(newDocumentAndInstallOperation.Item3);
+                    return result.ToImmutableAndFree();
                 }
 
                 /// <summary>
@@ -53,19 +60,34 @@ protected override async Task<IEnumerable<CodeActionOperation>> ComputePreviewOp
                 protected override async Task<IEnumerable<CodeActionOperation>> ComputeOperationsAsync(
                     CancellationToken cancellationToken)
                 {
-                    var operations = await _getOperations.GetValueAsync(cancellationToken).ConfigureAwait(false);
-                    return ImmutableArray.Create<CodeActionOperation>(new CompoundOperation(operations.Item1, operations.Item2));
+                    var documentsAndInstallOperation = await _documentsAndInstallOperation.GetValueAsync(cancellationToken).ConfigureAwait(false);
+                    var oldDocument = documentsAndInstallOperation.Item1;
+                    var newDocument = documentsAndInstallOperation.Item2;
+
+                    var oldText = await oldDocument.GetTextAsync(cancellationToken).ConfigureAwait(false);
+                    var newText = await newDocument.GetTextAsync(cancellationToken).ConfigureAwait(false);
+
+                    return ImmutableArray.Create<CodeActionOperation>(new CompoundOperation(
+                        oldDocument.Id, oldText, newText, documentsAndInstallOperation.Item3));
                 }
             }
 
             private class CompoundOperation : CodeActionOperation
             {
-                private readonly ApplyChangesOperation _applyChanges;
+                private readonly DocumentId _changedDocumentId;
+                private readonly SourceText _oldText;
+                private readonly SourceText _newText;
                 private readonly InstallNugetPackageOperation _installNugetPackage;
 
-                public CompoundOperation(ApplyChangesOperation item1, InstallNugetPackageOperation item2)
+                public CompoundOperation(
+                    DocumentId changedDocumentId, 
+                    SourceText oldText,
+                    SourceText newText, 
+                    InstallNugetPackageOperation item2)
                 {
-                    _applyChanges = item1;
+                    _changedDocumentId = changedDocumentId;
+                    _oldText = oldText;
+                    _newText = newText;
                     _installNugetPackage = item2;
                 }
 
@@ -74,10 +96,11 @@ public CompoundOperation(ApplyChangesOperation item1, InstallNugetPackageOperati
 
                 internal override bool TryApply(Workspace workspace, IProgressTracker progressTracker, CancellationToken cancellationToken)
                 {
-                    var oldSolution = workspace.CurrentSolution;
+                    var newSolution = workspace.CurrentSolution.WithDocumentText(
+                        _changedDocumentId, _newText);
 
                     // First make the changes to add the import to the document.
-                    if (_applyChanges.TryApply(workspace, progressTracker, cancellationToken))
+                    if (workspace.TryApplyChanges(newSolution, progressTracker))
                     {
                         if (_installNugetPackage.TryApply(workspace, progressTracker, cancellationToken))
                         {
@@ -85,7 +108,9 @@ internal override bool TryApply(Workspace workspace, IProgressTracker progressTr
                         }
 
                         // Installing the nuget package failed.  Roll back the workspace.
-                        workspace.TryApplyChanges(oldSolution, progressTracker);
+                        var rolledBackSolution = workspace.CurrentSolution.WithDocumentText(
+                            _changedDocumentId, _oldText);
+                        workspace.TryApplyChanges(rolledBackSolution, progressTracker);
                     }
 
                     return false;

src/Features/Core/Portable/AddImport/CodeActions/PackageReference.ParentCodeAction.cs

@@ -78,7 +78,7 @@ private class ParentCodeAction : CodeAction
                             ? string.Format(FeaturesResources.Use_local_version_0, versionOpt)
                             : string.Format(FeaturesResources.Install_version_0, versionOpt);
 
-                    var getOperations = new AsyncLazy<ValueTuple<ApplyChangesOperation, InstallNugetPackageOperation>>(
+                    var getOperations = new AsyncLazy<ValueTuple<Document, Document, InstallNugetPackageOperation>>(
                         c => GetOperationsAsync(versionOpt, isLocal, document, node, placeSystemNamespaceFirst, c),
                         cacheResult: true);
 
@@ -87,7 +87,7 @@ private class ParentCodeAction : CodeAction
                         title, CodeActionPriority.Low, getOperations);
                 }
 
-                private async Task<ValueTuple<ApplyChangesOperation, InstallNugetPackageOperation>> GetOperationsAsync(
+                private async Task<ValueTuple<Document, Document, InstallNugetPackageOperation>> GetOperationsAsync(
                     string versionOpt, 
                     bool isLocal,
                     Document document, 
@@ -95,23 +95,23 @@ private class ParentCodeAction : CodeAction
                     bool placeSystemNamespaceFirst, 
                     CancellationToken cancellationToken)
                 {
+                    var oldDocument = document;
                     _reference.ReplaceNameNode(ref node, ref document, cancellationToken);
 
                     var newDocument = await _reference.provider.AddImportAsync(
                         node, _reference.SearchResult.NameParts, document, placeSystemNamespaceFirst, cancellationToken).ConfigureAwait(false);
-                    var newSolution = newDocument.Project.Solution;
 
-                    // Create a dummy code action here so that we go through the codepath
-                    // where the solution is 'preprocessed' (i.e. formatting/simplification/etc.
-                    // is run). 
-                    var codeAction = new SolutionChangeAction("", c => Task.FromResult(newSolution));
-                    var codeActionOperations =  await codeAction.GetOperationsAsync(cancellationToken).ConfigureAwait(false);
+                    // We're going to be manually applying this new document to the workspace
+                    // (so we can roll it back ourselves if installing the nuget package fails).
+                    // As such, we need to do the postprocessing ourselves of tihs document to 
+                    // ensure things like formatting/simplification happen to it.
+                    newDocument = await this.PostProcessChangesAsync(
+                        newDocument, cancellationToken).ConfigureAwait(false);
 
-                    var operation1 = (ApplyChangesOperation)codeActionOperations.Single();
-                    var operation2 = new InstallNugetPackageOperation(
+                    var installOperation = new InstallNugetPackageOperation(
                         _reference._installerService, document, _reference._source, _reference._packageName, versionOpt, isLocal);
 
-                    return ValueTuple.Create(operation1, operation2);
+                    return ValueTuple.Create(oldDocument, newDocument, installOperation);
                 }
             }
         }

src/Workspaces/Core/Portable/CodeActions/CodeAction.cs

@@ -259,7 +259,7 @@ protected async Task<Solution> PostProcessChangesAsync(Solution changedSolution,
         /// <param name="document">The document changed by the <see cref="CodeAction"/>.</param>
         /// <param name="cancellationToken">A cancellation token.</param>
         /// <returns>A document with the post processing changes applied.</returns>
-        protected virtual async Task<Document> PostProcessChangesAsync(Document document, CancellationToken cancellationToken)
+        protected async virtual Task<Document> PostProcessChangesAsync(Document document, CancellationToken cancellationToken)
         {
             if (document.SupportsSyntaxTree)
             {