From 9f6d74d0d0566fd1cc59cf99cc6c3d0a3ac5cb0a Mon Sep 17 00:00:00 2001
From: Shahar Prish <shaharp@microsoft.com>
Date: Thu, 26 Sep 2024 15:31:35 +0300
Subject: [PATCH] Fix issue with CodeFix.

---
 .../CorrectObsolescenceCodeFixTests.cs        |  6 +----
 .../CorrectObsolescenceCodeFix.cs             | 23 +++++++++++--------
 2 files changed, 15 insertions(+), 14 deletions(-)

diff --git a/src/SubtleEngineering.Analyzers.Tests/Obsolescence/CorrectObsolescenceCodeFixTests.cs b/src/SubtleEngineering.Analyzers.Tests/Obsolescence/CorrectObsolescenceCodeFixTests.cs
index 1509ece..2ce5bd4 100644
--- a/src/SubtleEngineering.Analyzers.Tests/Obsolescence/CorrectObsolescenceCodeFixTests.cs
+++ b/src/SubtleEngineering.Analyzers.Tests/Obsolescence/CorrectObsolescenceCodeFixTests.cs
@@ -15,7 +15,6 @@ public async Task ObsoleteMethodWithoutEditorBrowsable_ShouldAddAttribute()
     {
         const string code = """
             using System;
-            using System.ComponentModel;
             
             class TestClass
             {
@@ -61,7 +60,6 @@ public async Task ObsoleteClassWithoutEditorBrowsable_ShouldAddAttribute()
     {
         const string code = """
             using System;
-            using System.ComponentModel;
             
             [Obsolete]
             public class {|#0:ObsoleteClass|}
@@ -72,7 +70,7 @@ public class {|#0:ObsoleteClass|}
         const string fixedCode = """
             using System;
             using System.ComponentModel;
-
+            
             [Obsolete]
             [EditorBrowsable(EditorBrowsableState.Never)]
             public class ObsoleteClass
@@ -96,7 +94,6 @@ public async Task ObsoletePropertyWithoutEditorBrowsable_ShouldAddAttribute()
     {
         const string code = """
             using System;
-            using System.ComponentModel;
             
             class TestClass
             {
@@ -133,7 +130,6 @@ public async Task ObsoleteEventWithoutEditorBrowsable_ShouldAddAttribute()
     {
         const string code = """
             using System;
-            using System.ComponentModel;
             
             class TestClass
             {
diff --git a/src/SubtleEngineering.Analyzers/Obsolescence/CorrectObsolescenceCodeFix.cs b/src/SubtleEngineering.Analyzers/Obsolescence/CorrectObsolescenceCodeFix.cs
index 620132f..321be75 100644
--- a/src/SubtleEngineering.Analyzers/Obsolescence/CorrectObsolescenceCodeFix.cs
+++ b/src/SubtleEngineering.Analyzers/Obsolescence/CorrectObsolescenceCodeFix.cs
@@ -63,6 +63,15 @@ private async Task<Document> AddEditorBrowsableAttributeAsync(Document document,
         {
             var root = await document.GetSyntaxRootAsync(cancellationToken).ConfigureAwait(false);
 
+            // Track the declaration node so we can find it after modifying the root
+            root = root.TrackNodes(declarationNode);
+
+            // Add using directive for System.ComponentModel if it's not already present
+            var rootWithUsing = AddUsingDirectiveIfMissing(root, "System.ComponentModel");
+
+            // Get the updated declaration node from the new root
+            var updatedDeclarationNode = rootWithUsing.GetCurrentNode(declarationNode);
+
             // Create the [EditorBrowsable(EditorBrowsableState.Never)] attribute
             var editorBrowsableAttribute = SyntaxFactory.Attribute(
                 SyntaxFactory.ParseName("EditorBrowsable"),
@@ -75,16 +84,12 @@ private async Task<Document> AddEditorBrowsableAttributeAsync(Document document,
                                 SyntaxFactory.IdentifierName("Never"))))));
 
             var newAttributeList = SyntaxFactory.AttributeList(SyntaxFactory.SingletonSeparatedList(editorBrowsableAttribute));
-                //.WithTrailingTrivia(SyntaxFactory.CarriageReturnLineFeed);
-
-            // Add using directive for System.ComponentModel if it's not already present
-            var rootWithUsing = AddUsingDirectiveIfMissing(root, "System.ComponentModel");
 
-            // Add the attribute to the declaration node
-            SyntaxNode newDeclarationNode = AddAttributeToDeclaration(declarationNode, newAttributeList);
+            // Add the attribute to the updated declaration node
+            var modifiedDeclarationNode = AddAttributeToDeclaration(updatedDeclarationNode, newAttributeList);
 
             // Replace the old node with the new node in the syntax tree
-            var newRoot = rootWithUsing.ReplaceNode(declarationNode, newDeclarationNode);
+            var newRoot = rootWithUsing.ReplaceNode(updatedDeclarationNode, modifiedDeclarationNode);
 
             // Return the new document
             return document.WithSyntaxRoot(newRoot);
@@ -97,8 +102,8 @@ private SyntaxNode AddUsingDirectiveIfMissing(SyntaxNode root, string namespaceN
                 var hasUsingDirective = compilationUnit.Usings.Any(u => u.Name.ToString() == namespaceName);
                 if (!hasUsingDirective)
                 {
-                    var usingDirective = SyntaxFactory.UsingDirective(SyntaxFactory.ParseName(namespaceName))
-                        .WithTrailingTrivia(SyntaxFactory.CarriageReturnLineFeed);
+                    var usingDirective = SyntaxFactory.UsingDirective(SyntaxFactory.ParseName(namespaceName));
+                        //.WithTrailingTrivia(SyntaxFactory.CarriageReturnLineFeed);
 
                     compilationUnit = compilationUnit.AddUsings(usingDirective);
                     return compilationUnit;