AOT generator: self-contained reflection-free attribute registration

src/Adapter/MSTestAdapter.PlatformServices/PublicAPI/PublicAPI.Unshipped.txt

@@ -1,3 +1,4 @@
 #nullable enable
 Microsoft.VisualStudio.TestPlatform.MSTestAdapter.PlatformServices.SourceGeneration.ReflectionMetadataHook
 static Microsoft.VisualStudio.TestPlatform.MSTestAdapter.PlatformServices.SourceGeneration.ReflectionMetadataHook.Register(System.Reflection.Assembly! assembly, System.Type![]! types, System.Collections.Generic.IReadOnlyDictionary<System.Type!, System.Reflection.MethodInfo![]!>! testMethods) -> void
+static Microsoft.VisualStudio.TestPlatform.MSTestAdapter.PlatformServices.SourceGeneration.ReflectionMetadataHook.Register(System.Reflection.Assembly! assembly, System.Type![]! types, System.Collections.Generic.IReadOnlyDictionary<System.Type!, System.Reflection.MethodInfo![]!>! testMethods, System.Collections.Generic.IReadOnlyDictionary<System.Type!, System.Attribute![]!>! typeAttributes, object![]! assemblyAttributes) -> void

src/Adapter/MSTestAdapter.PlatformServices/SourceGeneration/CompositeSourceGeneratedReflectionDataProvider.cs

@@ -65,8 +65,8 @@ internal override void GetNavigationData(string className, string methodName, ou
     internal override object[] GetAssemblyAttributes(Assembly assembly)
     {
         CompositeState state = Volatile.Read(ref _state);
-        return state.ProvidersByAssembly.TryGetValue(assembly, out SourceGeneratedReflectionDataProvider? provider)
-            ? provider.AssemblyAttributes
+        return state.AssemblyAttributesByAssembly.TryGetValue(assembly, out object[]? merged)
+            ? merged
             : [];
     }
 
@@ -96,18 +96,21 @@ private sealed class CompositeState
             providers: [],
             providersByAssemblyName: new Dictionary<string, SourceGeneratedReflectionDataProvider>(StringComparer.OrdinalIgnoreCase),
             providersByAssembly: [],
+            assemblyAttributesByAssembly: [],
             mergedSnapshot: new SourceGeneratedReflectionDataProvider());
 #pragma warning restore IDE0028
 
         private CompositeState(
             IReadOnlyList<SourceGeneratedReflectionDataProvider> providers,
             Dictionary<string, SourceGeneratedReflectionDataProvider> providersByAssemblyName,
             Dictionary<Assembly, SourceGeneratedReflectionDataProvider> providersByAssembly,
+            Dictionary<Assembly, object[]> assemblyAttributesByAssembly,
             SourceGeneratedReflectionDataProvider mergedSnapshot)
         {
             Providers = providers;
             ProvidersByAssemblyName = providersByAssemblyName;
             ProvidersByAssembly = providersByAssembly;
+            AssemblyAttributesByAssembly = assemblyAttributesByAssembly;
             MergedSnapshot = mergedSnapshot;
         }
 
@@ -117,6 +120,12 @@ private CompositeState(
 
         public Dictionary<Assembly, SourceGeneratedReflectionDataProvider> ProvidersByAssembly { get; }
 
+        // Cumulative per-assembly assembly-level attribute union. Unlike ProvidersByAssembly (which
+        // keeps the last provider per assembly for type/method lookups), this dictionary preserves
+        // the union of AssemblyAttributes across every Register call for the same assembly so a
+        // future re-registration doesn't silently drop attributes published by an earlier one.
+        public Dictionary<Assembly, object[]> AssemblyAttributesByAssembly { get; }
+
         public SourceGeneratedReflectionDataProvider MergedSnapshot { get; }
 
         public CompositeState With(SourceGeneratedReflectionDataProvider added)
@@ -131,15 +140,32 @@ public CompositeState With(SourceGeneratedReflectionDataProvider added)
             byName[added.AssemblyName] = added;
 
             var byAssembly = new Dictionary<Assembly, SourceGeneratedReflectionDataProvider>(ProvidersByAssembly);
+            var attributesByAssembly = new Dictionary<Assembly, object[]>(AssemblyAttributesByAssembly);
             if (added.Assembly is { } addedAssembly)
             {
                 byAssembly[addedAssembly] = added;
+
+                if (attributesByAssembly.TryGetValue(addedAssembly, out object[]? existing) && existing.Length > 0)
+                {
+                    if (added.AssemblyAttributes.Length > 0)
+                    {
+                        object[] merged = new object[existing.Length + added.AssemblyAttributes.Length];
+                        Array.Copy(existing, 0, merged, 0, existing.Length);
+                        Array.Copy(added.AssemblyAttributes, 0, merged, existing.Length, added.AssemblyAttributes.Length);
+                        attributesByAssembly[addedAssembly] = merged;
+                    }
+                }
+                else if (added.AssemblyAttributes.Length > 0)
+                {
+                    attributesByAssembly[addedAssembly] = added.AssemblyAttributes;
+                }
             }
 
             return new CompositeState(
                 providers,
                 byName,
                 byAssembly,
+                attributesByAssembly,
                 BuildMergedSnapshot(providers));
         }
 

