[TrimableTypeMap] Reuse base UCO wrappers for inherited overrides

src/Microsoft.Android.Sdk.TrimmableTypeMap/Generator/Model/TypeMapAssemblyData.cs

@@ -165,7 +165,7 @@ sealed class JavaPeerProxyData
 	public List<UcoConstructorData> UcoConstructors { get; } = new ();
 
 	/// <summary>
-	/// RegisterNatives registrations (method name, JNI signature, wrapper name).
+	/// RegisterNatives registrations (method name, JNI signature, wrapper target).
 	/// </summary>
 	public List<NativeRegistrationData> NativeRegistrations { get; } = new ();
 }
@@ -232,6 +232,24 @@ sealed record UcoMethodData
 	public bool UsesExportMethodDispatch => ExportMethodDispatch != null;
 }
 
+sealed record UcoWrapperTargetData
+{
+	/// <summary>
+	/// Namespace of the generated proxy type containing the wrapper method.
+	/// </summary>
+	public required string TypeNamespace { get; init; }
+
+	/// <summary>
+	/// Name of the generated proxy type containing the wrapper method.
+	/// </summary>
+	public required string TypeName { get; init; }
+
+	/// <summary>
+	/// Name of the UCO wrapper method whose function pointer to register.
+	/// </summary>
+	public required string MethodName { get; init; }
+}
+
 sealed record ExportMethodDispatchData
 {
 	/// <summary>
@@ -329,6 +347,13 @@ sealed record NativeRegistrationData
 	/// Name of the UCO wrapper method whose function pointer to register.
 	/// </summary>
 	public required string WrapperMethodName { get; init; }
+
+	/// <summary>
+	/// Generated proxy wrapper target to register. This may point at a wrapper
+	/// emitted for a different proxy when inherited virtual overrides share the
+	/// same base callback.
+	/// </summary>
+	public required UcoWrapperTargetData WrapperTarget { get; init; }
 }
 
 /// <summary>

src/Microsoft.Android.Sdk.TrimmableTypeMap/Generator/ModelBuilder.cs

@@ -99,6 +99,8 @@ public static TypeMapAssemblyData Build (IReadOnlyList<JavaPeerInfo> peers, stri
 			}
 		}
 
+		BuildNativeRegistrations (model);
+
 		// Compute IgnoresAccessChecksTo from cross-assembly references
 		var referencedAssemblies = new SortedSet<string> (StringComparer.Ordinal);
 		foreach (var proxy in model.ProxyTypes) {
@@ -314,7 +316,6 @@ static JavaPeerProxyData BuildProxyType (JavaPeerInfo peer, string jniName, Hash
 		if (isAcw) {
 			BuildUcoMethods (peer, proxy);
 			BuildUcoConstructors (peer, proxy);
-			BuildNativeRegistrations (proxy);
 		}
 
 		return proxy;
@@ -374,31 +375,104 @@ static void BuildUcoConstructors (JavaPeerInfo peer, JavaPeerProxyData proxy)
 		}
 	}
 
