[NativeAOT] Objective-C: SetMessageSendPendingException and SetMessageSendCallback (#77956)

Co-authored-by: Jan Kotas <jkotas@microsoft.com>

10 files changed, 160 insertions, 4 deletions

diff --git a/src/coreclr/nativeaot/Directory.Build.props b/src/coreclr/nativeaot/Directory.Build.props
index 953cf276bee..94ee1363ee0 100644
--- a/src/coreclr/nativeaot/Directory.Build.props
+++ b/src/coreclr/nativeaot/Directory.Build.props
@@ -59,6 +59,13 @@
     <FeatureComWrappers>false</FeatureComWrappers>
     <FeatureComWrappers Condition="'$(TargetsWindows)' == 'true'">true</FeatureComWrappers>
   </PropertyGroup>
+  <PropertyGroup>
+    <FeatureObjCMarshal>false</FeatureObjCMarshal>
+    <FeatureObjCMarshal Condition="'$(TargetsOSX)' == 'true'">true</FeatureObjCMarshal>
+  </PropertyGroup>
+  <PropertyGroup>
+    <DefineConstants Condition="'$(FeatureObjCMarshal)' == 'true'">FEATURE_OBJCMARSHAL;$(DefineConstants)</DefineConstants>
+  </PropertyGroup>
 
   <!-- Platform specific properties -->
   <PropertyGroup Condition="'$(Platform)' == 'x64'">
diff --git a/src/coreclr/nativeaot/System.Private.CoreLib/src/Internal/Runtime/CompilerHelpers/InteropHelpers.cs b/src/coreclr/nativeaot/System.Private.CoreLib/src/Internal/Runtime/CompilerHelpers/InteropHelpers.cs
index 8e383b44d85..f5367d848d4 100644
--- a/src/coreclr/nativeaot/System.Private.CoreLib/src/Internal/Runtime/CompilerHelpers/InteropHelpers.cs
+++ b/src/coreclr/nativeaot/System.Private.CoreLib/src/Internal/Runtime/CompilerHelpers/InteropHelpers.cs
@@ -9,6 +9,7 @@ using System.Reflection;
 using System.Runtime;
 using System.Runtime.CompilerServices;
 using System.Runtime.InteropServices;
+using System.Runtime.InteropServices.ObjectiveC;
 using System.Runtime.Loader;
 using System.Text;
 using System.Threading;
@@ -342,6 +343,17 @@ namespace Internal.Runtime.CompilerHelpers
             byte* methodName = (byte*)pCell->MethodName;
             IntPtr pTarget;
 
+#if FEATURE_OBJCMARSHAL
+#pragma warning disable CA1416
+            if (pCell->IsObjectiveCMessageSend && ObjectiveCMarshal.TryGetGlobalMessageSendCallback(pCell->ObjectiveCMessageSendFunction, out pTarget))
+            {
+                Debug.Assert(pTarget != IntPtr.Zero);
+                pCell->Target = pTarget;
+                return;
+            }
+#pragma warning restore CA1416
+#endif
+
 #if TARGET_WINDOWS
             CharSet charSetMangling = pCell->CharSetMangling;
             if (charSetMangling == 0)
@@ -613,7 +625,11 @@ namespace Internal.Runtime.CompilerHelpers
             public IntPtr Target;
             public IntPtr MethodName;
             public ModuleFixupCell* Module;
-            public CharSet CharSetMangling;
+            private int Flags;
+
+            public CharSet CharSetMangling => (CharSet)(Flags & MethodFixupCellFlagsConstants.CharSetMask);
+            public bool IsObjectiveCMessageSend => (Flags & MethodFixupCellFlagsConstants.IsObjectiveCMessageSendMask) != 0;
+            public int ObjectiveCMessageSendFunction => (Flags & MethodFixupCellFlagsConstants.ObjectiveCMessageSendFunctionMask) >> MethodFixupCellFlagsConstants.ObjectiveCMessageSendFunctionShift;
         }
 
         internal unsafe struct CustomMarshallerKey : IEquatable<CustomMarshallerKey>
diff --git a/src/coreclr/nativeaot/System.Private.CoreLib/src/System.Private.CoreLib.csproj b/src/coreclr/nativeaot/System.Private.CoreLib/src/System.Private.CoreLib.csproj
index 0c0e84b3b7d..64926e9eeca 100644
--- a/src/coreclr/nativeaot/System.Private.CoreLib/src/System.Private.CoreLib.csproj
+++ b/src/coreclr/nativeaot/System.Private.CoreLib/src/System.Private.CoreLib.csproj
@@ -212,6 +212,7 @@
     <Compile Include="System\Runtime\InteropServices\Marshal.NativeAot.cs" />
     <Compile Include="System\Runtime\InteropServices\Marshal.Com.cs" Condition="'$(FeatureCominterop)' == 'true'" />
     <Compile Include="System\Runtime\InteropServices\MemoryMarshal.NativeAot.cs" />