src/Adapter/MSTestAdapter.PlatformServices/SourceGeneration/ReflectionMetadataHook.cs

@@ -21,7 +21,7 @@ namespace Microsoft.VisualStudio.TestPlatform.MSTestAdapter.PlatformServices.Sou
 /// that needs to call it across the assembly boundary, and module initializers cannot use
 /// <c>internal</c> APIs from a different assembly. The signature and behaviour of this hook are
 /// implementation details that may evolve with the generator without a major version bump; do
-/// not hand-roll a call to <see cref="Register"/> from your own code.
+/// not hand-roll a call to <see cref="Register(Assembly, Type[], IReadOnlyDictionary{Type, MethodInfo[]})"/> from your own code.
 /// </para>
 /// <para>
 /// <b>Discovery limitation.</b> The MSTest source generator only enumerates types that carry
@@ -62,6 +62,33 @@ public static class ReflectionMetadataHook
     /// </remarks>
     [EditorBrowsable(EditorBrowsableState.Never)]
     public static void Register(Assembly assembly, Type[] types, IReadOnlyDictionary<Type, MethodInfo[]> testMethods)
+        => Register(assembly, types, testMethods, EmptyTypeAttributes, []);
+
+    /// <summary>
+    /// <b>Infrastructure.</b> Publishes source-generated metadata for <paramref name="assembly"/>
+    /// to the MSTest adapter, including pre-materialized type-level and assembly-level attributes
+    /// so the adapter serves them without runtime reflection. Safe to call from multiple module
+    /// initializers; later registrations are merged with earlier ones.
+    /// </summary>
+    /// <param name="assembly">The test assembly the metadata describes.</param>
+    /// <param name="types">All types directly annotated with <c>[TestClass]</c> in the assembly.</param>
+    /// <param name="testMethods">A map from each test class to its <c>[TestMethod]</c> set.</param>
+    /// <param name="typeAttributes">
+    /// A map from each test class to its pre-inflated <see cref="Attribute"/> instances. The adapter
+    /// returns these from <c>GetCustomAttributes(Type)</c> instead of reflecting at runtime.
+    /// </param>
+    /// <param name="assemblyAttributes">Pre-inflated assembly-level attribute instances.</param>
+    /// <remarks>
+    /// Do not call this method from hand-written code; it is meant to be invoked exclusively from
+    /// the <c>[ModuleInitializer]</c> emitted by the MSTest source generator.
+    /// </remarks>
+    [EditorBrowsable(EditorBrowsableState.Never)]
+    public static void Register(
+        Assembly assembly,
+        Type[] types,
+        IReadOnlyDictionary<Type, MethodInfo[]> testMethods,
+        IReadOnlyDictionary<Type, Attribute[]> typeAttributes,
+        object[] assemblyAttributes)
     {
         if (assembly is null)
         {
@@ -78,6 +105,16 @@ public static void Register(Assembly assembly, Type[] types, IReadOnlyDictionary
             throw new ArgumentNullException(nameof(testMethods));
         }
 
+        if (typeAttributes is null)
+        {
+            throw new ArgumentNullException(nameof(typeAttributes));
+        }
+
+        if (assemblyAttributes is null)
+        {
+            throw new ArgumentNullException(nameof(assemblyAttributes));
+        }
+
         var typesCopy = (Type[])types.Clone();
 
         var testMethodsCopy = new Dictionary<Type, MethodInfo[]>(testMethods.Count);
@@ -86,6 +123,14 @@ public static void Register(Assembly assembly, Type[] types, IReadOnlyDictionary
             testMethodsCopy[kvp.Key] = (MethodInfo[])kvp.Value.Clone();
         }
 
+        var typeAttributesCopy = new Dictionary<Type, Attribute[]>(typeAttributes.Count);
+        foreach (KeyValuePair<Type, Attribute[]> kvp in typeAttributes)
+        {
+            typeAttributesCopy[kvp.Key] = (Attribute[])kvp.Value.Clone();
+        }
+
+        object[] assemblyAttributesCopy = (object[])assemblyAttributes.Clone();
+
         // TypesByName must always match Type.FullName at runtime (see comment in the source
         // generator emitter): compute it on the runtime side from typeof(T).FullName so the
         // generator emits less code and the same FullName conventions are honored for nested
@@ -106,6 +151,8 @@ public static void Register(Assembly assembly, Type[] types, IReadOnlyDictionary
             Types = typesCopy,
             TypesByName = typesByName,
             TypeMethods = testMethodsCopy,
+            TypeAttributes = typeAttributesCopy,
+            AssemblyAttributes = assemblyAttributesCopy,
         };
 
         lock (Lock)
@@ -121,4 +168,6 @@ public static void Register(Assembly assembly, Type[] types, IReadOnlyDictionary
             }
         }
     }
