Fix a few issues related to trivia generation for add/remove suppression fixes...

Fix a few issues related to trivia generation for add/remove suppression fixes from the error list. Unit tests have been added for the same.

src/EditorFeatures/CSharpTest/Diagnostics/Suppression/RemoveSuppressionTests.cs

@@ -212,6 +212,39 @@ int Method()
 }");
             }
 
+            [Fact, Trait(Traits.Feature, Traits.Features.CodeActionsSuppression)]
+            public void TestRemoveLocalAttributeSuppression2()
+            {
+                Test(
+        $@"
+using System;
+
+class Class1
+{{
+    [System.Diagnostics.CodeAnalysis.SuppressMessage(""InfoDiagnostic"", ""InfoDiagnostic:InfoDiagnostic"", Justification = ""{FeaturesResources.SuppressionPendingJustification}"")]           
+    [|class Class2|]
+    {{
+        int Method()
+        {{
+            int x = 0;
+        }}
+    }}
+}}",
+        @"
+using System;
+
+class Class1
+{
+    class Class2
+    {
+        int Method()
+        {
+            int x = 0;
+        }
+    }
+}");
+            }
+
             [Fact, Trait(Traits.Feature, Traits.Features.CodeActionsSuppression)]
             public void TestRemoveGlobalAttributeSuppression()
             {

src/EditorFeatures/CSharpTest/Diagnostics/Suppression/SuppressionTest_FixAllTests.cs

@@ -40,6 +40,7 @@ int Method()
 class Class2
 {
 }
+class Class3 { }
         </Document>
         <Document>
 class Class3
@@ -85,6 +86,9 @@ class Class2
 #pragma warning restore InfoDiagnostic // InfoDiagnostic Title
 {
 }
+#pragma warning disable InfoDiagnostic // InfoDiagnostic Title
+class Class3 { }
+#pragma warning restore InfoDiagnostic // InfoDiagnostic Title
         </Document>
         <Document>
 class Class3

src/Features/Core/Portable/CodeFixes/Suppression/AbstractSuppressionCodeFixProvider.PragmaHelpers.cs

@@ -169,6 +169,75 @@ internal static SyntaxToken GetNewEndTokenWithAddedPragma(SyntaxToken endToken,
                     return endToken.WithTrailingTrivia(trivia.InsertRange(index, pragmaTrivia));
                 };
             }
+
+            internal static void NormalizeTriviaOnTokens(AbstractSuppressionCodeFixProvider fixer, ref Document document, ref SuppressionTargetInfo suppressionTargetInfo)
+            {
+                // For pragma suppression fixes, we need to normalize the leading trivia on start token to account for
+                // the trailing trivia on its previous token (and similarly normalize trailing trivia for end token).
+
+                var startToken = suppressionTargetInfo.StartToken;
+                var endToken = suppressionTargetInfo.EndToken;
+                var nodeWithTokens = suppressionTargetInfo.NodeWithTokens;
+                var startAndEndTokensAreSame = startToken == endToken;
+                var isEndTokenEOF = fixer.IsEndOfFileToken(endToken);
+
+                var previousOfStart = startToken.GetPreviousToken(includeZeroWidth: true);
+                var nextOfEnd = !isEndTokenEOF ? endToken.GetNextToken(includeZeroWidth: true) : default(SyntaxToken);
+                if (!previousOfStart.HasTrailingTrivia && !nextOfEnd.HasLeadingTrivia)
+                {
+                    return;
+                }
+
+                var root = nodeWithTokens.SyntaxTree.GetRoot();
+                var spanEnd = !isEndTokenEOF ? nextOfEnd.FullSpan.End : endToken.FullSpan.End;
+                var subtreeRoot = root.FindNode(new TextSpan(previousOfStart.FullSpan.Start, spanEnd - previousOfStart.FullSpan.Start));
+
+                var currentStartToken = startToken;
+                var currentEndToken = endToken;
+                var newStartToken = startToken.WithLeadingTrivia(previousOfStart.TrailingTrivia.Concat(startToken.LeadingTrivia));
+
+                SyntaxToken newEndToken = currentEndToken;
+                if (startAndEndTokensAreSame)
+                {
+                    newEndToken = newStartToken;
+                }
+
+                newEndToken = newEndToken.WithTrailingTrivia(endToken.TrailingTrivia.Concat(nextOfEnd.LeadingTrivia));
+
+                var newPreviousOfStart = previousOfStart.WithTrailingTrivia();
+                var newNextOfEnd = nextOfEnd.WithLeadingTrivia();
+
+                var newSubtreeRoot = subtreeRoot.ReplaceTokens(new[] { startToken, previousOfStart, endToken, nextOfEnd },
+                    (o, n) =>
+                    {
+                        if (o == currentStartToken)
+                        {
+                            return startAndEndTokensAreSame ? newEndToken : newStartToken;
+                        }
+                        else if (o == previousOfStart)
+                        {
+                            return newPreviousOfStart;
+                        }
+                        else if (o == currentEndToken)
+                        {
+                            return newEndToken;
+                        }
+                        else if (o == nextOfEnd)
+                        {
+                            return newNextOfEnd;
+                        }
+                        else
+                        {
+                            return n;
+                        }
+                    });
+
+                root = root.ReplaceNode(subtreeRoot, newSubtreeRoot);
+                document = document.WithSyntaxRoot(root);
+                suppressionTargetInfo.StartToken = root.FindToken(startToken.SpanStart);
+                suppressionTargetInfo.EndToken = root.FindToken(endToken.SpanStart);
+                suppressionTargetInfo.NodeWithTokens = fixer.GetNodeWithTokens(suppressionTargetInfo.StartToken, suppressionTargetInfo.EndToken, root);
+            }
         }
     }
 }

