Support for abi-specific assemblies in marshal methods assembly rewriter

src/Xamarin.Android.Build.Tasks/Tasks/GenerateJavaStubs.cs

@@ -4,6 +4,7 @@
 using System.Collections.Generic;
 using System.Diagnostics;
 using System.IO;
+using System.IO.MemoryMappedFiles;
 using System.Linq;
 using System.Reflection;
 using System.Text;
@@ -145,9 +146,9 @@ void Run (DirectoryAssemblyResolver res, bool useMarshalMethods)
 			PackageNamingPolicy pnp;
 			JavaNativeTypeManager.PackageNamingPolicy = Enum.TryParse (PackageNamingPolicy, out pnp) ? pnp : PackageNamingPolicyEnum.LowercaseCrc64;
 
-			Dictionary<string, HashSet<string>> marshalMethodsAssemblyPaths = null;
+			Dictionary<string, List<ITaskItem>>? abiSpecificAssembliesByPath = null;
 			if (useMarshalMethods) {
-				marshalMethodsAssemblyPaths = new Dictionary<string, HashSet<string>> (StringComparer.Ordinal);
+				abiSpecificAssembliesByPath = new Dictionary<string, List<ITaskItem>> (StringComparer.Ordinal);
 			}
 
 			// Put every assembly we'll need in the resolver
@@ -162,8 +163,20 @@ void Run (DirectoryAssemblyResolver res, bool useMarshalMethods)
 					continue;
 				}
 
-				bool addAssembly = false;
 				string fileName = Path.GetFileName (assembly.ItemSpec);
+				if (abiSpecificAssembliesByPath != null) {
+					string? abi = assembly.GetMetadata ("Abi");
+					if (!String.IsNullOrEmpty (abi)) {
+						if (!abiSpecificAssembliesByPath.TryGetValue (fileName, out List<ITaskItem>? items)) {
+							items = new List<ITaskItem> ();
+							abiSpecificAssembliesByPath.Add (fileName, items);
+						}
+
+						items.Add (assembly);
+					}
+				}
+
+				bool addAssembly = false;
 				if (!hasExportReference && String.Compare ("Mono.Android.Export.dll", fileName, StringComparison.OrdinalIgnoreCase) == 0) {
 					hasExportReference = true;
 					addAssembly = true;
@@ -184,9 +197,6 @@ void Run (DirectoryAssemblyResolver res, bool useMarshalMethods)
 				}
 
 				res.Load (assembly.ItemSpec);
-				if (useMarshalMethods) {
-					StoreMarshalAssemblyPath (Path.GetFileNameWithoutExtension (assembly.ItemSpec), assembly);
-				}
 			}
 
 			// However we only want to look for JLO types in user code for Java stub code generation
@@ -200,7 +210,6 @@ void Run (DirectoryAssemblyResolver res, bool useMarshalMethods)
 				string name = Path.GetFileNameWithoutExtension (asm.ItemSpec);
 				if (!userAssemblies.ContainsKey (name))
 					userAssemblies.Add (name, asm.ItemSpec);
-				StoreMarshalAssemblyPath (name, asm);
 			}
 
 			// Step 1 - Find all the JLO types
@@ -237,27 +246,11 @@ void Run (DirectoryAssemblyResolver res, bool useMarshalMethods)
 				// in order to properly generate wrapper methods in the marshal methods assembly rewriter.
 				// We don't care about those generated by us, since they won't contain the `XA_BROKEN_EXCEPTION_TRANSITIONS` variable we look for.
 				var environmentParser = new EnvironmentFilesParser ();
