RewriteMarshalMethods.cs

#nullable enable
using System;
using System.Collections.Concurrent;
using System.Linq;
using Microsoft.Android.Build.Tasks;
using Microsoft.Build.Framework;
using Xamarin.Android.Tools;

namespace Xamarin.Android.Tasks;

/// <summary>
/// MSBuild task that rewrites .NET assemblies to use marshal methods instead of dynamic JNI registration.
/// This task modifies method implementations to use efficient native callbacks with [UnmanagedCallersOnly]
/// attributes, significantly improving startup performance and reducing runtime overhead for Android applications.
/// </summary>
/// <remarks>
/// This task operates on the marshal method classifications produced by earlier pipeline stages and:
/// 
/// 1. Retrieves marshal method classifications from the build pipeline state
/// 2. Parses environment files to determine exception transition behavior
/// 3. Rewrites assemblies to replace dynamic registration with static marshal methods
/// 4. Optionally builds managed lookup tables for runtime marshal method resolution
/// 5. Reports statistics on marshal method generation and any fallback to dynamic registration
/// 
/// The rewriting process creates native callback wrappers for methods that have non-blittable
/// parameters or return types, ensuring compatibility with the [UnmanagedCallersOnly] attribute
/// while maintaining proper marshaling semantics.
/// </remarks>
public class RewriteMarshalMethods : AndroidTask
{
	/// <summary>
	/// Gets the task prefix used for logging and error messages.
	/// </summary>
	public override string TaskPrefix => "RMM";

	/// <summary>
	/// Gets or sets whether to enable managed marshal methods lookup tables.
	/// When enabled, generates runtime lookup structures that allow dynamic resolution
	/// of marshal methods without string comparisons, improving runtime performance.
	/// </summary>
	public bool EnableManagedMarshalMethodsLookup { get; set; }

	/// <summary>
	/// Gets or sets the environment files to parse for configuration settings.
	/// These files may contain settings like XA_BROKEN_EXCEPTION_TRANSITIONS that
	/// affect how marshal method wrappers are generated.
	/// </summary>
	public ITaskItem [] Environments { get; set; } = [];

	/// <summary>
	/// Gets or sets the intermediate output directory path. Required for retrieving
	/// build state objects that contain marshal method classifications.
	/// </summary>
	[Required]
	public string IntermediateOutputDirectory { get; set; } = "";

