Remove unused usings when internalizing types

packages/http-client-csharp/generator/Microsoft.TypeSpec.Generator/src/PostProcessing/PostProcessor.cs

@@ -20,6 +20,11 @@ internal class PostProcessor
         private readonly HashSet<string> _typesToKeep;
         private INamedTypeSymbol? _modelFactorySymbol;
 
+        private static readonly string[] _experimentalAttributeNames = ["Experimental", "ExperimentalAttribute"];
+
+        // CS8019: Unnecessary using directive.
+        private const string UnnecessaryUsingDirectiveDiagnosticId = "CS8019";
+
         public PostProcessor(
             HashSet<string> typesToKeep,
             string? modelFactoryFullName = null,
@@ -156,20 +161,39 @@ public async Task<Project> InternalizeAsync(Project project)
                 project = MarkInternal(project, model, documentId);
             }
 
+            if (nodesToInternalize.Count > 0)
+            {
+                CodeModelGenerator.Instance.Emitter.Info(
+                    $"Internalized {nodesToInternalize.Count} unreferenced public type(s).");
+            }
+
             var modelNamesToRemove =
                 nodesToInternalize.Keys.Select(item => item.Identifier.Text);
-            project = await RemoveMethodsFromModelFactoryAsync(project, definitions, modelNamesToRemove.ToHashSet());
+            DocumentId? modelFactoryDocumentId;
+            (project, modelFactoryDocumentId) = await RemoveMethodsFromModelFactoryAsync(project, definitions, modelNamesToRemove.ToHashSet());
+
+            if (nodesToInternalize.Count > 0)
+            {
+                var documentsToClean = nodesToInternalize.Values.ToHashSet();
+                // Removing methods from the model factory can leave a using directive (for a model in a
+                // different namespace) unused, so include the model factory document in the cleanup pass.
+                if (modelFactoryDocumentId != null)
+                {
+                    documentsToClean.Add(modelFactoryDocumentId);
+                }
+                project = await RemoveUnusedUsingsAsync(project, documentsToClean);
+            }
 
             return project;
         }
 
