Update TypeMap attribute handling

src/tools/illink/src/linker/Linker.Steps/MarkStep.cs

@@ -254,14 +254,7 @@ protected virtual void Initialize()
         {
             InitializeCorelibAttributeXml();
             Context.Pipeline.InitializeMarkHandlers(Context, MarkContext);
-
-            if (Annotations.GetEntryPointAssembly() is AssemblyDefinition entryPoint)
-            {
-                _typeMapHandler = new TypeMapHandler(entryPoint);
-            }
-
-            _typeMapHandler.Initialize(Context, this);
-
+            _typeMapHandler.Initialize(Context, this, Annotations.GetEntryPointAssembly());
             ProcessMarkedPending();
         }
 
@@ -1458,7 +1451,7 @@ protected bool CheckProcessed(IMetadataTokenProvider provider)
             return !Annotations.SetProcessed(provider);
         }
 
-        protected virtual void MarkAssembly(AssemblyDefinition assembly, DependencyInfo reason, MessageOrigin origin)
+        public void MarkAssembly(AssemblyDefinition assembly, DependencyInfo reason, MessageOrigin origin)
         {
             Annotations.Mark(assembly, reason, origin);
             if (CheckProcessed(assembly))

src/tools/illink/src/linker/Linker.Steps/SweepStep.cs

@@ -261,7 +261,7 @@ protected virtual void SweepAssembly(AssemblyDefinition assembly)
 
         bool IsMarkedAssembly(AssemblyDefinition assembly)
         {
-            return Annotations.IsMarked(assembly.MainModule);
+            return Annotations.IsMarked(assembly) || Annotations.IsMarked(assembly.MainModule);
         }
 
         bool CanSweepNamesForMember(IMemberDefinition member, MetadataTrimming metadataTrimming)

src/tools/illink/src/linker/Linker/DependencyInfo.cs

@@ -147,6 +147,7 @@ public enum DependencyKind
         UnsafeAccessorTarget = 89, // the member is referenced via UnsafeAccessor attribute
 
         TypeMapEntry = 90, // external or proxy type map entry
+        TypeMapAssemblyTarget = 91, // TypeMapTargetAssembly for a used typeMapGroup
     }
 
     public readonly struct DependencyInfo : IEquatable<DependencyInfo>

src/tools/illink/src/linker/Linker/TypeReferenceEqualityComparer.cs