	/// <summary>
	/// Executes the marshal method rewriting task. This is the main entry point that
	/// coordinates the entire assembly rewriting process across all target architectures.
	/// </summary>
	/// <returns>
	/// true if the task completed successfully; false if errors occurred during processing.
	/// </returns>
	/// <remarks>
	/// The execution flow is:
	/// 
	/// 1. Retrieve native code generation state from previous pipeline stages
	/// 2. Parse environment files for configuration (e.g., broken exception transitions)
	/// 3. For each target architecture:
	///    - Rewrite assemblies to use marshal methods
	///    - Add special case methods (e.g., TypeManager methods)
	///    - Optionally build managed lookup tables
	/// 4. Report statistics on marshal method generation
	/// 5. Log warnings for methods that must fall back to dynamic registration
	/// 
	/// The task handles the ordering dependency between special case methods and managed
	/// lookup tables - special cases must be added first so they appear in the lookup tables.
	/// </remarks>
	public override bool RunTask ()
	{
		// Retrieve the stored NativeCodeGenState from the build pipeline
		// This contains marshal method classifications from earlier stages
		var nativeCodeGenStates = BuildEngine4.GetRegisteredTaskObjectAssemblyLocal<ConcurrentDictionary<AndroidTargetArch, NativeCodeGenState>> (
			MonoAndroidHelper.GetProjectBuildSpecificTaskObjectKey (GenerateJavaStubs.NativeCodeGenStateRegisterTaskKey, WorkingDirectory, IntermediateOutputDirectory),
			RegisteredTaskObjectLifetime.Build
		);

		if (nativeCodeGenStates is null)
			throw new InvalidOperationException ($"Internal error: {nameof (NativeCodeGenState)} not found");

		// Parse environment files to determine configuration settings
		// We need to parse the environment files supplied by the user to see if they want to use broken exception transitions. This information is needed
		// 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 ();
		bool brokenExceptionTransitionsEnabled = environmentParser.AreBrokenExceptionTransitionsEnabled (Environments);

		// Process each target architecture
		foreach (var kvp in nativeCodeGenStates) {
			NativeCodeGenState state = kvp.Value;

			if (state.Classifier is null) {
				Log.LogError ("state.Classifier cannot be null if marshal methods are enabled");
				return false;
			}

			// Handle the ordering dependency between special case methods and managed lookup tables
			if (!EnableManagedMarshalMethodsLookup) {
				// Standard path: rewrite first, then add special cases
				RewriteMethods (state, brokenExceptionTransitionsEnabled);
				state.Classifier.AddSpecialCaseMethods ();
			} else {
				// Managed lookup path: add special cases first so they appear in lookup tables
				// We need to run `AddSpecialCaseMethods` before `RewriteMarshalMethods` so that we can see the special case
				// methods (such as TypeManager.n_Activate_mm) when generating the managed lookup tables.
				state.Classifier.AddSpecialCaseMethods ();
				state.ManagedMarshalMethodsLookupInfo = new ManagedMarshalMethodsLookupInfo (Log);
				RewriteMethods (state, brokenExceptionTransitionsEnabled);
			}

			// Report statistics on marshal method generation
			Log.LogDebugMessage ($"[{state.TargetArch}] Number of generated marshal methods: {state.Classifier.MarshalMethods.Count}");
			if (state.Classifier.DynamicallyRegisteredMarshalMethods.Count > 0) {
				Log.LogWarning ($"[{state.TargetArch}] Number of methods in the project that will be registered dynamically: {state.Classifier.DynamicallyRegisteredMarshalMethods.Count}");
			}

			// Count and report methods that need blittable workaround wrappers
			var wrappedCount = state.Classifier.MarshalMethods.Sum (m => m.Value.Count (m2 => m2.NeedsBlittableWorkaround));

			if (wrappedCount > 0) {
				// TODO: change to LogWarning once the generator can output code which requires no non-blittable wrappers
				Log.LogDebugMessage ($"[{state.TargetArch}] Number of methods in the project that need marshal method wrappers: {wrappedCount}");
			}
		}

		return !Log.HasLoggedErrors;
	}

	/// <summary>
	/// Performs the actual assembly rewriting for a specific target architecture.
	/// Creates and executes the <see cref="MarshalMethodsAssemblyRewriter"/> that handles
	/// the low-level assembly modification operations.
	/// </summary>
	/// <param name="state">The native code generation state containing marshal method classifications and resolver.</param>
	/// <param name="brokenExceptionTransitionsEnabled">
	/// Whether to generate code compatible with broken exception transitions.
	/// This affects how wrapper methods handle exceptions during JNI calls.
	/// </param>
	/// <remarks>
	/// This method delegates the complex assembly rewriting logic to the specialized
	/// <see cref="MarshalMethodsAssemblyRewriter"/> class, which handles:
	/// - Adding [UnmanagedCallersOnly] attributes to native callbacks
	/// - Generating wrapper methods for non-blittable types
	/// - Modifying assembly references and imports
	/// - Building managed lookup table entries
	/// </remarks>
	void RewriteMethods (NativeCodeGenState state, bool brokenExceptionTransitionsEnabled)
	{
		if (state.Classifier == null) {
			return;
		}

		var rewriter = new MarshalMethodsAssemblyRewriter (Log, state.TargetArch, state.Classifier, state.Resolver, state.ManagedMarshalMethodsLookupInfo);
		rewriter.Rewrite (brokenExceptionTransitionsEnabled);
	}
}