-	static void BuildNativeRegistrations (JavaPeerProxyData proxy)
+	static void BuildNativeRegistrations (TypeMapAssemblyData model)
 	{
-		foreach (var uco in proxy.UcoMethods) {
-			proxy.NativeRegistrations.Add (new NativeRegistrationData {
-				JniMethodName = uco.CallbackMethodName,
-				JniSignature = uco.JniSignature,
-				WrapperMethodName = uco.WrapperName,
-			});
+		var sharedWrapperTargets = new Dictionary<UcoWrapperReuseKey, UcoWrapperTargetData> ();
+		foreach (var proxy in model.ProxyTypes) {
+			foreach (var uco in proxy.UcoMethods) {
+				if (!CanShareUcoWrapper (proxy, uco)) {
+					continue;
+				}
+
+				var reuseKey = CreateUcoWrapperReuseKey (uco);
+				if (!sharedWrapperTargets.ContainsKey (reuseKey)) {
+					sharedWrapperTargets.Add (reuseKey, CreateWrapperTarget (proxy, uco.WrapperName));
+				}
+			}
 		}
 
-		foreach (var uco in proxy.UcoConstructors) {
-			string jniName = uco.WrapperName;
-			int ucoSuffix = jniName.LastIndexOf ("_uco", StringComparison.Ordinal);
-			if (ucoSuffix >= 0) {
-				jniName = jniName.Substring (0, ucoSuffix);
+		foreach (var proxy in model.ProxyTypes) {
+			var reusedUcoMethods = new HashSet<UcoMethodData> ();
+
+			foreach (var uco in proxy.UcoMethods) {
+				var wrapperTarget = CreateWrapperTarget (proxy, uco.WrapperName);
+				if (CanReuseUcoWrapper (proxy, uco) &&
+				    sharedWrapperTargets.TryGetValue (CreateUcoWrapperReuseKey (uco), out var sharedWrapperTarget)) {
+					wrapperTarget = sharedWrapperTarget;
+					reusedUcoMethods.Add (uco);
+				}
+				proxy.NativeRegistrations.Add (new NativeRegistrationData {
+					JniMethodName = uco.CallbackMethodName,
+					JniSignature = uco.JniSignature,
+					WrapperMethodName = wrapperTarget.MethodName,
+					WrapperTarget = wrapperTarget,
+				});
 			}
 
-			proxy.NativeRegistrations.Add (new NativeRegistrationData {
-				JniMethodName = jniName,
-				JniSignature = uco.JniSignature,
-				WrapperMethodName = uco.WrapperName,
-			});
+			if (reusedUcoMethods.Count > 0) {
+				proxy.UcoMethods.RemoveAll (uco => reusedUcoMethods.Contains (uco));
+			}
+
+			foreach (var uco in proxy.UcoConstructors) {
+				string jniName = uco.WrapperName;
+				int ucoSuffix = jniName.LastIndexOf ("_uco", StringComparison.Ordinal);
+				if (ucoSuffix >= 0) {
+					jniName = jniName.Substring (0, ucoSuffix);
+				}
+
+				var wrapperTarget = CreateWrapperTarget (proxy, uco.WrapperName);
+				proxy.NativeRegistrations.Add (new NativeRegistrationData {
+					JniMethodName = jniName,
+					JniSignature = uco.JniSignature,
+					WrapperMethodName = wrapperTarget.MethodName,
+					WrapperTarget = wrapperTarget,
+				});
+			}
 		}
 	}
 
+	static bool CanShareUcoWrapper (JavaPeerProxyData proxy, UcoMethodData uco)
+	{
+		return !uco.UsesExportMethodDispatch &&
+			!proxy.IsGenericDefinition &&
+			!uco.CallbackType.ManagedTypeName.Contains ('`') &&
+			string.Equals (uco.CallbackType.ManagedTypeName, proxy.TargetType.ManagedTypeName, StringComparison.Ordinal) &&
+			string.Equals (uco.CallbackType.AssemblyName, proxy.TargetType.AssemblyName, StringComparison.Ordinal);
+	}
+
+	static bool CanReuseUcoWrapper (JavaPeerProxyData proxy, UcoMethodData uco)
+	{
+		return !uco.UsesExportMethodDispatch &&
+			!proxy.IsGenericDefinition &&
+			!uco.CallbackType.ManagedTypeName.Contains ('`') &&
+			(!string.Equals (uco.CallbackType.ManagedTypeName, proxy.TargetType.ManagedTypeName, StringComparison.Ordinal) ||
+			 !string.Equals (uco.CallbackType.AssemblyName, proxy.TargetType.AssemblyName, StringComparison.Ordinal));
+	}
+
+	static UcoWrapperTargetData CreateWrapperTarget (JavaPeerProxyData proxy, string methodName)
+	{
+		return new UcoWrapperTargetData {
+			TypeNamespace = proxy.Namespace,
+			TypeName = proxy.TypeName,
+			MethodName = methodName,
+		};
+	}
+
+	static UcoWrapperReuseKey CreateUcoWrapperReuseKey (UcoMethodData uco)
+	{
+		return new UcoWrapperReuseKey (
+			uco.CallbackType.ManagedTypeName,
+			uco.CallbackType.AssemblyName,
+			uco.CallbackMethodName,
+			uco.JniSignature);
+	}
+
+	readonly record struct UcoWrapperReuseKey (
+		string CallbackTypeName,
+		string CallbackAssemblyName,
+		string CallbackMethodName,
+		string JniSignature);
+
 	static TypeMapAttributeData BuildEntry (JavaPeerInfo peer, JavaPeerProxyData? proxy,
 		string outputAssemblyName, string jniName)
 	{

src/Microsoft.Android.Sdk.TrimmableTypeMap/Generator/TypeMapAssemblyEmitter.cs

@@ -203,10 +203,10 @@ void EmitCore (TypeMapAssemblyData model, bool useSharedTypemapUniverse)
 		EmitRankSentinels (model);
 		EmitMemberReferences ();
 
-		// Track wrapper method names → handles for RegisterNatives
-		var wrapperHandles = new Dictionary<string, MethodDefinitionHandle> ();
+		// Track wrapper targets → handles for RegisterNatives.
+		var wrapperHandles = new Dictionary<UcoWrapperTargetData, MethodDefinitionHandle> ();
 
-		foreach (var proxy in model.ProxyTypes) {
+		foreach (var proxy in OrderProxiesForWrapperTargets (model.ProxyTypes)) {
 			EmitProxyType (proxy, wrapperHandles);
 		}
 
@@ -225,6 +225,52 @@ void EmitCore (TypeMapAssemblyData model, bool useSharedTypemapUniverse)
 		_pe.EmitIgnoresAccessChecksToAttribute (model.IgnoresAccessChecksTo);
 	}
 
+	static List<JavaPeerProxyData> OrderProxiesForWrapperTargets (IReadOnlyList<JavaPeerProxyData> proxies)
+	{
+		var proxyByType = new Dictionary<(string Namespace, string TypeName), JavaPeerProxyData> ();
+		foreach (var proxy in proxies) {
+			proxyByType [(proxy.Namespace, proxy.TypeName)] = proxy;
+		}
+
+		var ordered = new List<JavaPeerProxyData> (proxies.Count);
+		var states = new Dictionary<JavaPeerProxyData, int> ();
+
+		foreach (var proxy in proxies) {
+			Visit (proxy);
+		}
+
+		return ordered;
+
+		void Visit (JavaPeerProxyData proxy)
+		{
+			if (states.TryGetValue (proxy, out int state)) {
+				if (state == 2) {
+					return;
+				}
+
+				// A cycle would indicate invalid wrapper-target data. Avoid recursing
+				// forever and keep the original relative order for the cyclic edge.
+				return;
+			}
+
+			states [proxy] = 1;
+
+			foreach (var registration in proxy.NativeRegistrations) {
+				var target = registration.WrapperTarget;
+				if (target.TypeNamespace == proxy.Namespace && target.TypeName == proxy.TypeName) {
+					continue;
+				}
+
+				if (proxyByType.TryGetValue ((target.TypeNamespace, target.TypeName), out var targetProxy)) {
+					Visit (targetProxy);
+				}
+			}
+
+			states [proxy] = 2;
+			ordered.Add (proxy);
+		}
+	}
+
 	void EmitTypeReferences ()
 	{
 		var metadata = _pe.Metadata;
@@ -598,7 +644,7 @@ ExportMethodDispatchEmitter GetExportMethodDispatchEmitter ()
 		return _exportMethodDispatchEmitter;
 	}
 
-	void EmitProxyType (JavaPeerProxyData proxy, Dictionary<string, MethodDefinitionHandle> wrapperHandles)
+	void EmitProxyType (JavaPeerProxyData proxy, Dictionary<UcoWrapperTargetData, MethodDefinitionHandle> wrapperHandles)
 	{
 		if (proxy.IsAcw) {
 			// RegisterNatives uses RVA-backed UTF-8 fields under <PrivateImplementationDetails>.
@@ -693,12 +739,12 @@ void EmitProxyType (JavaPeerProxyData proxy, Dictionary<string, MethodDefinition
 			var handle = uco.UsesExportMethodDispatch
 				? GetExportMethodDispatchEmitter ().EmitUcoMethod (uco)
 				: EmitUcoMethod (uco, proxy);
-			wrapperHandles [uco.WrapperName] = handle;
+			wrapperHandles [CreateWrapperTarget (proxy, uco.WrapperName)] = handle;
 		}
 
 		foreach (var uco in proxy.UcoConstructors) {
 			var handle = EmitUcoConstructor (uco, proxy);
-			wrapperHandles [uco.WrapperName] = handle;
+			wrapperHandles [CreateWrapperTarget (proxy, uco.WrapperName)] = handle;
 		}
 
 		// RegisterNatives
@@ -1565,14 +1611,23 @@ void EncodeUcoConstructorLocals_DefaultConstructor (BlobBuilder blob, EntityHand
 		blob.WriteCompressedInteger (CodedIndex.TypeDefOrRefOrSpec (targetTypeRef));
 	}
 
+	static UcoWrapperTargetData CreateWrapperTarget (JavaPeerProxyData proxy, string methodName)
+	{
+		return new UcoWrapperTargetData {
+			TypeNamespace = proxy.Namespace,
+			TypeName = proxy.TypeName,
+			MethodName = methodName,
+		};
+	}
+
 	void EmitRegisterNatives (JavaPeerProxyData proxy,
-		Dictionary<string, MethodDefinitionHandle> wrapperHandles)
+		Dictionary<UcoWrapperTargetData, MethodDefinitionHandle> wrapperHandles)
 	{
 		// Filter to only registrations that have corresponding wrapper methods
 		var registrations = proxy.NativeRegistrations;
 		var validRegs = new List<(NativeRegistrationData Reg, MethodDefinitionHandle Wrapper)> (registrations.Count);
 		foreach (var reg in registrations) {
-			if (wrapperHandles.TryGetValue (reg.WrapperMethodName, out var wrapperHandle)) {
+			if (wrapperHandles.TryGetValue (reg.WrapperTarget, out var wrapperHandle)) {
 				validRegs.Add ((reg, wrapperHandle));
 			}
 		}