Share implementation of ComWrappers between CoreCLR and NativeAOT

src/coreclr/System.Private.CoreLib/System.Private.CoreLib.csproj

@@ -248,7 +248,8 @@
     <Compile Include="src\System\RuntimeType.GenericCache.cs" />
   </ItemGroup>
   <ItemGroup Condition="'$(FeatureComWrappers)' == 'true'">
-    <Compile Include="$(BclSourcesRoot)\System\Runtime\InteropServices\ComWrappers.cs" />
+    <Compile Include="$(BclSourcesRoot)\System\Runtime\InteropServices\ComWrappers.CoreCLR.cs" />
+    <Compile Include="$(BclSourcesRoot)\System\Runtime\InteropServices\TrackerObjectManager.CoreCLR.cs" />
   </ItemGroup>
   <ItemGroup Condition="'$(FeatureCominterop)' == 'true'">
     <Compile Include="$(CommonPath)System\Runtime\InteropServices\IDispatch.cs">

src/coreclr/System.Private.CoreLib/src/System/ComAwareWeakReference.CoreCLR.cs

@@ -10,13 +10,21 @@ namespace System
     internal sealed partial class ComAwareWeakReference
     {
         [LibraryImport(RuntimeHelpers.QCall, EntryPoint = "ComWeakRefToObject")]
-        private static partial void ComWeakRefToObject(IntPtr pComWeakRef, long wrapperId, ObjectHandleOnStack retRcw);
+        private static partial void ComWeakRefToObject(IntPtr pComWeakRef, ObjectHandleOnStack retRcw);
 
-        internal static object? ComWeakRefToObject(IntPtr pComWeakRef, long wrapperId)
+        internal static object? ComWeakRefToObject(IntPtr pComWeakRef, object? context)
         {
-            object? retRcw = null;
-            ComWeakRefToObject(pComWeakRef, wrapperId, ObjectHandleOnStack.Create(ref retRcw));
-            return retRcw;
+#if FEATURE_COMINTEROP
+            if (context is null)
+            {
+                // This wrapper was not created by ComWrappers, so we try to rehydrate using built-in COM.
+                object? retRcw = null;
+                ComWeakRefToObject(pComWeakRef, ObjectHandleOnStack.Create(ref retRcw));
+                return retRcw;
+            }
+#endif // FEATURE_COMINTEROP
+
+            return ComWeakRefToComWrappersObject(pComWeakRef, context);
         }
 
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
@@ -40,16 +48,33 @@ internal static unsafe bool PossiblyComObject(object target)
         internal static extern bool HasInteropInfo(object target);
 
         [LibraryImport(RuntimeHelpers.QCall, EntryPoint = "ObjectToComWeakRef")]
-        private static partial IntPtr ObjectToComWeakRef(ObjectHandleOnStack retRcw, out long wrapperId);
+        private static partial IntPtr ObjectToComWeakRef(ObjectHandleOnStack retRcw);
 
-        internal static nint ObjectToComWeakRef(object target, out long wrapperId)
+        internal static nint ObjectToComWeakRef(object target, out object? context)
         {
-            if (HasInteropInfo(target))
+            if (!HasInteropInfo(target))
+            {
+                context = null;
+                return IntPtr.Zero;
+            }
+
+#if FEATURE_COMINTEROP
+            if (target is __ComObject)
+            {
+                // This object is using built-in COM, so use built-in COM to create the weak reference.
+                context = null;
+                return ObjectToComWeakRef(ObjectHandleOnStack.Create(ref target));
+            }
+#endif // FEATURE_COMINTEROP
+
+            if (PossiblyComWrappersObject(target))
             {
-                return ObjectToComWeakRef(ObjectHandleOnStack.Create(ref target), out wrapperId);
+                return ComWrappersObjectToComWeakRef(target, out context);
             }
 
-            wrapperId = 0;
+            // This object is not produced using built-in COM or ComWrappers
+            // or is an aggregated object, so we cannot create a weak reference.
+            context = null;
             return IntPtr.Zero;
         }
     }

src/coreclr/System.Private.CoreLib/src/System/GC.CoreCLR.cs

@@ -106,7 +106,7 @@ internal enum GC_ALLOC_FLAGS
         private static partial long GetTotalMemory();
 
         [LibraryImport(RuntimeHelpers.QCall, EntryPoint = "GCInterface_Collect")]
-        private static partial void _Collect(int generation, int mode);
+        private static partial void _Collect(int generation, int mode, [MarshalAs(UnmanagedType.U1)] bool lowMemoryPressure);
 
         [MethodImpl(MethodImplOptions.InternalCall)]
         private static extern int GetMaxGeneration();