src/Features/Core/Portable/CodeFixes/Suppression/AbstractSuppressionCodeFixProvider.PragmaWarningCodeAction.cs

@@ -15,7 +15,20 @@ internal sealed class PragmaWarningCodeAction : AbstractSuppressionCodeAction, I
             private readonly Diagnostic _diagnostic;
             private readonly bool _forFixMultipleContext;
 
-            internal PragmaWarningCodeAction(
+            public static PragmaWarningCodeAction Create(
+                    SuppressionTargetInfo suppressionTargetInfo,
+                    Document document,
+                    Diagnostic diagnostic,
+                    AbstractSuppressionCodeFixProvider fixer)
+            {
+                // We need to normalize the leading trivia on start token to account for
+                // the trailing trivia on its previous token (and similarly normalize trailing trivia for end token).
+                PragmaHelpers.NormalizeTriviaOnTokens(fixer, ref document, ref suppressionTargetInfo);
+
+                return new PragmaWarningCodeAction(suppressionTargetInfo, document, diagnostic, fixer);
+            }
+
+            private PragmaWarningCodeAction(
                 SuppressionTargetInfo suppressionTargetInfo,
                 Document document,
                 Diagnostic diagnostic,

src/Features/Core/Portable/CodeFixes/Suppression/AbstractSuppressionCodeFixProvider.RemoveSuppressionCodeAction.BatchFixer.cs

@@ -130,7 +130,7 @@ public override async Task<CodeAction> TryGetMergedFixAsync(IEnumerable<CodeActi
                             var attributesToRemove = await GetAttributeNodesToFixAsync(attributeRemoveFixesForTree, cancellationToken).ConfigureAwait(false);
                             var document = oldSolution.GetDocument(tree);
                             var root = await document.GetSyntaxRootAsync(cancellationToken).ConfigureAwait(false);
-                            var newRoot = root.RemoveNodes(attributesToRemove, SyntaxRemoveOptions.KeepLeadingTrivia);
+                            var newRoot = root.RemoveNodes(attributesToRemove, SyntaxRemoveOptions.KeepLeadingTrivia | SyntaxRemoveOptions.AddElasticMarker);
                             currentSolution = currentSolution.WithDocumentSyntaxRoot(document.Id, newRoot);
                         }
 

src/Features/Core/Portable/CodeFixes/Suppression/AbstractSuppressionCodeFixProvider.RemoveSuppressionCodeAction_Pragma.cs

@@ -32,7 +32,7 @@ private class PragmaRemoveAction : RemoveSuppressionCodeAction, IPragmaBasedCode
                 {
                     // We need to normalize the leading trivia on start token to account for
                     // the trailing trivia on its previous token (and similarly normalize trailing trivia for end token).
-                    NormalizeTriviaOnTokens(fixer, ref document, ref suppressionTargetInfo);
+                    PragmaHelpers.NormalizeTriviaOnTokens(fixer, ref document, ref suppressionTargetInfo);
 
                     return new PragmaRemoveAction(suppressionTargetInfo, document, diagnostic, fixer);
                 }
@@ -228,70 +228,6 @@ private SyntaxNode FormatNode(SyntaxNode node)
                 {
                     return Formatter.Format(node, _document.Project.Solution.Workspace);
                 }
-
-                private static void NormalizeTriviaOnTokens(AbstractSuppressionCodeFixProvider fixer, ref Document document, ref SuppressionTargetInfo suppressionTargetInfo)
-                {
-                    var startToken = suppressionTargetInfo.StartToken;
-                    var endToken = suppressionTargetInfo.EndToken;
-                    var nodeWithTokens = suppressionTargetInfo.NodeWithTokens;
-                    var startAndEndTokensAreSame = startToken == endToken;
-
-                    var previousOfStart = startToken.GetPreviousToken();
-                    var nextOfEnd = endToken.GetNextToken();
-                    if (!previousOfStart.HasTrailingTrivia && !nextOfEnd.HasLeadingTrivia)
-                    {
-                        return;
-                    }
-
-                    var root = nodeWithTokens.SyntaxTree.GetRoot();
-                    var subtreeRoot = root.FindNode(new TextSpan(previousOfStart.FullSpan.Start, nextOfEnd.FullSpan.End - previousOfStart.FullSpan.Start));
-
-                    var currentStartToken = startToken;
-                    var currentEndToken = endToken;
-                    var newStartToken = startToken.WithLeadingTrivia(previousOfStart.TrailingTrivia.Concat(startToken.LeadingTrivia));
-
-                    SyntaxToken newEndToken = currentEndToken;
-                    if (startAndEndTokensAreSame)
-                    {
-                        newEndToken = newStartToken;
-                    }
-
-                    newEndToken = newEndToken.WithTrailingTrivia(endToken.TrailingTrivia.Concat(nextOfEnd.LeadingTrivia));
-
-                    var newPreviousOfStart = previousOfStart.WithTrailingTrivia();
-                    var newNextOfEnd = nextOfEnd.WithLeadingTrivia();
-
-                    var newSubtreeRoot = subtreeRoot.ReplaceTokens(new[] { startToken, previousOfStart, endToken, nextOfEnd },
-                        (o, n) =>
-                        {
-                            if (o == currentStartToken)
-                            {
-                                return newStartToken;
-                            }
-                            else if (o == previousOfStart)
-                            {
-                                return newPreviousOfStart;
-                            }
-                            else if (o == currentEndToken)
-                            {
-                                return newEndToken;
-                            }
-                            else if (o == nextOfEnd)
-                            {
-                                return newNextOfEnd;
-                            }
-                            else
-                            {
-                                return n;
-                            }
-                        });
-
-                    root = root.ReplaceNode(subtreeRoot, newSubtreeRoot);
-                    document = document.WithSyntaxRoot(root);
-                    suppressionTargetInfo.StartToken = root.FindToken(startToken.SpanStart);
-                    suppressionTargetInfo.EndToken = root.FindToken(endToken.SpanStart);
-                    suppressionTargetInfo.NodeWithTokens = fixer.GetNodeWithTokens(suppressionTargetInfo.StartToken, suppressionTargetInfo.EndToken, root);
-                }
             }
         }
     }

