Fixing COM Aggregation and XAML reference tracking on .NET AOT

src/coreclr/nativeaot/Runtime.Base/src/System/Runtime/InternalCalls.cs

@@ -60,14 +60,14 @@ internal static class InternalCalls
 
         // Force a garbage collection.
         [RuntimeExport("RhCollect")]
-        internal static void RhCollect(int generation, InternalGCCollectionMode mode)
+        internal static void RhCollect(int generation, InternalGCCollectionMode mode, bool lowMemoryP = false)
         {
-            RhpCollect(generation, mode);
+            RhpCollect(generation, mode, lowMemoryP);
         }
 
         [DllImport(Redhawk.BaseName)]
         [UnmanagedCallConv(CallConvs = new Type[] { typeof(CallConvCdecl) })]
-        private static extern void RhpCollect(int generation, InternalGCCollectionMode mode);
+        private static extern void RhpCollect(int generation, InternalGCCollectionMode mode, bool lowMemoryP);
 
         [RuntimeExport("RhGetGcTotalMemory")]
         internal static long RhGetGcTotalMemory()

src/coreclr/nativeaot/Runtime/GCHelpers.cpp

@@ -25,7 +25,7 @@
 #include "threadstore.inl"
 #include "thread.inl"
 
-EXTERN_C NATIVEAOT_API void __cdecl RhpCollect(uint32_t uGeneration, uint32_t uMode)
+EXTERN_C NATIVEAOT_API void __cdecl RhpCollect(uint32_t uGeneration, uint32_t uMode, UInt32_BOOL lowMemoryP)
 {
     // This must be called via p/invoke rather than RuntimeImport to make the stack crawlable.
 
@@ -35,7 +35,7 @@ EXTERN_C NATIVEAOT_API void __cdecl RhpCollect(uint32_t uGeneration, uint32_t uM
     pCurThread->DisablePreemptiveMode();
 
     ASSERT(!pCurThread->IsDoNotTriggerGcSet());
-    GCHeapUtilities::GetGCHeap()->GarbageCollect(uGeneration, FALSE, uMode);
+    GCHeapUtilities::GetGCHeap()->GarbageCollect(uGeneration, lowMemoryP, uMode);
 
     pCurThread->EnablePreemptiveMode();
 }
@@ -128,6 +128,11 @@ COOP_PINVOKE_HELPER(int32_t, RhSetGcLatencyMode, (int32_t newLatencyMode))
     return GCHeapUtilities::GetGCHeap()->SetGcLatencyMode(newLatencyMode);
 }
 
+COOP_PINVOKE_HELPER(FC_BOOL_RET, RhIsPromoted, (OBJECTREF obj))
+{
+    FC_RETURN_BOOL(GCHeapUtilities::GetGCHeap()->IsPromoted(obj));
+}
+
 COOP_PINVOKE_HELPER(FC_BOOL_RET, RhIsServerGc, ())
 {
     FC_RETURN_BOOL(GCHeapUtilities::IsServerHeap());

src/coreclr/nativeaot/Runtime/gcrhenv.cpp

@@ -989,15 +989,19 @@ bool GCToEEInterface::EagerFinalized(Object* obj)
     // Managed code should not be running.
     ASSERT(GCHeapUtilities::GetGCHeap()->IsGCInProgressHelper());
 
-    // the lowermost 1 bit is reserved for storing additional info about the handle
-    const uintptr_t HandleTagBits = 1;
+    // the lowermost 2 bits are reserved for storing additional info about the handle
+    // we can use these bits because handle is at least 4 byte aligned
+    const uintptr_t HandleTagBits = 3;
 
     WeakReference* weakRefObj = (WeakReference*)obj;
     OBJECTHANDLE handle = (OBJECTHANDLE)(weakRefObj->m_taggedHandle & ~HandleTagBits);
-    _ASSERTE((weakRefObj->m_taggedHandle & 2) == 0);
-    HandleType handleType = (weakRefObj->m_taggedHandle & 1) ? HandleType::HNDTYPE_WEAK_LONG : HandleType::HNDTYPE_WEAK_SHORT;
+    HandleType handleType = (weakRefObj->m_taggedHandle & 2) ?
+        HandleType::HNDTYPE_STRONG :
+        (weakRefObj->m_taggedHandle & 1) ?
+        HandleType::HNDTYPE_WEAK_LONG :
+        HandleType::HNDTYPE_WEAK_SHORT;
     // keep the bit that indicates whether this reference was tracking resurrection, clear the rest.
-    weakRefObj->m_taggedHandle &= HandleTagBits;
+    weakRefObj->m_taggedHandle &= (uintptr_t)1;
     GCHandleUtilities::GetGCHandleManager()->DestroyHandleOfType(handle, handleType);
     return true;
 }

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

@@ -210,6 +210,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\TrackerObjectManager.NativeAot.cs" Condition="'$(FeatureComWrappers)' == 'true'" />    
     <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'" />
@@ -273,6 +274,9 @@
     <Compile Include="$(CommonPath)\Interop\Windows\Ole32\Interop.CoGetApartmentType.cs">
       <Link>Interop\Windows\Ole32\Interop.CoGetApartmentType.cs</Link>
     </Compile>