@@ -174,7 +174,7 @@ public static void Collect(int generation)
         public static void Collect()
         {
             // -1 says to GC all generations.
-            _Collect(-1, (int)InternalGCCollectionMode.Blocking);
+            _Collect(-1, (int)InternalGCCollectionMode.Blocking, lowMemoryPressure: false);
         }
 
         public static void Collect(int generation, GCCollectionMode mode)
@@ -189,6 +189,11 @@ public static void Collect(int generation, GCCollectionMode mode, bool blocking)
         }
 
         public static void Collect(int generation, GCCollectionMode mode, bool blocking, bool compacting)
+        {
+            Collect(generation, mode, blocking, compacting, lowMemoryPressure: false);
+        }
+
+        internal static void Collect(int generation, GCCollectionMode mode, bool blocking, bool compacting, bool lowMemoryPressure)
         {
             ArgumentOutOfRangeException.ThrowIfNegative(generation);
 
@@ -197,7 +202,6 @@ public static void Collect(int generation, GCCollectionMode mode, bool blocking,
                 throw new ArgumentOutOfRangeException(nameof(mode), SR.ArgumentOutOfRange_Enum);
             }
 
-
             int iInternalModes = 0;
 
             if (mode == GCCollectionMode.Optimized)
@@ -222,7 +226,9 @@ public static void Collect(int generation, GCCollectionMode mode, bool blocking,
             }
 
             if (compacting)
+            {
                 iInternalModes |= (int)InternalGCCollectionMode.Compacting;
+            }
 
             if (blocking)
             {
@@ -233,7 +239,7 @@ public static void Collect(int generation, GCCollectionMode mode, bool blocking,
                 iInternalModes |= (int)InternalGCCollectionMode.NonBlocking;
             }
 
-            _Collect(generation, iInternalModes);
+            _Collect(generation, (int)iInternalModes, lowMemoryPressure);
         }
 
         public static int CollectionCount(int generation)

src/coreclr/System.Private.CoreLib/src/System/Runtime/InteropServices/ComWrappers.CoreCLR.cs

@@ -0,0 +1,148 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System.Threading;
+using System.Runtime.CompilerServices;
+using System.Collections.Generic;
+using System.Collections.Concurrent;
+
+namespace System.Runtime.InteropServices
+{
+    /// <summary>
+    /// Class for managing wrappers of COM IUnknown types.
+    /// </summary>
+    public abstract partial class ComWrappers
+    {
+        /// <summary>
+        /// Get the runtime provided IUnknown implementation.
+        /// </summary>
+        /// <param name="fpQueryInterface">Function pointer to QueryInterface.</param>
+        /// <param name="fpAddRef">Function pointer to AddRef.</param>
+        /// <param name="fpRelease">Function pointer to Release.</param>
+        public static unsafe void GetIUnknownImpl(out IntPtr fpQueryInterface, out IntPtr fpAddRef, out IntPtr fpRelease)
+            => GetIUnknownImplInternal(out fpQueryInterface, out fpAddRef, out fpRelease);
+
+        [LibraryImport(RuntimeHelpers.QCall, EntryPoint = "ComWrappers_GetIUnknownImpl")]
+        [SuppressGCTransition]
+        private static partial void GetIUnknownImplInternal(out IntPtr fpQueryInterface, out IntPtr fpAddRef, out IntPtr fpRelease);
+
+        internal static IntPtr DefaultIUnknownVftblPtr { get; } = CreateDefaultIUnknownVftbl();
+        internal static IntPtr TaggedImplVftblPtr { get; } = CreateTaggedImplVftbl();
+        internal static IntPtr DefaultIReferenceTrackerTargetVftblPtr { get; } = CreateDefaultIReferenceTrackerTargetVftbl();
+
+        private static unsafe IntPtr CreateDefaultIUnknownVftbl()
+        {
+            IntPtr* vftbl = (IntPtr*)RuntimeHelpers.AllocateTypeAssociatedMemory(typeof(ComWrappers), 3 * sizeof(IntPtr));
+            GetIUnknownImpl(out vftbl[0], out vftbl[1], out vftbl[2]);
+            return (IntPtr)vftbl;
+        }
+
+        private static unsafe IntPtr CreateTaggedImplVftbl()
+        {
+            IntPtr* vftbl = (IntPtr*)RuntimeHelpers.AllocateTypeAssociatedMemory(typeof(ComWrappers), 4 * sizeof(IntPtr));
+            GetIUnknownImpl(out vftbl[0], out vftbl[1], out vftbl[2]);
+            vftbl[3] = GetTaggedImplCurrentVersion();
+            return (IntPtr)vftbl;
+        }
+
+        private static readonly ConditionalWeakTable<object, List<ManagedObjectWrapperHolder>> s_allManagedObjectWrapperTable = [];
+
+        private static unsafe void RegisterManagedObjectWrapperForDiagnostics(object instance, ManagedObjectWrapperHolder wrapper)
+        {
+            // Record the relationship between the managed object and this wrapper.
+            // This will ensure that we keep all ManagedObjectWrapperHolders alive as long as the managed object is alive,
+            // even if the ComWrappers instance dies.
+            // We must do this as the runtime's recording of the relationship between the managed object and the wrapper
+            // lives as long as the object is alive, and we record the raw pointer to the wrapper.
+            List<ManagedObjectWrapperHolder> allWrappersForThisInstance = s_allManagedObjectWrapperTable.GetOrCreateValue(instance);
+            lock (allWrappersForThisInstance)
+            {
+                allWrappersForThisInstance.Add(wrapper);
+            }
+            RegisterManagedObjectWrapperForDiagnosticsInternal(ObjectHandleOnStack.Create(ref instance), wrapper.Wrapper);
+        }
+
+        [LibraryImport(RuntimeHelpers.QCall, EntryPoint = "ComWrappers_RegisterManagedObjectWrapperForDiagnostics")]
+        private static unsafe partial void RegisterManagedObjectWrapperForDiagnosticsInternal(ObjectHandleOnStack instance, ManagedObjectWrapper* wrapper);
+
+        private static void RegisterNativeObjectWrapperForDiagnostics(NativeObjectWrapper registeredWrapper)
+        {
+            object target = registeredWrapper.ProxyHandle.Target!;
+            RegisterNativeObjectWrapperForDiagnosticsInternal(ObjectHandleOnStack.Create(ref target), ObjectHandleOnStack.Create(ref registeredWrapper));
+        }
+
+        [LibraryImport(RuntimeHelpers.QCall, EntryPoint = "ComWrappers_RegisterNativeObjectWrapperForDiagnostics")]
+        private static unsafe partial void RegisterNativeObjectWrapperForDiagnosticsInternal(ObjectHandleOnStack target, ObjectHandleOnStack wrapper);
+
+        internal static int CallICustomQueryInterface(ManagedObjectWrapperHolder holder, ref Guid iid, out IntPtr ppObject)
+        {
+            if (holder.WrappedObject is ICustomQueryInterface customQueryInterface)
+            {
+                return (int)customQueryInterface.GetInterface(ref iid, out ppObject);
+            }
+
+            ppObject = IntPtr.Zero;
+            return -1; // See TryInvokeICustomQueryInterfaceResult
+        }
+
+        internal static IntPtr GetOrCreateComInterfaceForObjectWithGlobalMarshallingInstance(object obj)
+        {
+            try
+            {
+                return s_globalInstanceForMarshalling is null
+                    ? IntPtr.Zero
+                    : s_globalInstanceForMarshalling.GetOrCreateComInterfaceForObject(obj, CreateComInterfaceFlags.TrackerSupport);
+            }
+            catch (ArgumentException)
+            {
+                // We've failed to create a COM interface for the object.
+                // Fallback to built-in COM.
+                return IntPtr.Zero;
+            }
+        }
+
+        internal static object? GetOrCreateObjectForComInstanceWithGlobalMarshallingInstance(IntPtr comObject, CreateObjectFlags flags)
+        {
+            try
+            {
+                return s_globalInstanceForMarshalling?.GetOrCreateObjectForComInstance(comObject, flags);
+            }
+            catch (ArgumentNullException)
+            {
+                // We've failed to create a managed object for the COM instance.
+                // Fallback to built-in COM.
+                return null;
+            }
+        }
+
+        [LibraryImport(RuntimeHelpers.QCall, EntryPoint = "ComWrappers_GetIReferenceTrackerTargetVftbl")]
+        [SuppressGCTransition]
+        private static partial IntPtr GetDefaultIReferenceTrackerTargetVftbl();
+
+        private static IntPtr CreateDefaultIReferenceTrackerTargetVftbl()
+            => GetDefaultIReferenceTrackerTargetVftbl();
+
+        private static IntPtr GetTaggedImplCurrentVersion()
+        {
+            return GetTaggedImpl();
+        }
+
+        [LibraryImport(RuntimeHelpers.QCall, EntryPoint = "ComWrappers_GetTaggedImpl")]
+        [SuppressGCTransition]
+        private static partial IntPtr GetTaggedImpl();
+
+        internal sealed partial class ManagedObjectWrapperHolder
+        {
+            [LibraryImport(RuntimeHelpers.QCall, EntryPoint = "ComWrappers_RegisterIsRootedCallback")]
+            private static partial void RegisterIsRootedCallback();
+
+            private static IntPtr AllocateRefCountedHandle(ManagedObjectWrapperHolder holder)
+            {
+                return AllocateRefCountedHandle(ObjectHandleOnStack.Create(ref holder));
+            }
+
+            [LibraryImport(RuntimeHelpers.QCall, EntryPoint = "ComWrappers_AllocateRefCountedHandle")]
+            private static partial IntPtr AllocateRefCountedHandle(ObjectHandleOnStack obj);
+        }
+    }
+}