Mrwcontext fixes (#7998)

Fixes #7987 and #7990

Author: JoshLove-msft

packages/http-client-csharp/generator/Microsoft.TypeSpec.Generator.ClientModel/src/Providers/ModelReaderWriterContextDefinition.cs

@@ -15,9 +15,9 @@ namespace Microsoft.TypeSpec.Generator.ClientModel.Providers
 {
  internal class ModelReaderWriterContextDefinition : TypeProvider
  {
- internal static string s_name = $"{RemovePeriods(ScmCodeModelGenerator.Instance.TypeFactory.PrimaryNamespace)}Context";
+ private static readonly string _name = $"{RemovePeriods(ScmCodeModelGenerator.Instance.TypeFactory.PrimaryNamespace)}Context";
 
- protected override string BuildName() => s_name;
+ protected override string BuildName() => _name;
 
  protected override string BuildRelativeFilePath() => Path.Combine("src", "Generated", "Models", $"{Name}.cs");
 

packages/http-client-csharp/generator/Microsoft.TypeSpec.Generator.ClientModel/src/ScmCodeModelGenerator.cs

@@ -20,6 +20,8 @@ public class ScmCodeModelGenerator : CodeModelGenerator
  private ScmOutputLibrary? _scmOutputLibrary;
  public override OutputLibrary OutputLibrary => _scmOutputLibrary ??= new();
 
+ internal ModelReaderWriterContextDefinition MrwContextDefinition { get; } = new ModelReaderWriterContextDefinition();
+
  public override ScmTypeFactory TypeFactory { get; }
 
  [ImportingConstructor]
@@ -36,7 +38,8 @@ protected override void Configure()
  AddMetadataReference(MetadataReference.CreateFromFile(typeof(ClientResult).Assembly.Location));
  AddMetadataReference(MetadataReference.CreateFromFile(typeof(BinaryData).Assembly.Location));
  AddMetadataReference(MetadataReference.CreateFromFile(typeof(JsonSerializer).Assembly.Location));
- AddTypeToKeepPublic(ModelReaderWriterContextDefinition.s_name);
+ AddTypeToKeepPublic(MrwContextDefinition);
+ AddNonRootType(MrwContextDefinition);
  }
  }
 }

packages/http-client-csharp/generator/Microsoft.TypeSpec.Generator.ClientModel/src/ScmOutputLibrary.cs

@@ -77,7 +77,7 @@ protected override TypeProvider[] BuildTypeProviders()
  new BinaryContentHelperDefinition(),
  new PipelineRequestHeadersExtensionsDefinition(),
  .. GetMultipartFormDataBinaryContentDefinition(),
- new ModelReaderWriterContextDefinition()
+ ScmCodeModelGenerator.Instance.MrwContextDefinition
  ];
  }
 

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

@@ -43,6 +43,9 @@ public async Task ExecuteAsync()
  await customCodeWorkspace.GetCompilationAsync(),
  await GeneratedCodeWorkspace.LoadBaselineContract());
 
+ // Configure must be called after the SourceInputModel is set.
+ CodeModelGenerator.Instance.Configure();
+
  GeneratedCodeWorkspace generatedCodeWorkspace = await GeneratedCodeWorkspace.Create();
 
  var output = CodeModelGenerator.Instance.OutputLibrary;

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

@@ -128,6 +128,7 @@ public void AddSharedSourceDirectory(string sharedSourceDirectory)
  internal HashSet<string> TypesToKeep { get; } = [];
 
  internal HashSet<string> TypesToKeepPublic { get; } = [];
+ internal HashSet<string> NonRootTypes { get; } = [];
 
  /// <summary>
  /// Adds a type to the list of types to keep.
@@ -147,7 +148,23 @@ public void AddTypeToKeep(string typeName)
  /// <summary>
  /// Adds a type to the list of types to keep as public.
  /// </summary>
- /// <param name="typeName">The type provider representing the type.</param>
- public void AddTypeToKeepPublic(string typeName) => TypesToKeepPublic.Add(typeName);
+ /// <param name="type">The type provider representing the type.</param>
+ public void AddTypeToKeepPublic(TypeProvider type) => TypesToKeepPublic.Add(type.Type.FullyQualifiedName);
+
+ /// <summary>
+ /// Adds a type to the list of non-root type providers. Non root type providers are types whose
+ /// references do not contribute to usages of the generated code. Therefore if the 'unreferenced-types-handling' property
+ /// is not set to 'keepAll', any types referenced by non-root type providers will not automatically be kept.
+ /// </summary>
+ /// <param name="type">The fully qualified type name.</param>
+ public void AddNonRootType(string type) => NonRootTypes.Add(type);
+
+ /// <summary>
+ /// Adds a type to the list of non-root type providers. Non root type providers are types whose
+ /// references do not contribute to usages of the generated code. Therefore if the 'unreferenced-types-handling' property
+ /// is not set to 'keepAll', any types referenced by non-root type providers will not automatically be kept.
+ /// </summary>
+ /// <param name="type">The type provider representing the type</param>
+ public void AddNonRootType(TypeProvider type) => NonRootTypes.Add(type.Type.FullyQualifiedName);
  }
 }

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

