From 2ebb7888f9b4440b8157754babb59694fa42505c Mon Sep 17 00:00:00 2001
From: Rolf Bjarne Kvinge <rolf@xamarin.com>
Date: Thu, 16 Apr 2026 09:19:38 +0200
Subject: [PATCH] Fix illink TypeMap proxy retention. Fixes #127004.

Keep matching TypeMapAssociation entries when an external TypeMap entry preserves the same source type.

Fixes https://github.com/dotnet/runtime/issues/127004.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 .../src/linker/Linker/TypeMapHandler.cs       | 34 +++++++++++++++++++
 .../Reflection/TypeMap.cs                     | 16 +++++++++
 2 files changed, 50 insertions(+)

diff --git a/src/tools/illink/src/linker/Linker/TypeMapHandler.cs b/src/tools/illink/src/linker/Linker/TypeMapHandler.cs
index 5cadd46b65851b..4780cdd8d5db9a 100644
--- a/src/tools/illink/src/linker/Linker/TypeMapHandler.cs
+++ b/src/tools/illink/src/linker/Linker/TypeMapHandler.cs
@@ -108,7 +108,41 @@ void MarkTypeMapAttribute(CustomAttributeWithOrigin entry, DependencyInfo info)
 
             // Mark the target type as instantiated
             if (entry.TargetType is { } targetType && _context.Resolve(targetType) is TypeDefinition targetTypeDef)
+            {
                 _context.Annotations.MarkInstantiated(targetTypeDef);
+                if (entry.Attribute.AttributeType is GenericInstanceType { Name: "TypeMapAttribute`1", GenericArguments: [TypeReference typeMapGroup] })
+                    MarkAssociatedProxyTypeMapAttributes(targetTypeDef, typeMapGroup);
+            }
+        }
+
+        void MarkAssociatedProxyTypeMapAttributes(TypeDefinition sourceType, TypeReference typeMapGroup)
+        {
+            if (_unmarkedProxyTypeMapEntries.TryGetValue(sourceType, out Dictionary<TypeReference, List<CustomAttributeWithOrigin>>? entriesByGroup) &&
+                entriesByGroup.Remove(typeMapGroup, out List<CustomAttributeWithOrigin>? unmarkedAttributes))
+            {
+                foreach (var attr in unmarkedAttributes)
+                    MarkTypeMapAttribute(attr, new DependencyInfo(DependencyKind.TypeMapEntry, sourceType));
+
+                if (entriesByGroup.Count == 0)
+                    _unmarkedProxyTypeMapEntries.Remove(sourceType);
+            }
+
+            if (_pendingProxyTypeMapEntries.TryGetValue(typeMapGroup, out List<CustomAttributeWithOrigin>? pendingAttributes))
+            {
+                for (int i = pendingAttributes.Count - 1; i >= 0; i--)
+                {
+                    CustomAttributeWithOrigin attr = pendingAttributes[i];
+                    if (attr.Attribute.ConstructorArguments is [{ Value: TypeReference pendingSourceType }, _] &&
+                        _context.Resolve(pendingSourceType) == sourceType)
+                    {
+                        pendingAttributes.RemoveAt(i);
+                        MarkTypeMapAttribute(attr, new DependencyInfo(DependencyKind.TypeMapEntry, sourceType));
+                    }
+                }
+
+                if (pendingAttributes.Count == 0)
+                    _pendingProxyTypeMapEntries.Remove(typeMapGroup);
+            }
         }
 
         public void ProcessType(TypeDefinition definition)
diff --git a/src/tools/illink/test/Mono.Linker.Tests.Cases/Reflection/TypeMap.cs b/src/tools/illink/test/Mono.Linker.Tests.Cases/Reflection/TypeMap.cs
index 56867194edec4c..62170062da609b 100644
--- a/src/tools/illink/test/Mono.Linker.Tests.Cases/Reflection/TypeMap.cs
+++ b/src/tools/illink/test/Mono.Linker.Tests.Cases/Reflection/TypeMap.cs
@@ -82,6 +82,12 @@
 [assembly: TypeMapAssemblyTarget<UsedTypeMapUniverse>("library")]
 [assembly: TypeMapAssemblyTarget<UnusedTypeMap2>("library")] // Should be removed
 
+// Verify that a type can be kept if it's used for both TypeMap and TypeMapAssociation
+[assembly: TypeMap<UsedExternalTypeMap>("BothInExternalAndProxy", typeof(BothInExternalAndProxy), typeof(BothInExternalAndProxy))] // Kept
+[assembly: TypeMapAssociation<UsedExternalTypeMap>(typeof(BothInExternalAndProxy), typeof(BothInExternalAndProxyTarget))] // Kept
+[assembly: KeptAttributeAttribute(typeof(TypeMapAttribute<UsedExternalTypeMap>), "BothInExternalAndProxy", typeof(BothInExternalAndProxy), typeof(BothInExternalAndProxy))]
+[assembly: KeptAttributeAttribute(typeof(TypeMapAssociationAttribute<UsedExternalTypeMap>), typeof(BothInExternalAndProxy), typeof(BothInExternalAndProxyTarget))]
+
 namespace Mono.Linker.Tests.Cases.Reflection
 {
     [SetupLinkerAction("link", "System.Private.CoreLib")] // Needed to get the RemoveAttributeInstances in embedded xml
@@ -227,6 +233,9 @@ static void ConstrainedStaticCall<T>(T t) where T : IStaticInterface
             _ = new int();
             _ = TypeMapping.GetOrCreateExternalTypeMapping<string>();
             _ = TypeMapping.GetOrCreateProxyTypeMapping<string>();
+
+            // Use BothInExternalAndProxy in a way that preserves any corresponding typemap entries.
+            Console.WriteLine(new BothInExternalAndProxy());
         }
 
         [ExpectBodyModified]
@@ -554,4 +563,11 @@ class UsedProxyTarget2;
 
     [Kept]
     class PreservedTargetType;
+
+    [Kept]
+    [KeptMember(".ctor()")]
+    class BothInExternalAndProxy;
+    [Kept]
+    class BothInExternalAndProxyTarget;
+
 }