+    <Compile Include="System\Runtime\InteropServices\ObjectiveCMarshal.NativeAot.cs" Condition="'$(FeatureObjCMarshal)' == 'true'" />
     <Compile Include="System\Runtime\InteropServices\UnsafeGCHandle.cs" />
     <Compile Include="System\Runtime\Intrinsics\X86\X86Base.NativeAot.cs" Condition="'$(SupportsX86Intrinsics)' == 'true'" />
     <Compile Include="System\Runtime\JitInfo.NativeAot.cs" />
diff --git a/src/coreclr/nativeaot/System.Private.CoreLib/src/System/Runtime/InteropServices/ObjectiveCMarshal.NativeAot.cs b/src/coreclr/nativeaot/System.Private.CoreLib/src/System/Runtime/InteropServices/ObjectiveCMarshal.NativeAot.cs
new file mode 100644
index 00000000000..a7ab0249acb
--- /dev/null
+++ b/src/coreclr/nativeaot/System.Private.CoreLib/src/System/Runtime/InteropServices/ObjectiveCMarshal.NativeAot.cs
@@ -0,0 +1,70 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System;
+using System.Diagnostics;
+using System.Runtime.ExceptionServices;
+using System.Threading;
+
+namespace System.Runtime.InteropServices.ObjectiveC
+{
+    public static unsafe partial class ObjectiveCMarshal
+    {
+        private static readonly IntPtr[] s_ObjcMessageSendFunctions = new IntPtr[(int)MessageSendFunction.MsgSendSuperStret + 1];
+
+        [ThreadStatic]
+        private static Exception? t_pendingExceptionObject;
+
+        /// <summary>
+        /// Sets a pending exception to be thrown the next time the runtime is entered from an Objective-C msgSend P/Invoke.
+        /// </summary>
+        /// <param name="exception">The exception.</param>
+        /// <remarks>
+        /// If <c>null</c> is supplied any pending exception is discarded.
+        /// </remarks>
+        public static void SetMessageSendPendingException(Exception? exception)
+        {
+            t_pendingExceptionObject = exception;
+        }
+
+        private static bool TrySetGlobalMessageSendCallback(
+            MessageSendFunction msgSendFunction,
+            IntPtr func)
+        {
+            return Interlocked.CompareExchange(ref s_ObjcMessageSendFunctions[(int)msgSendFunction], func, IntPtr.Zero) == IntPtr.Zero;
+        }
+
+        internal static bool TryGetGlobalMessageSendCallback(int msgSendFunction, out IntPtr func)
+        {
+            func = s_ObjcMessageSendFunctions[msgSendFunction];
+            return func != IntPtr.Zero;
+        }
+
+        [StackTraceHidden]
+        internal static void ThrowPendingExceptionObject()
+        {
+            Exception? ex = t_pendingExceptionObject;
+            if (ex != null)
+            {
+                t_pendingExceptionObject = null;
+                ExceptionDispatchInfo.Throw(ex);
+            }
+        }
+
+        private static bool TryInitializeReferenceTracker(
+            delegate* unmanaged<void> beginEndCallback,
+            delegate* unmanaged<IntPtr, int> isReferencedCallback,
+            delegate* unmanaged<IntPtr, void> trackedObjectEnteredFinalization)
+        {
+            throw new NotImplementedException();
+        }
+
+        private static IntPtr CreateReferenceTrackingHandleInternal(
+            object obj,
+            out int memInSizeT,
+            out IntPtr mem)
+        {
+            throw new NotImplementedException();
+        }
+    }
+}
diff --git a/src/coreclr/tools/Common/Internal/Runtime/RuntimeConstants.cs b/src/coreclr/tools/Common/Internal/Runtime/RuntimeConstants.cs
index f01bd3f30fe..24d4076b202 100644
--- a/src/coreclr/tools/Common/Internal/Runtime/RuntimeConstants.cs
+++ b/src/coreclr/tools/Common/Internal/Runtime/RuntimeConstants.cs
@@ -78,4 +78,15 @@ namespace Internal.Runtime
         CastClass,
         AllocateArray,
     }
+
+    /// <summary>
+    /// Constants that describe the bits of the Flags field of MethodFixupCell.
+    /// </summary>
+    internal static class MethodFixupCellFlagsConstants
+    {
+        public const int CharSetMask = 0x7;
+        public const int IsObjectiveCMessageSendMask = 0x8;
+        public const int ObjectiveCMessageSendFunctionMask = 0x70;
+        public const int ObjectiveCMessageSendFunctionShift = 4;
+    }
 }