+    <Compile Include="$(CommonPath)\Interop\Windows\Ole32\Interop.CoGetContextToken.cs">
+      <Link>Interop\Windows\Ole32\Interop.CoGetContextToken.cs</Link>
+    </Compile>    
     <Compile Include="$(CommonPath)\Interop\Windows\OleAut32\Interop.VariantClear.cs">
       <Link>Interop\Windows\OleAut32\Interop.VariantClear.cs</Link>
     </Compile>

src/coreclr/nativeaot/System.Private.CoreLib/src/System/ComAwareWeakReference.NativeAot.cs

@@ -2,32 +2,94 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System;
+using System.Runtime.CompilerServices;
+using System.Runtime.InteropServices;
 
 #if FEATURE_COMINTEROP || FEATURE_COMWRAPPERS
 
-#pragma warning disable IDE0060
-
 namespace System
 {
     internal sealed partial class ComAwareWeakReference
     {
+        private static readonly Guid IID_IInspectable = new Guid(0xAF86E2E0, 0xB12D, 0x4c6a, 0x9C, 0x5A, 0xD7, 0xAA, 0x65, 0x10, 0x1E, 0x90);
+        private static readonly Guid IID_IWeakReferenceSource = new Guid(0x00000038, 0, 0, 0xC0, 0, 0, 0, 0, 0, 0, 0x46);
+
+        // We don't want to consult the ComWrappers if no RCW objects have been created.
+        // So we instead use this variable which is set by ComWrappers to determine
+        // whether RCW objects have been created before consulting it as part of PossiblyComObject.
+        internal static bool ComWrappersRcwInitialized;
+
         internal static object? ComWeakRefToObject(IntPtr pComWeakRef, long wrapperId)
         {
-            // NativeAOT support for COM WeakReference is NYI
-            throw new NotImplementedException();
+            if (wrapperId == 0)
+            {
+                return null;
+            }
+
+            // Using the IWeakReference*, get ahold of the target native COM object's IInspectable*.  If this resolve fails or
+            // returns null, then we assume that the underlying native COM object is no longer alive, and thus we cannot create a
+            // new RCW for it.
+            if (IWeakReference.Resolve(pComWeakRef, IID_IInspectable, out IntPtr targetPtr) == HResults.S_OK &&
+                targetPtr != IntPtr.Zero)
+            {
+                using ComHolder target = new ComHolder(targetPtr);
+                if (Marshal.QueryInterface(target.Ptr, in ComWrappers.IID_IUnknown, out IntPtr targetIdentityPtr) == HResults.S_OK)
+                {
+                    using ComHolder targetIdentity = new ComHolder(targetIdentityPtr);
+                    return ComWrappers.GetOrCreateObjectFromWrapper(wrapperId, targetIdentity.Ptr);
+                }
+            }
+
+            return null;
         }
 
-        internal static bool PossiblyComObject(object target)
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        internal static unsafe bool PossiblyComObject(object target)
         {
-            // NativeAOT support for COM WeakReference is NYI
-            return false;
+            return ComWrappersRcwInitialized && ComWrappers.IsRcwObject(target);
         }
 
         internal static IntPtr ObjectToComWeakRef(object target, out long wrapperId)
         {
-            // NativeAOT support for COM WeakReference is NYI
+            if (ComWrappers.TryGetComInstanceForIID(
+                target,
+                IID_IWeakReferenceSource,
+                out IntPtr weakReferenceSourcePtr,
+                out bool isAggregated,
+                out wrapperId))
+            {
+                // If the RCW is an aggregated RCW, then the managed object cannot be recreated from the IUnknown
+                // as the outer IUnknown wraps the managed object. In this case, don't create a weak reference backed
+                // by a COM weak reference.
+                using ComHolder weakReferenceSource = new ComHolder(weakReferenceSourcePtr);
+                if (!isAggregated && IWeakReferenceSource.GetWeakReference(weakReferenceSource.Ptr, out IntPtr weakReference) == HResults.S_OK)
+                {
+                    return weakReference;
+                }
+            }
+
             wrapperId = 0;
-            return 0;
+            return IntPtr.Zero;
+        }
+    }
+
+    // Wrapper for IWeakReference
+    internal static unsafe class IWeakReference
+    {
+        public static int Resolve(IntPtr pThis, Guid guid, out IntPtr inspectable)
+        {
+            fixed (IntPtr* inspectablePtr = &inspectable)
+                return (*(delegate* unmanaged<IntPtr, Guid*, IntPtr*, int>**)pThis)[3](pThis, &guid, inspectablePtr);
+        }
+    }
+
+    // Wrapper for IWeakReferenceSource
+    internal static unsafe class IWeakReferenceSource
+    {
+        public static int GetWeakReference(IntPtr pThis, out IntPtr weakReference)
+        {
+            fixed (IntPtr* weakReferencePtr = &weakReference)
+                return (*(delegate* unmanaged<IntPtr, IntPtr*, int>**)pThis)[3](pThis, weakReferencePtr);
         }
     }
 }