-				var targetPaths = new List<string> ();
-
-				if (!LinkingEnabled) {
-					targetPaths.Add (Path.GetDirectoryName (ResolvedAssemblies[0].ItemSpec));
-				} else {
-					if (String.IsNullOrEmpty (IntermediateOutputDirectory)) {
-						throw new InvalidOperationException ($"Internal error: marshal methods require the `IntermediateOutputDirectory` property of the `GenerateJavaStubs` task to have a value");
-					}
 
-					// If the <ResourceIdentifiers> property is set then, even if we have just one RID, the linked assemblies path will include the RID
-					if (!HaveMultipleRIDs && SupportedAbis.Length == 1) {
-						targetPaths.Add (Path.Combine (IntermediateOutputDirectory, "linked"));
-					} else {
-						foreach (string abi in SupportedAbis) {
-							targetPaths.Add (Path.Combine (IntermediateOutputDirectory, AbiToRid (abi), "linked"));
-						}
-					}
-				}
+				Dictionary<AssemblyDefinition, string> assemblyPaths = AddMethodsFromAbiSpecificAssemblies (classifier, res, abiSpecificAssembliesByPath);
 
-				var rewriter = new MarshalMethodsAssemblyRewriter (classifier.MarshalMethods, classifier.Assemblies, marshalMethodsAssemblyPaths, Log);
-				rewriter.Rewrite (res, targetPaths, environmentParser.AreBrokenExceptionTransitionsEnabled (Environments));
+				var rewriter = new MarshalMethodsAssemblyRewriter (classifier.MarshalMethods, classifier.Assemblies, assemblyPaths, Log);
+				rewriter.Rewrite (res, environmentParser.AreBrokenExceptionTransitionsEnabled (Environments));
 			}
 
 			// Step 3 - Generate type maps
@@ -413,40 +406,35 @@ void Run (DirectoryAssemblyResolver res, bool useMarshalMethods)
 					Log.LogDebugMessage ($"Number of methods in the project that need marshal method wrappers: {classifier.WrappedMethodCount}");
 				}
 			}
+		}
 
-			void StoreMarshalAssemblyPath (string name, ITaskItem asm)
-			{
-				if (!useMarshalMethods) {
-					return;
-				}
-
-				// TODO: we need to keep paths to ALL the assemblies, we need to rewrite them for all RIDs eventually. Right now we rewrite them just for one RID
-				if (!marshalMethodsAssemblyPaths.TryGetValue (name, out HashSet<string> assemblyPaths)) {
-					assemblyPaths = new HashSet<string> ();
-					marshalMethodsAssemblyPaths.Add (name, assemblyPaths);
-				}
-
-				assemblyPaths.Add (asm.ItemSpec);
-			}
-
-			string AbiToRid (string abi)
-			{
-				switch (abi) {
-					case "arm64-v8a":
-						return "android-arm64";
+		AssemblyDefinition LoadAssembly (string path, DirectoryAssemblyResolver resolver)
+		{
+			string pdbPath = Path.ChangeExtension (path, ".pdb");
+			var readerParameters = new ReaderParameters {
+				AssemblyResolver                = resolver,
+				InMemory                        = false,
+				ReadingMode                     = ReadingMode.Immediate,
+				ReadSymbols                     = File.Exists (pdbPath),
+				ReadWrite                       = false,
+			};
 
-					case "armeabi-v7a":
-						return "android-arm";
+			MemoryMappedViewStream? viewStream = null;
+			try {
+				// Create stream because CreateFromFile(string, ...) uses FileShare.None which is too strict
+				using var fileStream = new FileStream (path, FileMode.Open, FileAccess.Read, FileShare.Read, 4096, false);
+				using var mappedFile = MemoryMappedFile.CreateFromFile (
+					fileStream, null, fileStream.Length, MemoryMappedFileAccess.Read, HandleInheritability.None, true);
+				viewStream = mappedFile.CreateViewStream (0, 0, MemoryMappedFileAccess.Read);
 
-					case "x86":
-						return "android-x86";
+				AssemblyDefinition result = ModuleDefinition.ReadModule (viewStream, readerParameters).Assembly;
 
-					case "x86_64":
-						return "android-x64";
+				// We transferred the ownership of the viewStream to the collection.
+				viewStream = null;
 
-					default:
-						throw new InvalidOperationException ($"Internal error: unsupported ABI '{abi}'");
-				}
+				return result;
+			} finally {
+				viewStream?.Dispose ();
 			}
 		}
 
@@ -573,5 +561,137 @@ void WriteTypeMappings (List<TypeDefinition> types, TypeDefinitionCache cache)
 			GeneratedBinaryTypeMaps = tmg.GeneratedBinaryTypeMaps.ToArray ();
 			BuildEngine4.RegisterTaskObjectAssemblyLocal (ProjectSpecificTaskObjectKey (ApplicationConfigTaskState.RegisterTaskObjectKey), appConfState, RegisteredTaskObjectLifetime.Build);
 		}