diff --git a/src/coreclr/tools/Common/TypeSystem/IL/Stubs/PInvokeILEmitter.cs b/src/coreclr/tools/Common/TypeSystem/IL/Stubs/PInvokeILEmitter.cs
index 57f724815fb..b1bbb25130f 100644
--- a/src/coreclr/tools/Common/TypeSystem/IL/Stubs/PInvokeILEmitter.cs
+++ b/src/coreclr/tools/Common/TypeSystem/IL/Stubs/PInvokeILEmitter.cs
@@ -345,6 +345,13 @@ namespace Internal.IL.Stubs
                             InteropTypes.GetPInvokeMarshal(context)
                             .GetKnownMethod("SaveLastError", null)));
             }
+
+            if (MarshalHelpers.ShouldCheckForPendingException(context.Target, _pInvokeMetadata))
+            {
+                MetadataType lazyHelperType = context.SystemModule.GetKnownType("System.Runtime.InteropServices.ObjectiveC", "ObjectiveCMarshal");
+                callsiteSetupCodeStream.Emit(ILOpcode.call, emitter.NewToken(lazyHelperType
+                    .GetKnownMethod("ThrowPendingExceptionObject", null)));
+            }
         }
 
         private void EmitCalli(PInvokeILCodeStreams ilCodeStreams, CalliMarshallingMethodThunk calliThunk)
diff --git a/src/coreclr/tools/Common/TypeSystem/Interop/IL/MarshalHelpers.cs b/src/coreclr/tools/Common/TypeSystem/Interop/IL/MarshalHelpers.cs
index ebcaee6ddfe..91a98346037 100644
--- a/src/coreclr/tools/Common/TypeSystem/Interop/IL/MarshalHelpers.cs
+++ b/src/coreclr/tools/Common/TypeSystem/Interop/IL/MarshalHelpers.cs
@@ -3,6 +3,7 @@
 
 using System;
 using Debug = System.Diagnostics.Debug;
+using System.Runtime.InteropServices.ObjectiveC;
 using Internal.TypeSystem.Ecma;
 
 namespace Internal.TypeSystem.Interop
@@ -920,12 +921,13 @@ namespace Internal.TypeSystem.Interop
             return MarshallerKind.Invalid;
         }
 
+        private const string ObjectiveCLibrary = "/usr/lib/libobjc.dylib";
+
         internal static bool ShouldCheckForPendingException(TargetDetails target, PInvokeMetadata metadata)
         {
             if (!target.IsOSX)
                 return false;
 
-            const string ObjectiveCLibrary = "/usr/lib/libobjc.dylib";
             const string ObjectiveCMsgSend = "objc_msgSend";
 
             // This is for the objc_msgSend suite of functions.
@@ -938,6 +940,24 @@ namespace Internal.TypeSystem.Interop
                 && metadata.Name.StartsWith(ObjectiveCMsgSend);
         }
 