@@ -11,6 +11,9 @@ namespace Mono.Linker
     // Copied from https://github.com/jbevain/cecil/blob/master/Mono.Cecil/TypeReferenceComparer.cs
     internal sealed class TypeReferenceEqualityComparer : EqualityComparer<TypeReference>
     {
+        [Obsolete("Default will point to default object equality comparer. Use the constructor to create an instance with a resolver.")]
+        public static new object Default => throw new InvalidOperationException();
+
         public readonly ITryResolveMetadata _resolver;
 
         public TypeReferenceEqualityComparer(ITryResolveMetadata resolver)

src/tools/illink/test/Mono.Linker.Tests.Cases/Reflection/Dependencies/TypeMapReferencedAssembly.cs

@@ -0,0 +1,51 @@
+// Copyright (c) .NET Foundation and contributors. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+
+using System.Runtime.InteropServices;
+using Mono.Linker.Tests.Cases.Reflection.Dependencies;
+using Mono.Linker.Tests.Cases.Reflection.Dependencies.Library;
+
+[assembly: TypeMap<UsedTypeMapUniverse>("UnimportantString", typeof(TargetTypeUnconditional1))]
+[assembly: TypeMap<UsedTypeMapUniverse>("UnimportantString", typeof(TargetTypeConditional1), typeof(TrimTarget1))]
+[assembly: TypeMapAssociation<UsedTypeMapUniverse>(typeof(ProxySource1), typeof(ProxyTarget1))]
+
+[assembly: TypeMap<UnusedTypeMapUniverse>("UnimportantString", typeof(TargetTypeUnconditional2))]
+[assembly: TypeMap<UnusedTypeMapUniverse>("UnimportantString", typeof(TargetTypeConditional2), typeof(TrimTarget2))]
+[assembly: TypeMapAssociation<UnusedTypeMapUniverse>(typeof(ProxySource2), typeof(ProxyTarget2))]
+[assembly: TypeMapAssemblyTarget<UsedTypeMapUniverse>("library2")]
+
+namespace Mono.Linker.Tests.Cases.Reflection.Dependencies
+{
+    public class TypeMapReferencedAssembly
+    {
+        public static void Main()
+        {
+            // Mark expected trim targets
+            _ = new TrimTarget1();
+            _ = new ProxySource1();
+            _ = new TrimTarget2();
+            _ = new ProxySource2();
+
+            // Mark expected type map universe
+            _ = TypeMapping.GetOrCreateExternalTypeMapping<UsedTypeMapUniverse>();
+            _ = TypeMapping.GetOrCreateProxyTypeMapping<UsedTypeMapUniverse>();
+        }
+    }
+
+    public class UsedTypeMapUniverse;
+    public class UnusedTypeMapUniverse;
+}
+
+namespace Mono.Linker.Tests.Cases.Reflection.Dependencies.Library
+{
+    public class TargetTypeUnconditional1;
+    public class TargetTypeConditional1;
+    public class ProxySource1;
+    public class ProxyTarget1;
+    public class TargetTypeUnconditional2;
+    public class TargetTypeConditional2;
+    public class ProxySource2;
+    public class ProxyTarget2;
+    public class TrimTarget1;
+    public class TrimTarget2;
+}

src/tools/illink/test/Mono.Linker.Tests.Cases/Reflection/Dependencies/TypeMapSecondOrderReference.cs

@@ -0,0 +1,43 @@
+
+// Copyright (c) .NET Foundation and contributors. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+
+using System.Runtime.InteropServices;
+using Mono.Linker.Tests.Cases.Reflection.Dependencies;
+using Mono.Linker.Tests.Cases.Reflection.Dependencies.Library2;
+
+// Nothing in this assembly should be referenced or rooted directly
+// This should validate the if only the type map attributes are kept, the assembly is still preserved
+[assembly: TypeMap<string>("UnimportantString", typeof(long))]
+[assembly: TypeMap<string>("UnimportantString", typeof(uint), typeof(string))]
+[assembly: TypeMapAssociation<string>(typeof(int), typeof(uint))]
+
+[assembly: TypeMap<UnusedTypeMapUniverse2>("UnimportantString", typeof(TargetTypeUnconditional2))]
+[assembly: TypeMap<UnusedTypeMapUniverse2>("UnimportantString", typeof(TargetTypeConditional2), typeof(TrimTarget2))]
+[assembly: TypeMapAssociation<UnusedTypeMapUniverse2>(typeof(ProxySource2), typeof(ProxyTarget2))]
+
+[assembly: TypeMapAssemblyTarget<string>("library")] // Circular reference
+
+namespace Mono.Linker.Tests.Cases.Reflection.Dependencies
+{
+    public class TypeMapReferencedAssembly2
+    {
+    }
+
+    public class UsedTypeMapUniverse2;
+    public class UnusedTypeMapUniverse2;
+}
+
+namespace Mono.Linker.Tests.Cases.Reflection.Dependencies.Library2
+{
+    public class TargetTypeUnconditional1;
+    public class TargetTypeConditional1;
+    public class ProxySource1;
+    public class ProxyTarget1;
+    public class TargetTypeUnconditional2;
+    public class TargetTypeConditional2;
+    public class ProxySource2;
+    public class ProxyTarget2;
+    public class TrimTarget1;
+    public class TrimTarget2;
+}

src/tools/illink/test/Mono.Linker.Tests.Cases/Reflection/TypeMap.cs

@@ -11,9 +11,13 @@
 using Mono.Linker.Tests.Cases.Expectations.Assertions;
 using Mono.Linker.Tests.Cases.Expectations.Metadata;
 using Mono.Linker.Tests.Cases.Reflection;
+using Mono.Linker.Tests.Cases.Reflection.Dependencies;
+using Mono.Linker.Tests.Cases.Reflection.Dependencies.Library;
 
-[assembly: KeptAttributeAttribute(typeof(TypeMapAttribute<UsedTypeMap>), By = Tool.Trimmer)]
-[assembly: KeptAttributeAttribute(typeof(TypeMapAssociationAttribute<UsedTypeMap>), By = Tool.Trimmer)]
+[assembly: KeptAttributeAttribute(typeof(TypeMapAttribute<UsedTypeMap>))]
+[assembly: KeptAttributeAttribute(typeof(TypeMapAssociationAttribute<UsedTypeMap>))]
+[assembly: KeptAttributeAttribute(typeof(TypeMapAssociationAttribute<UsedProxyTypeMap>))]
+[assembly: KeptAttributeAttribute(typeof(TypeMapAttribute<UsedExternalTypeMap>))]
 [assembly: TypeMap<UsedTypeMap>("TrimTargetIsTarget", typeof(TargetAndTrimTarget), typeof(TargetAndTrimTarget))]
 [assembly: TypeMap<UsedTypeMap>("TrimTargetIsUnrelated", typeof(TargetType), typeof(TrimTarget))]
 [assembly: TypeMap<UsedTypeMap>(nameof(AllocatedNoTypeCheckClassTarget), typeof(AllocatedNoTypeCheckClassTarget), typeof(AllocatedNoTypeCheckClass))]
@@ -28,6 +32,15 @@
 [assembly: TypeMap<UsedTypeMap>("Ldobj", typeof(LdobjTarget), typeof(LdobjType))]
 [assembly: TypeMap<UsedTypeMap>("ArrayElement", typeof(ArrayElementTarget), typeof(ArrayElement))]
 [assembly: TypeMap<UsedTypeMap>("TrimTargetIsAllocatedNoTypeCheckNoBoxStruct", typeof(ConstructedNoTypeCheckOrBoxTarget), typeof(ConstructedNoTypeCheckNoBoxStruct))]
+
+// The TypeMap Universes are kept separate such that Proxy attributes shouldn't be kept if only the External type map is needed.
+[assembly: TypeMap<UsedExternalTypeMap>("UsedOnlyForExternalTypeMap", typeof(UsedExternalTarget), typeof(UsedTrimTarget))] // Kept
+[assembly: TypeMapAssociation<UsedExternalTypeMap>(typeof(UsedProxySource), typeof(UsedProxyTarget))] // Removed
+
+// The TypeMap Universes are kept separate such that External attributes shouldn't be kept if only the Proxy type map is needed.
+[assembly: TypeMap<UsedProxyTypeMap>("UsedOnlyForExternalTypeMap", typeof(UsedExternalTarget2), typeof(UsedTrimTarget2))] // Removed
+[assembly: TypeMapAssociation<UsedProxyTypeMap>(typeof(UsedProxySource2), typeof(UsedProxyTarget2))] // Kept
+
 [assembly: TypeMapAssociation<UsedTypeMap>(typeof(SourceClass), typeof(ProxyType))]
 [assembly: TypeMapAssociation<UsedTypeMap>(typeof(TypeCheckOnlyClass), typeof(TypeCheckOnlyProxy))]
 [assembly: TypeMapAssociation<UsedTypeMap>(typeof(AllocatedNoBoxStructType), typeof(AllocatedNoBoxProxy))]
@@ -40,10 +53,61 @@
 [assembly: TypeMap<UsedTypeMap>("ClassWithStaticMethod", typeof(TargetType4), typeof(ClassWithStaticMethod))]
 [assembly: TypeMap<UsedTypeMap>("ClassWithStaticMethodAndField", typeof(TargetType5), typeof(ClassWithStaticMethodAndField))]
 
+[assembly: KeptAttributeAttribute(typeof(TypeMapAssemblyTargetAttribute<UsedTypeMap>))]
+[assembly: TypeMapAssemblyTarget<UsedTypeMap>("library")]
+// TypeMapAssemblyTarget is kept regardless of which type map the program needs (External or Proxy)
+[assembly: KeptAttributeAttribute(typeof(TypeMapAssemblyTargetAttribute<UsedProxyTypeMap>))]
+[assembly: KeptAttributeAttribute(typeof(TypeMapAssemblyTargetAttribute<UsedExternalTypeMap>))]
+[assembly: TypeMapAssemblyTarget<UsedProxyTypeMap>("library")]
+[assembly: TypeMapAssemblyTarget<UsedExternalTypeMap>("library")]
+[assembly: TypeMapAssemblyTarget<UnusedTypeMap2>("library")] // Should be removed
+
 namespace Mono.Linker.Tests.Cases.Reflection
 {
-    [Kept]
+    [SetupLinkerAction("link", "System.Private.CoreLib")] // Needed to get the RemoveAttributeInstances in embedded xml
+    [SetupLinkerArgument("--ignore-link-attributes", "false")]
     [SetupCompileArgument("/unsafe")]
+    [SetupCompileBefore("library.dll", new[] { "Dependencies/TypeMapReferencedAssembly.cs" })]
+    [SetupCompileBefore("library2.dll", new[] { "Dependencies/TypeMapSecondOrderReference.cs" })]
+    [Kept]
+    [KeptAssembly("library.dll")]
+    [KeptAssembly("library2.dll")]
+    [KeptTypeInAssembly("library.dll", typeof(TypeMapReferencedAssembly))]
+    [KeptMemberInAssembly("library.dll", typeof(TypeMapReferencedAssembly), "Main()")]
+    [KeptTypeInAssembly("library.dll", typeof(TargetTypeUnconditional1), Tool = Tool.Trimmer)]
+    [KeptTypeInAssembly("library.dll", typeof(TrimTarget1))]
+    [KeptMemberInAssembly("library.dll", typeof(TrimTarget1), ".ctor()")]
+    [KeptTypeInAssembly("library.dll", typeof(TrimTarget2))]
+    [KeptMemberInAssembly("library.dll", typeof(TrimTarget2), ".ctor()")]
+    [KeptTypeInAssembly("library.dll", typeof(TargetTypeConditional1), Tool = Tool.Trimmer)]
+    [KeptTypeInAssembly("library.dll", typeof(ProxySource1))]
+    [KeptMemberInAssembly("library.dll", typeof(ProxySource1), ".ctor()")]
+    [KeptTypeInAssembly("library.dll", typeof(ProxyTarget1))]
+
+    [KeptAttributeInAssembly("library.dll", typeof(TypeMapAttribute<UsedTypeMapUniverse>))]
+    [KeptAttributeInAssembly("library.dll", typeof(TypeMapAssociationAttribute<UsedTypeMapUniverse>))]
+    [KeptAttributeInAssembly("library.dll", typeof(TypeMapAssemblyTargetAttribute<UsedTypeMapUniverse>))]
+    [RemovedAttributeInAssembly("library.dll", typeof(TypeMapAttribute<UnusedTypeMapUniverse>))]
+    [RemovedAttributeInAssembly("library.dll", typeof(TypeMapAssociationAttribute<UnusedTypeMapUniverse>))]
+
+    [KeptAttributeInAssembly("library2.dll", typeof(TypeMapAttribute<string>))]
+    [KeptAttributeInAssembly("library2.dll", typeof(TypeMapAssociationAttribute<string>))]
+    [KeptAttributeInAssembly("library2.dll", typeof(TypeMapAssemblyTargetAttribute<string>))]
+    [RemovedAttributeInAssembly("library2.dll", typeof(TypeMapAttribute<UnusedTypeMapUniverse2>))]
+    [RemovedAttributeInAssembly("library2.dll", typeof(TypeMapAssociationAttribute<UnusedTypeMapUniverse2>))]
+    // No types kept in library2, just TypeMap attributes
+    [RemovedTypeInAssembly("library2.dll", typeof(TypeMapReferencedAssembly2))]
+    [RemovedTypeInAssembly("library2.dll", typeof(UsedTypeMapUniverse2))]
+    [RemovedTypeInAssembly("library2.dll", typeof(Mono.Linker.Tests.Cases.Reflection.Dependencies.Library2.ProxySource1))]
+    [RemovedTypeInAssembly("library2.dll", typeof(Mono.Linker.Tests.Cases.Reflection.Dependencies.Library2.ProxySource2))]
+    [RemovedTypeInAssembly("library2.dll", typeof(Mono.Linker.Tests.Cases.Reflection.Dependencies.Library2.ProxyTarget1))]
+    [RemovedTypeInAssembly("library2.dll", typeof(Mono.Linker.Tests.Cases.Reflection.Dependencies.Library2.ProxyTarget2))]
+    [RemovedTypeInAssembly("library2.dll", typeof(Mono.Linker.Tests.Cases.Reflection.Dependencies.Library2.TargetTypeConditional1))]
+    [RemovedTypeInAssembly("library2.dll", typeof(Mono.Linker.Tests.Cases.Reflection.Dependencies.Library2.TargetTypeConditional2))]
+    [RemovedTypeInAssembly("library2.dll", typeof(Mono.Linker.Tests.Cases.Reflection.Dependencies.Library2.TargetTypeUnconditional1))]
+    [RemovedTypeInAssembly("library2.dll", typeof(Mono.Linker.Tests.Cases.Reflection.Dependencies.Library2.TargetTypeUnconditional2))]
+    [RemovedTypeInAssembly("library2.dll", typeof(Mono.Linker.Tests.Cases.Reflection.Dependencies.Library2.TrimTarget1))]
+    [RemovedTypeInAssembly("library2.dll", typeof(Mono.Linker.Tests.Cases.Reflection.Dependencies.Library2.TrimTarget2))]
     class TypeMap
     {
         [Kept]
@@ -114,6 +178,38 @@ static void ConstrainedStaticCall<T>(T t) where T : IStaticInterface
             Console.WriteLine(new ArrayElement[1]);
 
             Console.WriteLine(new ConstructedNoTypeCheckNoBoxStruct(42).Value);
+
+            TypeMapReferencedAssembly.Main();
+
+            // TypeMapUniverses are independent between External and Proxy type maps.
+            // That is, if the External type map is used for a given universe, that doesn't keep the Proxy type map, and vice versa.
+            // Since we only use UsedExternalTypeMap for its External type map, the Proxy type map and its attributes should be removed.
+            // And vice versa for UsedProxyTypeMap.
+            _ = TypeMapping.GetOrCreateExternalTypeMapping<UsedExternalTypeMap>();
+            _ = TypeMapping.GetOrCreateProxyTypeMapping<UsedProxyTypeMap>();
+            _ = new UsedProxySource();
+            _ = new UsedProxySource2();
+
+            UseTrimTargets();
+
+            // For the second order reference, instantiate int and string as the dependency source types to root the TypeMap attributes
+            // Do not directly root anything in the assembly to validate the assembly is kept even in only the typemap attributes are marked.
+            _ = new string('h', 2);
+            _ = new int();
+            _ = TypeMapping.GetOrCreateExternalTypeMapping<string>();
+            _ = TypeMapping.GetOrCreateProxyTypeMapping<string>();
+        }
+
+        [ExpectBodyModified]
+        [Kept]
+        static void UseTrimTargets()
+        {
+            object obj = null!;
+            // Rewritten as if these types are removed
+            if (obj is UsedTrimTarget)
+                obj = 1;
+            if (obj is UsedTrimTarget2)
+                obj = 2;
         }
 
         [Kept]
@@ -180,7 +276,7 @@ private static void CheckTrimTarget(object o)
         [Kept]
         private static UnboxedOnly Unbox(object o)
         {
-            return (UnboxedOnly) o;
+            return (UnboxedOnly)o;
         }
 
         [Kept]
@@ -216,6 +312,7 @@ class ProxyType;
 
     [Kept]
     class UnusedTypeMap;
+    class UnusedTypeMap2;
     class UnusedTargetType;
     class UnusedSourceClass;
     class UnusedProxyType;
@@ -400,57 +497,29 @@ public AllocatedNoBoxStructType(int value)
 
     [Kept]
     class AllocatedNoBoxProxy;
-}
 
-// Polyfill for the type map types until we use an LKG runtime that has them with an updated LinkAttributes XML.
-namespace System.Runtime.InteropServices
-{
-    [Kept(By = Tool.Trimmer)]
-    [KeptBaseType(typeof(Attribute), By = Tool.Trimmer)]
-    [KeptAttributeAttribute(typeof(AttributeUsageAttribute), By = Tool.Trimmer)]
-    [KeptAttributeAttribute(typeof(RemoveAttributeInstancesAttribute), By = Tool.Trimmer)]
-    [RemoveAttributeInstances]
-    [AttributeUsage(AttributeTargets.Assembly, AllowMultiple = true)]
-    public sealed class TypeMapAttribute<TTypeMapGroup> : Attribute
-    {
-        [Kept(By = Tool.Trimmer)]
-        public TypeMapAttribute(string value, Type target) { }
+    [Kept]
+    class UsedExternalTypeMap;
+    [Kept]
+    class UsedProxyTypeMap;
 
-        [Kept(By = Tool.Trimmer)]
-        [KeptAttributeAttribute(typeof(RequiresUnreferencedCodeAttribute), By = Tool.Trimmer)]
-        [RequiresUnreferencedCode("Interop types may be removed by trimming")]
-        public TypeMapAttribute(string value, Type target, Type trimTarget) { }
-    }
+    [Kept]
+    [KeptMember(".ctor()")]
+    class UsedProxySource;
+    [Kept]
+    [KeptMember(".ctor()")]
+    class UsedProxySource2;
+    [Kept]
+    class UsedTrimTarget;
+    [Kept(By=Tool.NativeAot)] // Associated attribute not kept
+    class UsedTrimTarget2;
 
-    [Kept(By = Tool.Trimmer)]
-    [KeptBaseType(typeof(Attribute), By = Tool.Trimmer)]
-    [KeptAttributeAttribute(typeof(AttributeUsageAttribute), By = Tool.Trimmer)]
-    [KeptAttributeAttribute(typeof(RemoveAttributeInstancesAttribute), By = Tool.Trimmer)]
-    [RemoveAttributeInstances]
-    [AttributeUsage(AttributeTargets.Assembly, AllowMultiple = true)]
-    public sealed class TypeMapAssociationAttribute<TTypeMapGroup> : Attribute
-    {
-        [Kept(By = Tool.Trimmer)]
-        public TypeMapAssociationAttribute(Type source, Type proxy) { }
-    }
+    [Kept]
+    class UsedExternalTarget;
+    class UsedExternalTarget2;
 
-    [Kept(By = Tool.Trimmer)]
-    public static class TypeMapping
-    {
-        [Kept(By = Tool.Trimmer)]
-        [KeptAttributeAttribute(typeof(RequiresUnreferencedCodeAttribute), By = Tool.Trimmer)]
-        [RequiresUnreferencedCode("Interop types may be removed by trimming")]
-        public static IReadOnlyDictionary<string, Type> GetOrCreateExternalTypeMapping<TTypeMapGroup>()
-        {
-            throw new NotImplementedException($"External type map for {typeof(TTypeMapGroup).Name}");
-        }
 
-        [Kept(By = Tool.Trimmer)]
-        [KeptAttributeAttribute(typeof(RequiresUnreferencedCodeAttribute), By = Tool.Trimmer)]
-        [RequiresUnreferencedCode("Interop types may be removed by trimming")]
-        public static IReadOnlyDictionary<Type, Type> GetOrCreateProxyTypeMapping<TTypeMapGroup>()
-        {
-            throw new NotImplementedException($"Proxy type map for {typeof(TTypeMapGroup).Name}");
-        }
-    }
+    class UsedProxyTarget;
+    [Kept]
+    class UsedProxyTarget2;
 }