src/Features/Core/Portable/CodeFixes/Suppression/AbstractSuppressionCodeFixProvider.cs

@@ -133,7 +133,7 @@ private async Task<IEnumerable<CodeFix>> GetSuppressionsAsync(Document documentO
                     if (diagnostic.Location.IsInSource && documentOpt != null)
                     {
                         // pragma warning disable.
-                        nestedActions.Add(new PragmaWarningCodeAction(suppressionTargetInfo, documentOpt, diagnostic, this));
+                        nestedActions.Add(PragmaWarningCodeAction.Create(suppressionTargetInfo, documentOpt, diagnostic, this));
                     }
 
                     // SuppressMessageAttribute suppression is not supported for compiler diagnostics.

src/VisualStudio/Core/Def/Implementation/TableDataSource/Suppression/VisualStudioSuppressionFixService.cs

@@ -237,7 +237,7 @@ private bool ApplySuppressionFix(Func<Project, bool> shouldFixInProject, bool se
                     newSolution,
                     fixAllPreviewChangesTitle: title,
                     fixAllTopLevelHeader: title,
-                    languageOpt: languages.SingleOrDefault(),
+                    languageOpt: languages.Count == 1 ? languages.Single() : null,
                     workspace: _workspace);
                 if (newSolution == null)
                 {