-        private async Task<Project> RemoveMethodsFromModelFactoryAsync(Project project,
+        private async Task<(Project Project, DocumentId? ModelFactoryDocumentId)> RemoveMethodsFromModelFactoryAsync(Project project,
             TypeSymbols definitions,
             HashSet<string> namesToRemove)
         {
             var modelFactorySymbol = definitions.ModelFactorySymbol;
             if (modelFactorySymbol == null)
-                return project;
+                return (project, null);
 
             var nodesToRemove = new List<SyntaxNode>();
 
@@ -203,7 +227,7 @@ private async Task<Project> RemoveMethodsFromModelFactoryAsync(Project project,
 
             // maybe this is possible, for instance, we could be adding the customization all entries previously inside the generated model factory so that the generated model factory is empty and removed.
             if (modelFactoryGeneratedDocument == null)
-                return project;
+                return (project, null);
 
             var root = await modelFactoryGeneratedDocument.GetSyntaxRootAsync();
             Debug.Assert(root is not null);
@@ -214,10 +238,10 @@ private async Task<Project> RemoveMethodsFromModelFactoryAsync(Project project,
             var methods = root.DescendantNodes().OfType<MethodDeclarationSyntax>();
             if (!methods.Any())
             {
-                return project.RemoveDocument(modelFactoryGeneratedDocument.Id);
+                return (project.RemoveDocument(modelFactoryGeneratedDocument.Id), null);
             }
 
-            return modelFactoryGeneratedDocument.Project;
+            return (modelFactoryGeneratedDocument.Project, modelFactoryGeneratedDocument.Id);
         }
 
         /// <summary>
@@ -336,7 +360,13 @@ private static IEnumerable<T> GetReferencedTypes<T>(T definition,
 
         private Project MarkInternal(Project project, BaseTypeDeclarationSyntax declarationNode, DocumentId documentId)
         {
+            CodeModelGenerator.Instance.Emitter.Debug(
+                $"Internalizing unreferenced public type '{declarationNode.Identifier.Text}'.");
+
             var newNode = ChangeModifier(declarationNode, SyntaxKind.PublicKeyword, SyntaxKind.InternalKeyword);
+            // The [Experimental] attribute is a public-API stability signal that is meaningless on a type that is
+            // being internalized, so strip it to avoid emitting it (and its use-site diagnostics) on internal types.
+            newNode = RemoveExperimentalAttribute(newNode, declarationNode.Identifier.Text);
             var tree = declarationNode.SyntaxTree;
             var document = project.GetDocument(documentId)!;
             var newRoot = tree.GetRoot().ReplaceNode(declarationNode, newNode)
@@ -345,6 +375,78 @@ private Project MarkInternal(Project project, BaseTypeDeclarationSyntax declarat
             return document.Project;
         }
 
+        /// <summary>
+        /// Removes any <c>[Experimental]</c> (<see cref="System.Diagnostics.CodeAnalysis.ExperimentalAttribute"/>)
+        /// attribute from <paramref name="declarationNode"/>. The declaration's leading trivia (such as documentation
+        /// comments and indentation) is re-attached to the resulting first token so that the surrounding formatting and
+        /// any documentation comment are preserved even when the attribute that carried them is removed.
+        /// </summary>
+        private static BaseTypeDeclarationSyntax RemoveExperimentalAttribute(
+            BaseTypeDeclarationSyntax declarationNode,
+            string typeName)
+        {
+            if (declarationNode.AttributeLists.Count == 0)
+            {
+                return declarationNode;
+            }
+
+            var newAttributeLists = new List<AttributeListSyntax>();
+            bool removed = false;
+
+            foreach (var attributeList in declarationNode.AttributeLists)
+            {
+                var keptAttributes = attributeList.Attributes
+                    .Where(attribute => !IsExperimentalAttribute(attribute))
+                    .ToList();
+
+                if (keptAttributes.Count == attributeList.Attributes.Count)
+                {
+                    // Nothing removed from this list.
+                    newAttributeLists.Add(attributeList);
+                }
+                else if (keptAttributes.Count > 0)
+                {
+                    // Keep the remaining (non-experimental) attributes in this list.
+                    removed = true;
+                    newAttributeLists.Add(attributeList.WithAttributes(SyntaxFactory.SeparatedList(keptAttributes)));
+                }
+                else
+                {
+                    // The entire list only contained experimental attributes and is dropped.
+                    removed = true;
+                }
+            }
+
+            if (!removed)
+            {
+                return declarationNode;
+            }
+
+            CodeModelGenerator.Instance.Emitter.Debug(
+                $"Removed [Experimental] attribute from '{typeName}' while internalizing it.");
+
+            // Preserve the original leading trivia (e.g. documentation comments and indentation) by re-attaching it to
+            // whatever token now leads the declaration. This keeps the doc comment even when it was attached to the
+            // attribute list that was removed.
+            var originalLeadingTrivia = declarationNode.GetLeadingTrivia();
+            var newNode = declarationNode.WithAttributeLists(SyntaxFactory.List(newAttributeLists));
+            var firstToken = newNode.GetFirstToken();
+
+            return newNode.ReplaceToken(firstToken, firstToken.WithLeadingTrivia(originalLeadingTrivia));
+        }
+
+        private static bool IsExperimentalAttribute(AttributeSyntax attribute)
+        {
+            var name = attribute.Name switch
+            {
+                QualifiedNameSyntax qualified => qualified.Right.Identifier.Text,
+                IdentifierNameSyntax identifier => identifier.Identifier.Text,
+                _ => attribute.Name.ToString()
+            };
+
+            return Array.IndexOf(_experimentalAttributeNames, name) >= 0;
+        }
+
         private async Task<Project> RemoveModelsAsync(Project project,
             IEnumerable<BaseTypeDeclarationSyntax> unusedModels)
         {
@@ -432,6 +534,70 @@ private async Task<Project> RemoveInvalidRefs(Project project)
             return solution.GetProject(project.Id)!;
         }
 
+        private async Task<Project> RemoveUnusedUsingsAsync(Project project, IEnumerable<DocumentId> documentIds)
+        {
+            var solution = project.Solution;
+            foreach (var documentId in documentIds)
+            {
+                solution = await RemoveUnusedUsings(solution, documentId);
+            }
+
+            return solution.GetProject(project.Id) ?? project;
+        }
+
+        private async Task<Solution> RemoveUnusedUsings(Solution solution, DocumentId documentId)
+        {
+            var document = solution.GetDocument(documentId);
+            if (document == null)
+            {
+                return solution;
+            }
+
+            document = await Simplifier.ReduceAsync(document);
+
+            var root = await document.GetSyntaxRootAsync();
+            var model = await document.GetSemanticModelAsync();
+
+            if (root is not CompilationUnitSyntax cu || model == null)
+            {
+                return solution;
+            }
+
+            // CS8019: Unnecessary using directive.
+            var unusedUsings = model.GetDiagnostics()
+                .Where(d => d.Id == UnnecessaryUsingDirectiveDiagnosticId)
+                .Select(d => cu.FindNode(d.Location.SourceSpan).FirstAncestorOrSelf<UsingDirectiveSyntax>())
+                .OfType<UsingDirectiveSyntax>()
+                .ToList();
+
+            if (unusedUsings.Count == 0)
+            {
+                return solution;
+            }
+
+            // Preserve any leading trivia on the first using directive (such as the file header and the
+            // #nullable directive) when that directive is removed, by carrying it over to the node that
+            // follows the removed directives.
+            var firstUsing = cu.Usings.FirstOrDefault();
+            var leadingTrivia = firstUsing is not null && unusedUsings.Contains(firstUsing)
+                ? firstUsing.GetLeadingTrivia()
+                : default;
+
+            var updatedRoot = cu.RemoveNodes(unusedUsings, SyntaxRemoveOptions.KeepNoTrivia);
+            if (updatedRoot == null)
+            {
+                return solution;
+            }
+
+            if (leadingTrivia.Count > 0)
+            {
+                var firstToken = updatedRoot.GetFirstToken();
+                updatedRoot = updatedRoot.ReplaceToken(firstToken, firstToken.WithLeadingTrivia(leadingTrivia.AddRange(firstToken.LeadingTrivia)));
+            }
+
+            return solution.WithDocumentSyntaxRoot(documentId, updatedRoot);
+        }
+
         private async Task<Solution> RemoveInvalidUsings(Solution solution, DocumentId documentId)
         {
             var document = solution.GetDocument(documentId)!;