From b1c092c52a093ca10e76ba2bae7487737a923030 Mon Sep 17 00:00:00 2001
From: Daniel Cazzulino <daniel@cazzulino.com>
Date: Tue, 10 Dec 2024 10:43:20 -0300
Subject: [PATCH] Preserve trivia when applying code templates

We should not lose valuable whitespaces and even more importantly, code comments.
---
 src/StructId.Analyzer/CodeTemplate.cs   | 39 ++++++++++++++++---------
 src/StructId.Tests/CodeTemplateTests.cs | 34 +++++++++++++++++++--
 2 files changed, 57 insertions(+), 16 deletions(-)

diff --git a/src/StructId.Analyzer/CodeTemplate.cs b/src/StructId.Analyzer/CodeTemplate.cs
index 42b1bd9..ef61a25 100644
--- a/src/StructId.Analyzer/CodeTemplate.cs
+++ b/src/StructId.Analyzer/CodeTemplate.cs
@@ -1,4 +1,5 @@
-using System.Collections.Generic;
+using System;
+using System.Collections.Generic;
 using System.Linq;
 using Microsoft.CodeAnalysis;
 using Microsoft.CodeAnalysis.CSharp;
@@ -115,16 +116,35 @@ class TemplateRewriter(string tself, string tid) : CSharpSyntaxRewriter
                         .WithOpenParenToken(parameters.OpenParenToken.WithoutTrivia()));
                 else
                     node = node.WithParameterList(null);
+
+                node = node.WithIdentifier(node.Identifier.WithTrailingTrivia(parameters.CloseParenToken.TrailingTrivia));
             }
 
             var visited = (RecordDeclarationSyntax)base.VisitRecordDeclaration(node)!;
+            var trivia = TriviaList();
+            // Rather than removing the empty attribute lists via the VisitAttributeList method, we do it here
+            // so we can preserve the trivia.
+            foreach (var list in visited.AttributeLists)
+            {
+                if (list.Attributes.Count == 0)
+                {
+                    trivia = trivia.AddRange(list.GetLeadingTrivia());
+                    visited = visited.RemoveNode(list, SyntaxRemoveOptions.KeepNoTrivia)!;
+                }
+            }
+
+            if (trivia.Count > 0)
+                visited = visited.WithLeadingTrivia(trivia);
 
             // remove file modifier from type declarations
             if (visited.Modifiers.FirstOrDefault(x => x.IsKind(SyntaxKind.FileKeyword)) is { } file)
-                // Preserve trivia, i.e. newline from original file modifier
-                return visited
-                    .WithLeadingTrivia(file.LeadingTrivia)
-                    .WithModifiers(visited.Modifiers.Remove(file));
+            {
+                // Preserve trivia, i.e. newline from original file modifier, as well as potentially 
+                // other trivia we might have added from removed attribute lists
+                visited = visited
+                    .WithModifiers(visited.Modifiers.Remove(file))
+                    .WithLeadingTrivia(file.LeadingTrivia);
+            }
 
             return visited;
         }
@@ -148,15 +168,6 @@ class TemplateRewriter(string tself, string tid) : CSharpSyntaxRewriter
             return base.VisitClassDeclaration(node);
         }
 
-        public override SyntaxNode? VisitAttributeList(AttributeListSyntax node)
-        {
-            node = (AttributeListSyntax)base.VisitAttributeList(node)!;
-            if (node.Attributes.Count == 0)
-                return null;
-
-            return node;
-        }
-
         public override SyntaxNode? VisitAttribute(AttributeSyntax node)
         {
             if (node.IsStructIdTemplate())
diff --git a/src/StructId.Tests/CodeTemplateTests.cs b/src/StructId.Tests/CodeTemplateTests.cs
index c6dc0df..e317660 100644
--- a/src/StructId.Tests/CodeTemplateTests.cs
+++ b/src/StructId.Tests/CodeTemplateTests.cs
@@ -122,8 +122,6 @@ file partial record struct TSelf
 
         var applied = CodeTemplate.Apply(template, "Foo", "string", normalizeWhitespace: true);
 
-        output.WriteLine(applied);
-
         Assert.Equal(
             CodeTemplate.Parse(
             """
@@ -136,4 +134,36 @@ partial record struct Foo
             """).NormalizeWhitespace().ToFullString().Trim().ReplaceLineEndings(),
             applied.ReplaceLineEndings());
     }
+
+    [Fact]
+    public void PreservesTrivia()
+    {
+        var template =
+            """
+            using System;
+
+            // Test
+            [TStructId]
+            file partial record struct TSelf(Ulid Value)
+            {
+                public static TSelf New() => new(Ulid.NewUlid());
+            }
+            """;
+
+        var applied = CodeTemplate.Apply(template, "ItemId", "Ulid");
+
+        Assert.Equal(
+            """
+            using System;
+            
+            // Test
+            partial record struct ItemId
+            {
+                public static ItemId New() => new(Ulid.NewUlid());
+            }
+            """,
+            applied);
+
+        output.WriteLine(applied);
+    }
 }