+        internal static int? GetObjectiveCMessageSendFunction(TargetDetails target, string pinvokeModule, string pinvokeFunction)
+        {
+            if (!target.IsOSX || pinvokeModule != ObjectiveCLibrary)
+                return null;
+
+#pragma warning disable CA1416
+            return pinvokeFunction switch
+            {
+                "objc_msgSend" => (int)ObjectiveCMarshal.MessageSendFunction.MsgSend,
+                "objc_msgSend_fpret" => (int)ObjectiveCMarshal.MessageSendFunction.MsgSendFpret,
+                "objc_msgSend_stret" => (int)ObjectiveCMarshal.MessageSendFunction.MsgSendStret,
+                "objc_msgSendSuper" => (int)ObjectiveCMarshal.MessageSendFunction.MsgSendSuper,
+                "objc_msgSendSuper_stret" => (int)ObjectiveCMarshal.MessageSendFunction.MsgSendSuperStret,
+                _ => null,
+            };
+#pragma warning restore CA1416
+        }
+
         public static bool IsRuntimeMarshallingEnabled(ModuleDesc module)
         {
             return module.Assembly is not EcmaAssembly assembly || !assembly.HasAssemblyCustomAttribute("System.Runtime.CompilerServices", "DisableRuntimeMarshallingAttribute");
diff --git a/src/coreclr/tools/aot/ILCompiler.Compiler/Compiler/DependencyAnalysis/PInvokeMethodFixupNode.cs b/src/coreclr/tools/aot/ILCompiler.Compiler/Compiler/DependencyAnalysis/PInvokeMethodFixupNode.cs
index 153f83ae080..92240c4e8c6 100644
--- a/src/coreclr/tools/aot/ILCompiler.Compiler/Compiler/DependencyAnalysis/PInvokeMethodFixupNode.cs
+++ b/src/coreclr/tools/aot/ILCompiler.Compiler/Compiler/DependencyAnalysis/PInvokeMethodFixupNode.cs
@@ -2,12 +2,15 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System;
+using System.Diagnostics;
 using System.Runtime.InteropServices;
 
 using Internal.IL.Stubs;
+using Internal.Runtime;
 using Internal.Text;
 using Internal.TypeSystem;
 using Internal.TypeSystem.Ecma;
+using Internal.TypeSystem.Interop;
 
 namespace ILCompiler.DependencyAnalysis
 {
@@ -73,7 +76,25 @@ namespace ILCompiler.DependencyAnalysis
             // Module fixup cell
             builder.EmitPointerReloc(factory.PInvokeModuleFixup(_pInvokeMethodData.ModuleData));
 
-            builder.EmitInt((int)_pInvokeMethodData.CharSetMangling);
+            int flags = 0;
+
+            int charsetFlags = (int)_pInvokeMethodData.CharSetMangling;
+            Debug.Assert((charsetFlags & MethodFixupCellFlagsConstants.CharSetMask) == charsetFlags);
+            charsetFlags &= MethodFixupCellFlagsConstants.CharSetMask;
+            flags |= charsetFlags;
+
+            int? objcFunction = MarshalHelpers.GetObjectiveCMessageSendFunction(factory.Target, _pInvokeMethodData.ModuleData.ModuleName, _pInvokeMethodData.EntryPointName);
+            if (objcFunction.HasValue)
+            {
+                flags |= MethodFixupCellFlagsConstants.IsObjectiveCMessageSendMask;
+
+                int objcFunctionFlags = objcFunction.Value << MethodFixupCellFlagsConstants.ObjectiveCMessageSendFunctionShift;
+                Debug.Assert((objcFunctionFlags & MethodFixupCellFlagsConstants.ObjectiveCMessageSendFunctionMask) == objcFunctionFlags);
+                objcFunctionFlags &= MethodFixupCellFlagsConstants.ObjectiveCMessageSendFunctionMask;
+                flags |= objcFunctionFlags;
+            }
+
+            builder.EmitInt(flags);
 
             return builder.ToObjectData();
         }
diff --git a/src/libraries/System.Private.CoreLib/src/System/Runtime/InteropServices/ObjectiveC/ObjectiveCMarshal.cs b/src/libraries/System.Private.CoreLib/src/System/Runtime/InteropServices/ObjectiveC/ObjectiveCMarshal.cs
index cbd6108ea4c..c2cb899e115 100644
--- a/src/libraries/System.Private.CoreLib/src/System/Runtime/InteropServices/ObjectiveC/ObjectiveCMarshal.cs
+++ b/src/libraries/System.Private.CoreLib/src/System/Runtime/InteropServices/ObjectiveC/ObjectiveCMarshal.cs
@@ -109,7 +109,11 @@ namespace System.Runtime.InteropServices.ObjectiveC
             ArgumentNullException.ThrowIfNull(obj);
 
             IntPtr refCountHandle = CreateReferenceTrackingHandleInternal(
+#if NATIVEAOT
+                obj,
+#else
                 ObjectHandleOnStack.Create(ref obj),
+#endif
                 out int memInSizeT,
                 out IntPtr mem);
 
diff --git a/src/libraries/System.Runtime.InteropServices/tests/System.Runtime.InteropServices.UnitTests/System/Runtime/InteropServices/ObjectiveC/MessageSendTests.cs b/src/libraries/System.Runtime.InteropServices/tests/System.Runtime.InteropServices.UnitTests/System/Runtime/InteropServices/ObjectiveC/MessageSendTests.cs
index 1493fc255bf..6c45152e8fe 100644
--- a/src/libraries/System.Runtime.InteropServices/tests/System.Runtime.InteropServices.UnitTests/System/Runtime/InteropServices/ObjectiveC/MessageSendTests.cs
+++ b/src/libraries/System.Runtime.InteropServices/tests/System.Runtime.InteropServices.UnitTests/System/Runtime/InteropServices/ObjectiveC/MessageSendTests.cs
@@ -16,7 +16,6 @@ namespace System.Runtime.InteropServices.Tests
 {
     [PlatformSpecific(TestPlatforms.OSX)]
     [SkipOnMono("Not currently implemented on Mono")]
-    [ConditionalClass(typeof(PlatformDetection), nameof(PlatformDetection.IsNotNativeAot))] // https://github.com/dotnet/runtimelab/issues/155
     public unsafe class MessageSendTests
     {
         private static int s_count = 1;