@@ -251,9 +251,12 @@ internal static Project AddDirectory(Project project, string directory, Func<str
  public async Task PostProcessAsync()
  {
  var modelFactory = CodeModelGenerator.Instance.OutputLibrary.ModelFactory.Value;
+ var nonRootTypes = CodeModelGenerator.Instance.NonRootTypes;
  var postProcessor = new PostProcessor(
  [.. CodeModelGenerator.Instance.TypeFactory.UnionVariantTypesToKeep, .. CodeModelGenerator.Instance.TypesToKeep],
- modelFactoryFullName: $"{modelFactory.Type.Namespace}.{modelFactory.Name}");
+ modelFactoryFullName: modelFactory.Type.FullyQualifiedName,
+ additionalNonRootTypeFullNames: nonRootTypes);
+
  switch (Configuration.UnreferencedTypesHandling)
  {
  case Configuration.UnreferencedTypesHandlingOption.KeepAll:

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

@@ -16,18 +16,18 @@ namespace Microsoft.TypeSpec.Generator
  internal class PostProcessor
  {
  private readonly string? _modelFactoryFullName;
- private readonly string? _aspExtensionClassName;
+ private readonly IEnumerable<string>? _additionalNonRootTypeFullNames;
  private readonly HashSet<string> _typesToKeep;
  private INamedTypeSymbol? _modelFactorySymbol;
 
  public PostProcessor(
  HashSet<string> typesToKeep,
  string? modelFactoryFullName = null,
- string? aspExtensionClassName = null)
+ IEnumerable<string>? additionalNonRootTypeFullNames = null)
  {
  _typesToKeep = typesToKeep;
  _modelFactoryFullName = modelFactoryFullName;
- _aspExtensionClassName = aspExtensionClassName;
+ _additionalNonRootTypeFullNames = additionalNonRootTypeFullNames;
  }
 
  private record TypeSymbols(
@@ -54,10 +54,22 @@ private async Task<TypeSymbols> GetTypeSymbolsAsync(Compilation compilation,
  var documentCache = new Dictionary<Document, HashSet<INamedTypeSymbol>>();
 
  if (_modelFactoryFullName != null)
+ {
  _modelFactorySymbol = compilation.GetTypeByMetadataName(_modelFactoryFullName);
- INamedTypeSymbol? aspDotNetExtensionSymbol = null;
- if (_aspExtensionClassName != null)
- aspDotNetExtensionSymbol = compilation.GetTypeByMetadataName(_aspExtensionClassName);
+ }
+
+ var additionalNonRootTypeSymbols = new HashSet<INamedTypeSymbol>(SymbolEqualityComparer.Default);
+ if (_additionalNonRootTypeFullNames != null)
+ {
+ foreach (var typeFullName in _additionalNonRootTypeFullNames)
+ {
+ var typeSymbol = compilation.GetTypeByMetadataName(typeFullName);
+ if (typeSymbol != null)
+ {
+ additionalNonRootTypeSymbols.Add(typeSymbol);
+ }
+ }
+ }
 
  foreach (var document in project.Documents)
  {
@@ -83,8 +95,10 @@ private async Task<TypeSymbols> GetTypeSymbolsAsync(Compilation compilation,
 
  // we do not add the model factory and aspDotNetExtension symbol to the declared symbol list so that it will never be included in any process of internalization or removal
  if (!SymbolEqualityComparer.Default.Equals(symbol, _modelFactorySymbol)
- && !SymbolEqualityComparer.Default.Equals(symbol, aspDotNetExtensionSymbol))
+ && !additionalNonRootTypeSymbols.Contains(symbol))
+ {
  result.Add(symbol);
+ }
 
  AddInList(declarationCache, symbol, typeDeclaration);
  AddInList(documentCache, document, symbol,
@@ -469,7 +483,39 @@ arg.Expression is TypeOfExpressionSyntax typeOfExpr &&
 
  if (invalidAttributes.Count > 0)
  {
+ // Check if any invalid attribute has type-level XML docs in its leading trivia
+ var attributeWithDocs = invalidAttributes
+ .OrderBy(a => a.SpanStart)
+ .FirstOrDefault(attr => attr.GetLeadingTrivia().Any(t =>
+ t.IsKind(SyntaxKind.SingleLineDocumentationCommentTrivia) ||
+ t.IsKind(SyntaxKind.MultiLineDocumentationCommentTrivia)));
+
+ SyntaxTriviaList? xmlDocs = null;
+ if (attributeWithDocs != null)
+ {
+ xmlDocs = attributeWithDocs.GetLeadingTrivia()
+ .Where(t => t.IsKind(SyntaxKind.SingleLineDocumentationCommentTrivia) ||
+ t.IsKind(SyntaxKind.MultiLineDocumentationCommentTrivia))
+ .ToSyntaxTriviaList();
+ }
+
+ // Remove all invalid attributes without keeping trivia
  cu = cu.RemoveNodes(invalidAttributes, SyntaxRemoveOptions.KeepNoTrivia)!;
+
+ // If we found XML docs, reattach them to the type declaration
+ if (xmlDocs?.Any() == true)
+ {
+ var typeDecl = cu.DescendantNodes()
+ .OfType<TypeDeclarationSyntax>()
+ .FirstOrDefault();
+
+ if (typeDecl != null)
+ {
+ cu = cu.ReplaceNode(typeDecl,
+ typeDecl.WithLeadingTrivia(xmlDocs.Value.AddRange(typeDecl.GetLeadingTrivia())));
+ }
+ }
+
  solution = solution.WithDocumentSyntaxRoot(documentId, cu);
  }
 

packages/http-client-csharp/generator/Microsoft.TypeSpec.Generator/src/StartUp/GeneratorHandler.cs

@@ -156,7 +156,6 @@ internal void SelectGenerator(CommandLineOptions options)
  }
  }
 
- CodeModelGenerator.Instance.Configure();
  loaded = true;
  break;
  }

packages/http-client-csharp/generator/Microsoft.TypeSpec.Generator/test/PostProcessing/PostProcessorTests.cs

@@ -1,6 +1,8 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+using System.ClientModel.Primitives;
+using System.Collections.Generic;
 using System.IO;
 using System.Linq;
 using System.Threading.Tasks;
@@ -58,6 +60,7 @@ public async Task RemovesInvalidUsings()
  CollectionAssert.Contains(usings, "System");
  }
 
+
  [Test]
  public async Task DoesNotRemoveValidUsings()
  {
@@ -102,11 +105,129 @@ public async Task DoesNotRemoveValidUsings()
  CollectionAssert.Contains(usings, "System");
  }
 
+ [Test]
+ public async Task RemovesInvalidAttributes()
+ {
+ MockHelpers.LoadMockGenerator();
+ var workspace = new AdhocWorkspace();
+ var projectInfo = ProjectInfo.Create(
+ ProjectId.CreateNewId(),
+ VersionStamp.Create(),
+ name: "TestProj",
+ assemblyName: "TestProj",
+ language: LanguageNames.CSharp)
+ .WithMetadataReferences(new[]
+ {
+ MetadataReference.CreateFromFile(typeof(object).Assembly.Location),
+ MetadataReference.CreateFromFile(typeof(ModelReaderWriterBuildableAttribute).Assembly.Location)
+ });
+
+ var project = workspace.AddProject(projectInfo);
+ var folder = Helpers.GetAssetFileOrDirectoryPath(false);
+ const string removesInvalidAttributesFileName = "RemovesInvalidAttributes.cs";
+ project = project.AddDocument(
+ removesInvalidAttributesFileName,
+ File.ReadAllText(Path.Join(folder, removesInvalidAttributesFileName))).Project;
+ project = project.AddDocument(
+ "Model.cs",
+ File.ReadAllText(Path.Join(folder, "Model.cs"))).Project;
+ project = project.AddDocument(
+ "RootClass.cs",
+ File.ReadAllText(Path.Join(folder, "RootClass.cs"))).Project;
+ var postProcessor = new TestPostProcessor("RootClass.cs", nonRootTypes: ["Sample.KeepMe"]);
+
+ var resultProject = await postProcessor.RemoveAsync(project);
+ var doc= resultProject.Documents
+ .Single(d => d.Name == removesInvalidAttributesFileName);
+ var root = await doc.GetSyntaxRootAsync();
+ var compilation = (CompilationUnitSyntax)root!;
+
+ var type = compilation
+ .DescendantNodes()
+ .OfType<BaseTypeDeclarationSyntax>()
+ .SingleOrDefault(t => t.Identifier.Text == "KeepMe");
+ Assert.IsNotNull(type, "The class 'KeepMe' should still exist.");
+
+ var attributes = type!.AttributeLists.SelectMany(a => a.Attributes).Select(al => al.ToString()).ToList();
+
+ // The invalid attribute should be removed
+ CollectionAssert.DoesNotContain(attributes, "ModelReaderWriterBuildable");
+
+ // The class documentation should be preserved
+ var classDocumentation = compilation.DescendantNodes()
+ .OfType<ClassDeclarationSyntax>()
+ .Single(c => c.Identifier.Text == "KeepMe")
+ .GetLeadingTrivia()
+ .Select(t => t.ToString())
+ .SingleOrDefault(t => t.Trim().Contains("<summary>"));
+
+ StringAssert.Contains("Class docs that should be kept.", classDocumentation);
+ }
+
+ [Test]
+ public async Task DoesNotRemoveValidAttributes()
+ {
+ MockHelpers.LoadMockGenerator();
+ var workspace = new AdhocWorkspace();
+ var projectInfo = ProjectInfo.Create(
+ ProjectId.CreateNewId(),
+ VersionStamp.Create(),
+ name: "TestProj",
+ assemblyName: "TestProj",
+ language: LanguageNames.CSharp)
+ .WithMetadataReferences(new[]
+ {
+ MetadataReference.CreateFromFile(typeof(object).Assembly.Location),
+ MetadataReference.CreateFromFile(typeof(ModelReaderWriterBuildableAttribute).Assembly.Location)
+ });
+
+ var project = workspace.AddProject(projectInfo);
+ var folder = Helpers.GetAssetFileOrDirectoryPath(false);
+ const string doesNotRemoveValidAttributesFileName = "DoesNotRemoveValidAttributes.cs";
+ project = project.AddDocument(
+ doesNotRemoveValidAttributesFileName,
+ File.ReadAllText(Path.Join(folder, doesNotRemoveValidAttributesFileName))).Project;
+ project = project.AddDocument(
+ "Model.cs",
+ File.ReadAllText(Path.Join(folder, "Model.cs"))).Project;
+ project = project.AddDocument(
+ "RootClass.cs",
+ File.ReadAllText(Path.Join(folder, "RootClass.cs"))).Project;
+ var postProcessor = new TestPostProcessor("RootClass.cs");
+
+ var resultProject = await postProcessor.RemoveAsync(project);
+ var doc= resultProject.Documents
+ .Single(d => d.Name == doesNotRemoveValidAttributesFileName);
+ var root = await doc.GetSyntaxRootAsync();
+ var compilation = (CompilationUnitSyntax)root!;
+
+ var type = compilation
+ .DescendantNodes()
+ .OfType<BaseTypeDeclarationSyntax>()
+ .SingleOrDefault(t => t.Identifier.Text == "KeepMe");
+ Assert.IsNotNull(type, "The class 'KeepMe' should still exist.");
+
+ var attributes = type!.AttributeLists.SelectMany(a => a.Attributes).Select(al => al.ToString().Trim()).ToList();
+
+ // The valid attribute should be retained
+ CollectionAssert.Contains(attributes, "ModelReaderWriterBuildable(typeof(Model))");
+
+ // The class documentation should be preserved
+ var classDocumentation = compilation.DescendantNodes()
+ .OfType<ClassDeclarationSyntax>()
+ .Single(c => c.Identifier.Text == "KeepMe")
+ .GetLeadingTrivia()
+ .Select(t => t.ToString())
+ .SingleOrDefault(t => t.Trim().Contains("<summary>"));
+
+ StringAssert.Contains("Class docs that should be kept.", classDocumentation);
+ }
+
  private class TestPostProcessor : PostProcessor
  {
  private readonly string _rootFile;
 
- public TestPostProcessor(string rootFile) : base([])
+ public TestPostProcessor(string rootFile, IEnumerable<string>? nonRootTypes = null) : base([], additionalNonRootTypeFullNames: nonRootTypes)
  {
  _rootFile = rootFile;
  }

packages/http-client-csharp/generator/Microsoft.TypeSpec.Generator/test/PostProcessing/TestData/PostProcessorTests/DoesNotRemoveValidAttributes/DoesNotRemoveValidAttributes.cs

@@ -0,0 +1,15 @@
+using System;
+using System.ClientModel.Primitives;
+using Sample.Models;
+
+namespace Sample
+{
+ /// <summary>
+ /// Class docs that should be kept.
+ /// </summary>
+ [ModelReaderWriterBuildable(typeof(Model))]
+ public class KeepMe
+ {
+ public void Foo() => Console.WriteLine(""hello"");
+ }
+}