+
+		/// <summary>
+		/// <para>
+		/// Classifier will see only unique assemblies, since that's what's processed by the JI type scanner - even though some assemblies may have
+		/// abi-specific features (e.g. inlined `IntPtr.Size` or processor-specific intrinsics), the **types** and **methods** will all be the same and, thus,
+		/// there's no point in scanning all of the additional copies of the same assembly.
+		/// </para>
+		/// <para>
+		/// This, however, doesn't work for the rewriter which needs to rewrite all of the copies so that they all have the same generated wrappers.  In
+		/// order to do that, we need to go over the list of assemblies found by the classifier, see if they are abi-specific ones and then add all the
+		/// marshal methods from the abi-specific assembly copies, so that the rewriter can easily rewrite them all.
+		/// </para>
+		/// <para>
+		/// This method returns a dictionary matching `AssemblyDefinition` instances to the path on disk to the assembly file they were loaded from.  It is necessary
+		/// because <see cref="LoadAssembly"/> uses a stream to load the data, in order to avoid later sharing violation issues when writing the assemblies.  Path
+		/// information is required by <see cref="MarshalMethodsAssemblyRewriter"/> to be available for each <see cref="MarshalMethodEntry"/>
+		/// </para>
+		/// </summary>
+		Dictionary<AssemblyDefinition, string> AddMethodsFromAbiSpecificAssemblies (MarshalMethodsClassifier classifier, DirectoryAssemblyResolver resolver, Dictionary<string, List<ITaskItem>> abiSpecificAssemblies)
+		{
+			IDictionary<string, IList<MarshalMethodEntry>> marshalMethods = classifier.MarshalMethods;
+			ICollection<AssemblyDefinition> assemblies = classifier.Assemblies;
+			var newAssemblies = new List<AssemblyDefinition> ();
+			var assemblyPaths = new Dictionary<AssemblyDefinition, string> ();
+
+			foreach (AssemblyDefinition asmdef in assemblies) {
+				string fileName = Path.GetFileName (asmdef.MainModule.FileName);
+				if (!abiSpecificAssemblies.TryGetValue (fileName, out List<ITaskItem>? abiAssemblyItems)) {
+					continue;
+				}
+
+				List<MarshalMethodEntry> assemblyMarshalMethods = FindMarshalMethodsForAssembly (marshalMethods, asmdef);;
+				Log.LogDebugMessage ($"Assembly {fileName} is ABI-specific");
+				foreach (ITaskItem abiAssemblyItem in abiAssemblyItems) {
+					if (String.Compare (abiAssemblyItem.ItemSpec, asmdef.MainModule.FileName, StringComparison.Ordinal) == 0) {
+						continue;
+					}
+
+					Log.LogDebugMessage ($"Looking for matching mashal methods in {abiAssemblyItem.ItemSpec}");
+					FindMatchingMethodsInAssembly (abiAssemblyItem, classifier, assemblyMarshalMethods, resolver, newAssemblies, assemblyPaths);
+				}
+			}
+
+			if (newAssemblies.Count > 0) {
+				foreach (AssemblyDefinition asmdef in newAssemblies) {
+					assemblies.Add (asmdef);
+				}
+			}
+
+			return assemblyPaths;
+		}
+
+		List<MarshalMethodEntry> FindMarshalMethodsForAssembly (IDictionary<string, IList<MarshalMethodEntry>> marshalMethods, AssemblyDefinition asm)
+		{
+			var seenNativeCallbacks = new HashSet<MethodDefinition> ();
+			var assemblyMarshalMethods = new List<MarshalMethodEntry> ();
+
+			foreach (var kvp in marshalMethods) {
+				foreach (MarshalMethodEntry method in kvp.Value) {
+					if (method.NativeCallback.Module.Assembly != asm) {
+						continue;
+					}
+
+					// More than one overriden method can use the same native callback method, we're interested only in unique native
+					// callbacks, since that's what gets rewritten.
+					if (seenNativeCallbacks.Contains (method.NativeCallback)) {
+						continue;
+					}
+
+					seenNativeCallbacks.Add (method.NativeCallback);
+					assemblyMarshalMethods.Add (method);
+				}
+			}
+
+			return assemblyMarshalMethods;
+		}
+
+		void FindMatchingMethodsInAssembly (ITaskItem assemblyItem, MarshalMethodsClassifier classifier, List<MarshalMethodEntry> assemblyMarshalMethods, DirectoryAssemblyResolver resolver, List<AssemblyDefinition> newAssemblies, Dictionary<AssemblyDefinition, string> assemblyPaths)
+		{
+			AssemblyDefinition asm = LoadAssembly (assemblyItem.ItemSpec, resolver);
+			newAssemblies.Add (asm);
+			assemblyPaths.Add (asm, assemblyItem.ItemSpec);
+
+			foreach (MarshalMethodEntry methodEntry in assemblyMarshalMethods) {
+				TypeDefinition wantedType = methodEntry.NativeCallback.DeclaringType;
+				TypeDefinition? type = asm.MainModule.FindType (wantedType.FullName);
+				if (type == null) {
+					throw new InvalidOperationException ($"Internal error: type '{wantedType.FullName}' not found in assembly '{assemblyItem.ItemSpec}', a linker error?");
+				}
+
+				if (type.MetadataToken != wantedType.MetadataToken) {
+					throw new InvalidOperationException ($"Internal error: type '{type.FullName}' in assembly '{assemblyItem.ItemSpec}' has a different token ID than the original type");
+				}
+
+				FindMatchingMethodInType (methodEntry, type, classifier);
+			}
+		}
+
+		void FindMatchingMethodInType (MarshalMethodEntry methodEntry, TypeDefinition type, MarshalMethodsClassifier classifier)
+		{
+			string callbackName = methodEntry.NativeCallback.FullName;
+
+			foreach (MethodDefinition typeNativeCallbackMethod in type.Methods) {
+				if (String.Compare (typeNativeCallbackMethod.FullName, callbackName, StringComparison.Ordinal) != 0) {
+					continue;
+				}
+
+				if (typeNativeCallbackMethod.Parameters.Count != methodEntry.NativeCallback.Parameters.Count) {
+					continue;
+				}
+
+				if (typeNativeCallbackMethod.MetadataToken != methodEntry.NativeCallback.MetadataToken) {
+					throw new InvalidOperationException ($"Internal error: tokens don't match for '{typeNativeCallbackMethod.FullName}'");
+				}
+
+				bool allMatch = true;
+				for (int i = 0; i < typeNativeCallbackMethod.Parameters.Count; i++) {
+					if (String.Compare (typeNativeCallbackMethod.Parameters[i].ParameterType.FullName, methodEntry.NativeCallback.Parameters[i].ParameterType.FullName, StringComparison.Ordinal) != 0) {
+						allMatch = false;
+						break;
+					}
+				}
+
+				if (!allMatch) {
+					continue;
+				}
+
+				Log.LogDebugMessage ($"Found match for '{typeNativeCallbackMethod.FullName}' in {type.Module.FileName}");
+				string methodKey = classifier.GetStoreMethodKey (methodEntry);
+				classifier.MarshalMethods[methodKey].Add (new MarshalMethodEntry (methodEntry, typeNativeCallbackMethod));
+			}
+		}
 	}
 }