+
+    private static readonly Dictionary<Type, Attribute[]> EmptyTypeAttributes = [];
 }

src/Adapter/MSTestAdapter.PlatformServices/SourceGeneration/SourceGeneratedReflectionOperations.cs

@@ -10,14 +10,15 @@ namespace Microsoft.VisualStudio.TestPlatform.MSTestAdapter.PlatformServices.Sou
 /// </summary>
 /// <remarks>
 /// <para>
-/// <b>Scope of the source-generated data.</b> The MSTest source generator's
-/// <c>[ModuleInitializer]</c> calls <see cref="ReflectionMetadataHook.Register"/> with only:
-/// the assembly, the <c>[TestClass]</c> types, and their <c>[TestMethod]</c>-annotated
-/// <see cref="MethodInfo"/>s. Everything else (type attributes, method attributes, assembly
-/// attributes, constructors, properties, navigation data) is intentionally <i>not</i>
-/// populated today; reads of those fields fall back to runtime reflection. The source-gen
-/// payload is therefore best understood as "type / test-method rooting + trimmer hints"
-/// rather than a full reflection replacement.
+/// <b>Scope of the source-generated data.</b> The shipping MSTest source generator's
+/// <c>[ModuleInitializer]</c> calls <see cref="ReflectionMetadataHook.Register(Assembly, Type[], IReadOnlyDictionary{Type, MethodInfo[]})"/>
+/// with only: the assembly, the <c>[TestClass]</c> types, and their <c>[TestMethod]</c>-annotated
+/// <see cref="MethodInfo"/>s. The AOT generator additionally publishes pre-materialized type-level
+/// and assembly-level attributes via the richer
+/// <see cref="ReflectionMetadataHook.Register(Assembly, Type[], IReadOnlyDictionary{Type, MethodInfo[]}, IReadOnlyDictionary{Type, Attribute[]}, object[])"/>
+/// overload. Anything still not populated (method attributes, constructors, properties, navigation
+/// data) falls back to runtime reflection, so the payload remains "type / test-method rooting +
+/// trimmer hints + materialized type/assembly attributes" rather than a full reflection replacement.
 /// </para>
 /// <para>
 /// <b>Why a fallback exists at all.</b> Each fallback in this class falls into one of three

src/Analyzers/MSTest.AotReflection.SourceGeneration/Generators/MSTestReflectionMetadataGenerator.cs

@@ -77,6 +77,20 @@ public void Initialize(IncrementalGeneratorInitializationContext context)
             string source = MetadataRegistryEmitter.EmitRegistry(assemblyName, payload.Metadata, payload.Classes);
             ctx.AddSource("MSTestReflectionMetadata.Registry.g.cs", SourceText.From(source, Encoding.UTF8));
         });
+
+        // Emit the [ModuleInitializer] that registers this assembly with the adapter. Without it,
+        // referencing this generator would emit metadata that nothing consumes. We skip emission
+        // when there are no test classes — there is nothing to register.
+        context.RegisterImplementationSourceOutput(combined, static (ctx, payload) =>
+        {
+            if (payload.Classes.IsDefaultOrEmpty)
+            {
+                return;
+            }
+
+            string source = RuntimeRegistrationEmitter.Emit(payload.Metadata, payload.Classes);
+            ctx.AddSource("MSTestReflectionMetadata.Registration.g.cs", SourceText.From(source, Encoding.UTF8));
+        });
     }
 
     private static TestClassTransformResult BuildResult(GeneratorAttributeSyntaxContext context, CancellationToken cancellationToken)

src/Analyzers/MSTest.AotReflection.SourceGeneration/Generators/MetadataRegistryEmitter.cs

@@ -416,7 +416,7 @@ private static void EmitAttributesProperty(IndentedStringBuilder sb, string prop
         }
     }
 
-    private static string BuildAttributeExpression(AttributeApplicationModel attribute)
+    internal static string BuildAttributeExpression(AttributeApplicationModel attribute)
     {
         string ctorArgs = string.Join(", ", attribute.ConstructorArguments.AsImmutableArray().Select(BuildConstantExpression));
         string ctorCall = $"new {attribute.FullyQualifiedAttributeType}({ctorArgs})";
@@ -502,7 +502,7 @@ private static string BuildArgumentsFromObjectArray(EquatableArray<TestParameter
 
     private static string Bool(bool value) => value ? "true" : "false";
 
-    private static string Escape(string value)
+    internal static string Escape(string value)
         => value.Replace("\\", "\\\\").Replace("\"", "\\\"");
 
     private static void AppendHeader(IndentedStringBuilder sb)