Skip closure allocation in ObjectReferenceWithContext<T> (#1587)

* Skip closure allocation in ObjectReferenceWithContext<T>

* Remove even more closure allocations in ObjectReferenceWithContext<T> (#1589)

* Skip closures for 'Context.CallInContext' calls

* Use function pointers instead of delegates

* Move context lambda to FOH

Author: Sergio0694

src/WinRT.Runtime/AgileReference.cs

@@ -94,7 +94,7 @@ protected virtual void Dispose(bool disposing)
                     {
                         Git.RevokeInterfaceFromGlobal(_cookie);
                     }
-                    catch(ArgumentException)
+                    catch (ArgumentException)
                     {
                         // Revoking cookie from GIT table may fail if apartment is gone.
                     }

src/WinRT.Runtime/Context.cs

@@ -29,17 +29,31 @@ public static unsafe IntPtr GetContextCallback()
         // On any exception, calls onFail callback if any set.
         // If not set, exception is handled due to today we don't
         // have any scenario to propagate it from here.
-        public unsafe static void CallInContext(IntPtr contextCallbackPtr, IntPtr contextToken, Action callback, Action onFailCallback)
+        //
+        // On modern .NET, we can use function pointers to avoid the
+        // small binary size increase from all generated fields and
+        // logic to cache delegates, since we don't need any of that.
+        public unsafe static void CallInContext(
+            IntPtr contextCallbackPtr,
+            IntPtr contextToken,
+#if NET && CsWinRT_LANG_11_FEATURES
+            delegate*<object, void> callback,
+            delegate*<object, void> onFailCallback,
+#else
+            Action<object> callback,
+            Action<object> onFailCallback,
+#endif
+            object state)
         {
             // Check if we are already on the same context, if so we do not need to switch.
             if(contextCallbackPtr == IntPtr.Zero || GetContextToken() == contextToken)
             {
-                callback();
+                callback(state);
                 return;
             }
 
 #if NET && CsWinRT_LANG_11_FEATURES
-            IContextCallbackVftbl.ContextCallback(contextCallbackPtr, callback, onFailCallback);
+            IContextCallbackVftbl.ContextCallback(contextCallbackPtr, callback, onFailCallback, state);
 #else
             ComCallData data = default;
             var contextCallback = new ABI.WinRT.Interop.IContextCallback(ObjectReference<ABI.WinRT.Interop.IContextCallback.Vftbl>.FromAbi(contextCallbackPtr));
@@ -48,13 +62,13 @@ public unsafe static void CallInContext(IntPtr contextCallbackPtr, IntPtr contex
             {
                 contextCallback.ContextCallback(_ =>
                 {
-                    callback();
+                    callback(state);
                     return 0;
                 }, &data, IID.IID_ICallbackWithNoReentrancyToApplicationSTA, 5);
             } 
-            catch(Exception)
+            catch (Exception)
             {
-                onFailCallback?.Invoke();
+                onFailCallback?.Invoke(state);
             }
 #endif
         }

src/WinRT.Runtime/Interop/IContextCallback.cs

@@ -17,6 +17,14 @@ internal struct ComCallData
         public IntPtr pUserDefined;
     }
 
+#if NET && CsWinRT_LANG_11_FEATURES
+    internal unsafe struct CallbackData
+    {
+        public delegate*<object, void> Callback;
+        public object State;
+    }
+#endif
+
 #if NET && CsWinRT_LANG_11_FEATURES
     internal unsafe struct IContextCallbackVftbl
     {
@@ -25,34 +33,31 @@ internal unsafe struct IContextCallbackVftbl
         private delegate* unmanaged[Stdcall]<IntPtr, IntPtr, ComCallData*, Guid*, int, IntPtr, int> ContextCallback_4;
 #pragma warning restore CS0649
 
-        public static void ContextCallback(IntPtr contextCallbackPtr, Action callback, Action onFailCallback)
+        public static void ContextCallback(IntPtr contextCallbackPtr, delegate*<object, void> callback, delegate*<object, void> onFailCallback, object state)
         {
             ComCallData comCallData;
             comCallData.dwDispid = 0;
             comCallData.dwReserved = 0;
 
-            // Copy the callback into a local to make sure it really is a local that
-            // gets marked as address taken, rather than something that could potentially
-            // be inlined into the caller into some state machine or anything else not safe.
-            Action callbackAddressTaken = callback;
+            CallbackData callbackData;
+            callbackData.Callback = callback;
+            callbackData.State = state;
 
             // We can just store a pointer to the callback to invoke in the context,
             // so we don't need to allocate another closure or anything. The callback
             // will be kept alive automatically, because 'comCallData' is address exposed.
             // We only do this if we can use C# 11, and if we're on modern .NET, to be safe.
             // In the callback below, we can then just retrieve the Action again to invoke it.
-            comCallData.pUserDefined = (IntPtr)(void*)&callbackAddressTaken;
+            comCallData.pUserDefined = (IntPtr)(void*)&callbackData;
 
             [UnmanagedCallersOnly]
             static int InvokeCallback(ComCallData* comCallData)
             {
                 try
                 {
-                    // Dereference the pointer to Action and invoke it (see notes above).
-                    // Once again, the pointer is not to the Action object, but just to the
-                    // local *reference* to the object, which is pinned (as it's a local).
-                    // That means that there's no pinning to worry about either.
-                    ((Action*)comCallData->pUserDefined)->Invoke();
+                    CallbackData* callbackData = (CallbackData*)comCallData->pUserDefined;
+
+                    callbackData->Callback(callbackData->State);
 
                     return 0; // S_OK
                 }
@@ -74,7 +79,10 @@ static int InvokeCallback(ComCallData* comCallData)
 
             if (hresult < 0)
             {
-                onFailCallback?.Invoke();
+                if (onFailCallback is not null)
+                {
+                    onFailCallback(state);
+                }
             }
         }
     }

src/WinRT.Runtime/ObjectReference.cs

@@ -698,17 +698,29 @@ private ConcurrentDictionary<IntPtr, IObjectReference> Make_CachedContext()
         private volatile bool _isAgileReferenceSet;
         private volatile AgileReference __agileReference;
         private AgileReference AgileReference => _isAgileReferenceSet ? __agileReference : Make_AgileReference();
-        private AgileReference Make_AgileReference()
+        private unsafe AgileReference Make_AgileReference()
         {
-            Context.CallInContext(_contextCallbackPtr, _contextToken, InitAgileReference, null);
+            Context.CallInContext(
+                _contextCallbackPtr,
+                _contextToken,
+#if NET && CsWinRT_LANG_11_FEATURES
+                &InitAgileReference,
+#else
+                InitAgileReference,
+#endif
+                null,
+                this);
 
             // Set after CallInContext callback given callback can fail to occur.
             _isAgileReferenceSet = true;
+
             return __agileReference;
 
-            void InitAgileReference()
+            static void InitAgileReference(object state)
             {
-                global::System.Threading.Interlocked.CompareExchange(ref __agileReference, new AgileReference(this), null);
+                ObjectReferenceWithContext<T> @this = Unsafe.As<ObjectReferenceWithContext<T>>(state);
+
+                global::System.Threading.Interlocked.CompareExchange(ref @this.__agileReference, new AgileReference(@this), null);
             }
         }
 
@@ -805,16 +817,32 @@ private ObjectReference<T> GetCurrentContext()
             // of ConcurrentDictionary<,> and transitively dependent types for every vtable type T, since it's not
             // something we actually need. Because the cache is private and we're the only ones using it, we can
             // just store the per-context agile references as IObjectReference values, and then cast them on return.
-            IObjectReference objectReference = CachedContext.GetOrAdd(currentContext, CreateForCurrentContext);
+#if NET
+            IObjectReference objectReference = CachedContext.GetOrAdd(currentContext, ContextCallbackHolder.Value, this);
+#else
+            IObjectReference objectReference = CachedContext.GetOrAdd(currentContext, ptr => ContextCallbackHolder.Value(ptr, this));
+#endif
 
             return Unsafe.As<ObjectReference<T>>(objectReference);
+        }
+
+        private static class ContextCallbackHolder
+        {
+            // We have a single lambda expression in this type, so we can manually rewrite it to a 'static readonly'
+            // field. This avoids the extra logic to lazily initialized it (it's already lazily initialized because
+            // it's in a 'beforefieldinit' type which is only used when the lambda is actually needed), and also it
+            // allows storing the entire delegate in the Frozen Object Heap (FOH) on modern runtimes.
+            public static readonly Func<IntPtr, IObjectReference, IObjectReference> Value = CreateForCurrentContext;
 
 #if NET
             [UnconditionalSuppressMessage("Trimming", "IL2087", Justification = "The '_iid' field is only empty when using annotated APIs not trim-safe.")]
 #endif
-            IObjectReference CreateForCurrentContext(IntPtr _)
+            private static IObjectReference CreateForCurrentContext(IntPtr _, IObjectReference state)
             {
-                var agileReference = AgileReference;
+                ObjectReferenceWithContext<T> @this = Unsafe.As<ObjectReferenceWithContext<T>>(state);
+
+                var agileReference = @this.AgileReference;
+
                 // We may fail to switch context and thereby not get an agile reference.
                 // In these cases, fallback to using the current context.
                 if (agileReference == null)
@@ -829,17 +857,17 @@ IObjectReference CreateForCurrentContext(IntPtr _)
                     // going through a trim-unsafe constructor, which is explicitly not supported in this configuration.
                     if (!RuntimeFeature.IsDynamicCodeCompiled)
                     {
-                        return agileReference.Get<T>(_iid);
+                        return agileReference.Get<T>(@this._iid);
                     }
 #endif
 
-                    if (_iid == Guid.Empty)
+                    if (@this._iid == Guid.Empty)
                     {
                         return agileReference.Get<T>(GuidGenerator.GetIID(typeof(T)));
                     }
                     else
                     {
-                        return agileReference.Get<T>(_iid);
+                        return agileReference.Get<T>(@this._iid);
                     }
                 }
                 catch (Exception)
@@ -858,8 +886,42 @@ protected override unsafe void Release()
                 CachedContext.Clear();
             }
 
-            Context.CallInContext(_contextCallbackPtr, _contextToken, base.Release, ReleaseWithoutContext);
+            Context.CallInContext(
+                _contextCallbackPtr,
+                _contextToken,
+#if NET && CsWinRT_LANG_11_FEATURES
+                &Release,
+                &ReleaseWithoutContext,
+#else
+                Release,
+                ReleaseWithoutContext,
+#endif
+                this);
+
             Context.DisposeContextCallback(_contextCallbackPtr);
+
+            static void Release(object state)
+            {
+                ObjectReferenceWithContext<T> @this = Unsafe.As<ObjectReferenceWithContext<T>>(state);
+
+                @this.ReleaseFromBase();
+            }
+
+            static void ReleaseWithoutContext(object state)
+            {
+                ObjectReferenceWithContext<T> @this = Unsafe.As<ObjectReferenceWithContext<T>>(state);
+
+                @this.ReleaseWithoutContext();
+            }
+        }
+
+        // Helper stub to invoke 'base.Release()' on a given 'ObjectReferenceWithContext<T>' input parameter.
+        // We can't just do 'param.base.Release()' (or something like that), so the only way to specifically
+        // invoke the base implementation of an overridden method on that object is to go through a helper
+        // instance method invoked on it that just calls the base implementation of the method we want.
+        private void ReleaseFromBase()
+        {
+            base.Release();
         }
 
         public override ObjectReference<IUnknownVftbl> AsKnownPtr(IntPtr ptr)