src/Xamarin.Android.Build.Tasks/Tests/Xamarin.Android.Build.Tests/Tasks/LinkerTests.cs

@@ -1,9 +1,12 @@
 using System;
 using System.IO;
 using System.Linq;
+using System.Text;
 using Java.Interop.Tools.Cecil;
 using Mono.Cecil;
+using Mono.Cecil.Cil;
 using Mono.Linker;
+using Mono.Tuner;
 using MonoDroid.Tuner;
 using NUnit.Framework;
 using Xamarin.ProjectTools;
@@ -514,7 +517,7 @@ public void AndroidUseNegotiateAuthentication ([Values (true, false, null)] bool
 		}
 
 		[Test]
-		public void DoNotErrorOnPerArchJavaTypeDuplicates ()
+		public void DoNotErrorOnPerArchJavaTypeDuplicates ([Values(true, false)] bool enableMarshalMethods)
 		{
 			if (!Builder.UseDotNet)
 				Assert.Ignore ("Test only valid on .NET");
@@ -525,26 +528,73 @@ public void DoNotErrorOnPerArchJavaTypeDuplicates ()
 			lib.Sources.Add (new BuildItem.Source ("Library1.cs") {
 				TextContent = () => @"
 namespace Lib1;
-public class Library1 : Java.Lang.Object {
+public class Library1 : Com.Example.Androidlib.MyRunner {
 	private static bool Is64Bits = IntPtr.Size >= 8;
 
 	public static bool Is64 () {
 		return Is64Bits;
 	}
+
+	public override void Run () => Console.WriteLine (Is64Bits);
 }",
+			});
+			lib.Sources.Add (new BuildItem ("AndroidJavaSource", "MyRunner.java") {
+				Encoding = new UTF8Encoding (encoderShouldEmitUTF8Identifier: false),
+				TextContent = () => @"
+package com.example.androidlib;
+
+public abstract class MyRunner {
+	public abstract void run();
+}"
 			});
 			var proj = new XamarinAndroidApplicationProject { IsRelease = true, ProjectName = "App1" };
 			proj.References.Add(new BuildItem.ProjectReference (Path.Combine ("..", "Lib1", "Lib1.csproj"), "Lib1"));
 			proj.MainActivity = proj.DefaultMainActivity.Replace (
 				"base.OnCreate (bundle);",
 				"base.OnCreate (bundle);\n" +
 				"if (Lib1.Library1.Is64 ()) Console.WriteLine (\"Hello World!\");");
+			proj.SetProperty ("AndroidEnableMarshalMethods", enableMarshalMethods.ToString ());
 
 
 			using var lb = CreateDllBuilder (Path.Combine (path, "Lib1"));
 			using var b = CreateApkBuilder (Path.Combine (path, "App1"));
 			Assert.IsTrue (lb.Build (lib), "build should have succeeded.");
 			Assert.IsTrue (b.Build (proj), "build should have succeeded.");
+
+			var intermediate = Path.Combine (Root, b.ProjectDirectory, proj.IntermediateOutputPath);
+			var dll = $"{lib.ProjectName}.dll";
+			Assert64Bit ("android-arm", expected64: false);
+			Assert64Bit ("android-arm64", expected64: true);
+			Assert64Bit ("android-x86", expected64: false);
+			Assert64Bit ("android-x64", expected64: true);
+
+			void Assert64Bit(string rid, bool expected64)
+			{
+				var assembly = AssemblyDefinition.ReadAssembly (Path.Combine (intermediate, rid, "linked", "shrunk", dll));
+				var type = assembly.MainModule.FindType ("Lib1.Library1");
+				Assert.NotNull (type, "Should find Lib1.Library1!");
+				var cctor = type.GetTypeConstructor ();
+				Assert.NotNull (type, "Should find Lib1.Library1.cctor!");
+				Assert.AreNotEqual (0, cctor.Body.Instructions.Count);
+
+				/*
+				 * IL snippet
+				 * .method private hidebysig specialname rtspecialname static 
+				 * void .cctor () cil managed 
+				 * {
+				 *   // Is64Bits = 4 >= 8;
+				 *   IL_0000: ldc.i4 4
+				 *   IL_0005: ldc.i4.8
+				 *   ...
+				 */
+				var instruction = cctor.Body.Instructions [0];
+				Assert.AreEqual (OpCodes.Ldc_I4, instruction.OpCode);
+				if (expected64) {
+					Assert.AreEqual (8, instruction.Operand, $"Expected 64-bit: {expected64}");
+				} else {
+					Assert.AreEqual (4, instruction.Operand, $"Expected 64-bit: {expected64}");
+				}
+			}
 		}
 	}
 }