JIT: don't make local copy when devirtualizing boxed interface call (#128270)

When devirtualizing an interface call on a GT_BOX(local), the importer
would replace the heap box with a stack-local copy and invoke the
unboxed entry on that copy. This is unsafe: if the callee returns an
interior managed pointer (for instance via [UnscopedRef]), the returned
byref aliases the temporary copy and dangles once the enclosing frame
goes away.

Always keep the heap box and rewrite the call to invoke the unboxed
entry on box + sizeof(MT). Object stack allocation already promotes
non-escaping heap boxes to the stack, so we recover the previous
optimization in the safe cases.

The BR_MAKE_LOCAL_COPY removal option and its implementation in
gtTryRemoveBoxUpstreamEffects are no longer reachable and are removed.

Fixes #122099.

---------

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: AndyAyersMS

src/coreclr/jit/compiler.h

@@ -3974,7 +3974,6 @@ class Compiler
         BR_REMOVE_BUT_NOT_NARROW,              // remove effects, return original source tree
         BR_DONT_REMOVE,                        // check if removal is possible, return copy source tree
         BR_DONT_REMOVE_WANT_TYPE_HANDLE,       // check if removal is possible, return type handle tree
-        BR_MAKE_LOCAL_COPY                     // revise box to copy to temp local and return local's address
     };
 
     GenTree* gtTryRemoveBoxUpstreamEffects(GenTree* tree, BoxRemovalOptions options = BR_REMOVE_AND_NARROW);

src/coreclr/jit/gentree.cpp

@@ -16536,10 +16536,9 @@ GenTree* Compiler::gtTryRemoveBoxUpstreamEffects(GenTree* op, BoxRemovalOptions
     Statement*  allocStmt = box->gtDefStmtWhenInlinedBoxValue;
     Statement*  copyStmt  = box->gtCopyStmtWhenInlinedBoxValue;
 
-    JITDUMP("gtTryRemoveBoxUpstreamEffects: %s to %s of BOX (valuetype)"
+    JITDUMP("gtTryRemoveBoxUpstreamEffects: %s of BOX (valuetype)"
             " [%06u] (assign/newobj " FMT_STMT " copy " FMT_STMT "\n",
-            (options == BR_DONT_REMOVE) ? "checking if it is possible" : "attempting",
-            (options == BR_MAKE_LOCAL_COPY) ? "make local unboxed version" : "remove side effects", dspTreeID(op),
+            (options == BR_DONT_REMOVE) ? "checking if it is possible" : "attempting", dspTreeID(op),
             allocStmt->GetID(), copyStmt->GetID());
 
     // If we don't recognize the form of the store, bail.
@@ -16615,65 +16614,6 @@ GenTree* Compiler::gtTryRemoveBoxUpstreamEffects(GenTree* op, BoxRemovalOptions
         return nullptr;
     }
 
-    // Handle case where we are optimizing the box into a local copy
-    if (options == BR_MAKE_LOCAL_COPY)
-    {
-        // Drill into the box to get at the box temp local and the box type
-        GenTree* boxTemp = box->BoxOp();
-        assert(boxTemp->IsLocal());
-        const unsigned boxTempLcl = boxTemp->AsLclVar()->GetLclNum();
-        assert(lvaTable[boxTempLcl].lvType == TYP_REF);
-        CORINFO_CLASS_HANDLE boxClass = lvaTable[boxTempLcl].lvClassHnd;
-        assert(boxClass != nullptr);
-
-        // Verify that the copy has the expected shape
-        // (store_blk|store_ind (add (boxTempLcl, ptr-size)))
-        //
-        // The shape here is constrained to the patterns we produce
-        // over in impImportAndPushBox for the inlined box case.
-        //
-        GenTree* copyDstAddr = copy->AsIndir()->Addr();
-        if (!copyDstAddr->OperIs(GT_ADD))
-        {
-            JITDUMP("Unexpected copy dest address tree\n");
-            return nullptr;
-        }
-
-        GenTree* copyDstAddrOp1 = copyDstAddr->AsOp()->gtOp1;
-        if (!copyDstAddrOp1->OperIs(GT_LCL_VAR) || (copyDstAddrOp1->AsLclVarCommon()->GetLclNum() != boxTempLcl))
-        {
-            JITDUMP("Unexpected copy dest address 1st addend\n");
-            return nullptr;
-        }
-
-        GenTree* copyDstAddrOp2 = copyDstAddr->AsOp()->gtOp2;
-        if (!copyDstAddrOp2->IsIntegralConst(TARGET_POINTER_SIZE))
-        {
-            JITDUMP("Unexpected copy dest address 2nd addend\n");
-            return nullptr;
-        }
-
-        // Screening checks have all passed. Do the transformation.
-        //
-        // Retype the box temp to be a struct
-        JITDUMP("Retyping box temp V%02u to struct %s\n", boxTempLcl, eeGetClassName(boxClass));
-        lvaTable[boxTempLcl].lvType   = TYP_UNDEF;
-        const bool isUnsafeValueClass = false;
-        lvaSetStruct(boxTempLcl, boxClass, isUnsafeValueClass);
-
-        // Remove the newobj and store to box temp
-        JITDUMP("Bashing NEWOBJ [%06u] to NOP\n", dspTreeID(boxLclDef));
-        boxLclDef->gtBashToNOP();
-
-        // Update the copy from the value to be boxed to the box temp
-        copy->AsIndir()->Addr() = gtNewLclVarAddrNode(boxTempLcl, TYP_BYREF);
-
-        // Return the address of the now-struct typed box temp
-        GenTree* retValue = gtNewLclVarAddrNode(boxTempLcl, TYP_BYREF);
-
-        return retValue;
-    }
-
     // If the copy is a struct copy, make sure we know how to isolate any source side effects.
     GenTree* copySrc = copy->Data();
 

src/coreclr/jit/importercalls.cpp

@@ -9284,189 +9284,67 @@ void Compiler::impDevirtualizeCall(GenTreeCall*            call,
 
             if (unboxedEntryMethod != nullptr)
             {
-                bool optimizedTheBox = false;
-
-                // If the 'this' object is a local box, see if we can revise things
-                // to not require boxing.
+                // Rewrite the call to target the unboxed entry on the box payload. Keep the heap box,
+                // since the callee may return an interior managed pointer into it; object stack allocation
+                // can later promote the box to the stack when escape analysis proves the receiver does not
+                // escape.
                 //
-                if (thisObj->IsBoxedValue() && !isExplicitTailCall)
+                if (requiresInstMethodTableArg)
                 {
-                    // Since the call is the only consumer of the box, we know the box can't escape
-                    // since it is being passed an interior pointer.
-                    //
-                    // So, revise the box to simply create a local copy, use the address of that copy
-                    // as the this pointer, and update the entry point to the unboxed entry.
+                    // Get the method table from the boxed object.
                     //
-                    // Ideally, we then inline the boxed method and and if it turns out not to modify
-                    // the copy, we can undo the copy too.
-                    GenTree* localCopyThis = nullptr;
+                    // TODO-CallArgs-REVIEW: Use thisObj here? Differs by gtEffectiveVal.
+                    GenTree* const clonedThisArg = gtClone(thisArg->GetEarlyNode());
 
-                    if (requiresInstMethodTableArg)
+                    if (clonedThisArg == nullptr)
                     {
-                        // Perform a trial box removal and ask for the type handle tree that fed the box.
-                        //
-                        JITDUMP("Unboxed entry needs method table arg...\n");
-                        GenTree* methodTableArg =
-                            gtTryRemoveBoxUpstreamEffects(thisObj, BR_DONT_REMOVE_WANT_TYPE_HANDLE);
-
-                        if (methodTableArg != nullptr)
-                        {
-                            // If that worked, turn the box into a copy to a local var
-                            //
-                            JITDUMP("Found suitable method table arg tree [%06u]\n", dspTreeID(methodTableArg));
-                            localCopyThis = gtTryRemoveBoxUpstreamEffects(thisObj, BR_MAKE_LOCAL_COPY);
-
-                            if (localCopyThis != nullptr)
-                            {
-                                // Pass the local var as this and the type handle as a new arg
-                                //
-                                JITDUMP("Success! invoking unboxed entry point on local copy, and passing method table "
-                                        "arg\n");
-                                // TODO-CallArgs-REVIEW: Might discard commas otherwise?
-                                assert(thisObj == thisArg->GetEarlyNode());
-                                thisArg->SetEarlyNode(localCopyThis);
-                                INDEBUG(call->gtCallDebugFlags |= GTF_CALL_MD_UNBOXED);
-
-                                call->gtArgs.InsertInstParam(this, methodTableArg);
-
-                                call->gtCallMethHnd   = unboxedEntryMethod;
-                                derivedMethod         = unboxedEntryMethod;
-                                pDerivedResolvedToken = &dvInfo.resolvedTokenDevirtualizedUnboxedMethod;
-
-                                // Method attributes will differ because unboxed entry point is shared
-                                //
-                                const DWORD unboxedMethodAttribs =
-                                    info.compCompHnd->getMethodAttribs(unboxedEntryMethod);
-                                JITDUMP("Updating method attribs from 0x%08x to 0x%08x\n", derivedMethodAttribs,
-                                        unboxedMethodAttribs);
-                                derivedMethodAttribs = unboxedMethodAttribs;
-                                optimizedTheBox      = true;
-                            }
-                            else
-                            {
-                                JITDUMP("Sorry, failed to undo the box -- can't convert to local copy\n");
-                            }
-                        }
-                        else
-                        {
-                            JITDUMP("Sorry, failed to undo the box -- can't find method table arg\n");
-                        }
+                        JITDUMP("unboxed entry needs MT arg, but `this` was too complex to clone. Deferring update.\n");
                     }
                     else
                     {
-                        JITDUMP("Found unboxed entry point, trying to simplify box to a local copy\n");
-                        localCopyThis = gtTryRemoveBoxUpstreamEffects(thisObj, BR_MAKE_LOCAL_COPY);
-
-                        if (localCopyThis != nullptr)
-                        {
-                            JITDUMP("Success! invoking unboxed entry point on local copy\n");
-                            assert(thisObj == thisArg->GetEarlyNode());
-                            // TODO-CallArgs-REVIEW: Might discard commas otherwise?
-                            thisArg->SetEarlyNode(localCopyThis);
-                            call->gtCallMethHnd = unboxedEntryMethod;
-                            INDEBUG(call->gtCallDebugFlags |= GTF_CALL_MD_UNBOXED);
-                            derivedMethod         = unboxedEntryMethod;
-                            pDerivedResolvedToken = &dvInfo.resolvedTokenDevirtualizedUnboxedMethod;
-
-                            optimizedTheBox = true;
-                        }
-                        else
-                        {
-                            JITDUMP("Sorry, failed to undo the box\n");
-                        }
-                    }
-
-                    if (optimizedTheBox)
-                    {
-                        assert(localCopyThis->IsLclVarAddr());
-
-                        // We may end up inlining this call, so the local copy must be marked as "aliased",
-                        // making sure the inlinee importer will know when to spill references to its value.
-                        lvaGetDesc(localCopyThis->AsLclFld())->lvHasLdAddrOp = true;
-                        Metrics.DevirtualizedCallRemovedBox++;
-                        Metrics.DevirtualizedCallUnboxedEntry++;
-
-#if FEATURE_TAILCALL_OPT
-                        if (call->IsImplicitTailCall())
-                        {
-                            JITDUMP("Clearing the implicit tail call flag\n");
+                        JITDUMP("revising call to invoke unboxed entry with additional method table arg\n");
 
-                            // If set, we clear the implicit tail call flag
-                            // as we just introduced a new address taken local variable
-                            //
-                            call->gtCallMoreFlags &= ~GTF_CALL_M_IMPLICIT_TAILCALL;
-                        }
-#endif // FEATURE_TAILCALL_OPT
-                    }
-                }
+                        GenTree* const methodTableArg = gtNewMethodTableLookup(clonedThisArg);
 
-                if (!optimizedTheBox)
-                {
-                    // If we get here, we have a boxed value class that either wasn't boxed
-                    // locally, or was boxed locally but we were unable to remove the box for
-                    // various reasons.
-                    //
-                    // We can still update the call to invoke the unboxed entry, if the
-                    // boxed value is simple.
-                    //
-                    if (requiresInstMethodTableArg)
-                    {
-                        // Get the method table from the boxed object.
+                        // Update the 'this' pointer to refer to the box payload
                         //
-                        // TODO-CallArgs-REVIEW: Use thisObj here? Differs by gtEffectiveVal.
-                        GenTree* const clonedThisArg = gtClone(thisArg->GetEarlyNode());
-
-                        if (clonedThisArg == nullptr)
-                        {
-                            JITDUMP(
-                                "unboxed entry needs MT arg, but `this` was too complex to clone. Deferring update.\n");
-                        }
-                        else
-                        {
-                            JITDUMP("revising call to invoke unboxed entry with additional method table arg\n");
-
-                            GenTree* const methodTableArg = gtNewMethodTableLookup(clonedThisArg);
-
-                            // Update the 'this' pointer to refer to the box payload
-                            //
-                            GenTree* const payloadOffset = gtNewIconNode(TARGET_POINTER_SIZE, TYP_I_IMPL);
-                            GenTree* const boxPayload =
-                                gtNewOperNode(GT_ADD, TYP_BYREF, thisArg->GetEarlyNode(), payloadOffset);
-
-                            assert(thisObj == thisArg->GetEarlyNode());
-                            thisArg->SetEarlyNode(boxPayload);
-                            call->gtCallMethHnd = unboxedEntryMethod;
-                            INDEBUG(call->gtCallDebugFlags |= GTF_CALL_MD_UNBOXED);
-
-                            // Method attributes will differ because unboxed entry point is shared
-                            //
-                            const DWORD unboxedMethodAttribs = info.compCompHnd->getMethodAttribs(unboxedEntryMethod);
-                            JITDUMP("Updating method attribs from 0x%08x to 0x%08x\n", derivedMethodAttribs,
-                                    unboxedMethodAttribs);
-                            derivedMethod         = unboxedEntryMethod;
-                            pDerivedResolvedToken = &dvInfo.resolvedTokenDevirtualizedUnboxedMethod;
-                            derivedMethodAttribs  = unboxedMethodAttribs;
-
-                            call->gtArgs.InsertInstParam(this, methodTableArg);
-                            Metrics.DevirtualizedCallUnboxedEntry++;
-                        }
-                    }
-                    else
-                    {
-                        JITDUMP("revising call to invoke unboxed entry\n");
-
                         GenTree* const payloadOffset = gtNewIconNode(TARGET_POINTER_SIZE, TYP_I_IMPL);
                         GenTree* const boxPayload =
                             gtNewOperNode(GT_ADD, TYP_BYREF, thisArg->GetEarlyNode(), payloadOffset);
 
+                        assert(thisObj == thisArg->GetEarlyNode());
                         thisArg->SetEarlyNode(boxPayload);
                         call->gtCallMethHnd = unboxedEntryMethod;
                         INDEBUG(call->gtCallDebugFlags |= GTF_CALL_MD_UNBOXED);
+
+                        // Method attributes will differ because unboxed entry point is shared
+                        //
+                        const DWORD unboxedMethodAttribs = info.compCompHnd->getMethodAttribs(unboxedEntryMethod);
+                        JITDUMP("Updating method attribs from 0x%08x to 0x%08x\n", derivedMethodAttribs,
+                                unboxedMethodAttribs);
                         derivedMethod         = unboxedEntryMethod;
                         pDerivedResolvedToken = &dvInfo.resolvedTokenDevirtualizedUnboxedMethod;
+                        derivedMethodAttribs  = unboxedMethodAttribs;
+
+                        call->gtArgs.InsertInstParam(this, methodTableArg);
                         Metrics.DevirtualizedCallUnboxedEntry++;
                     }
                 }
+                else
+                {
+                    JITDUMP("revising call to invoke unboxed entry\n");
+
+                    GenTree* const payloadOffset = gtNewIconNode(TARGET_POINTER_SIZE, TYP_I_IMPL);
+                    GenTree* const boxPayload =
+                        gtNewOperNode(GT_ADD, TYP_BYREF, thisArg->GetEarlyNode(), payloadOffset);
+
+                    thisArg->SetEarlyNode(boxPayload);
+                    call->gtCallMethHnd = unboxedEntryMethod;
+                    INDEBUG(call->gtCallDebugFlags |= GTF_CALL_MD_UNBOXED);
+                    derivedMethod         = unboxedEntryMethod;
+                    pDerivedResolvedToken = &dvInfo.resolvedTokenDevirtualizedUnboxedMethod;
+                    Metrics.DevirtualizedCallUnboxedEntry++;
+                }
             }
             else
             {

src/coreclr/jit/jitmetadatalist.h

@@ -57,7 +57,6 @@ JITMETADATAMETRIC(ImporterBranchFold,                    int,              0)
 JITMETADATAMETRIC(ImporterSwitchFold,                    int,              0)
 JITMETADATAMETRIC(DevirtualizedCall,                     int,              0)
 JITMETADATAMETRIC(DevirtualizedCallUnboxedEntry,         int,              0)
-JITMETADATAMETRIC(DevirtualizedCallRemovedBox,           int,              0)
 JITMETADATAMETRIC(GDV,                                   int,              0)
 JITMETADATAMETRIC(ClassGDV,                              int,              0)
 JITMETADATAMETRIC(MethodGDV,                             int,              0)

src/tests/JIT/Regression/JitBlue/Runtime_122099/Runtime_122099.cs

@@ -0,0 +1,73 @@
+// 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.CodeAnalysis;
+using System.Runtime.CompilerServices;
+using Xunit;
+using TestLibrary;
+
+namespace Runtime_122099;
+
+public class Runtime_122099
+{
+    // Verifies that an interface call on a boxed value type whose target returns an
+    // `[UnscopedRef]` ref to an instance field keeps the boxed receiver alive long
+    // enough for the returned ref to remain valid after the call returns.
+    [ActiveIssue("https://github.com/dotnet/runtime/issues/122099", typeof(PlatformDetection), nameof(PlatformDetection.IsMonoRuntime))]
+    [Fact]
+    public static int TestEntryPoint()
+    {
+        ref int r = ref EscapeRef();
+        r = 42;
+
+        // Reuse the same stack region the callees occupied.
+        Clobber();
+
+        return r == 42 ? 100 : 101;
+    }
+
+    [MethodImpl(MethodImplOptions.NoInlining | MethodImplOptions.AggressiveOptimization)]
+    private static ref int EscapeRef()
+    {
+        Struct1 s = new Struct1 { Field = 7 };
+        return ref BoxedRefEscape(s);
+    }
+
+    [MethodImpl(MethodImplOptions.NoInlining | MethodImplOptions.AggressiveOptimization)]
+    private static ref int BoxedRefEscape(Struct1 x)
+    {
+        return ref ((I1)x).Method1();
+    }
+
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    private static void Clobber()
+    {
+        Span<int> spam = stackalloc int[64];
+        for (int i = 0; i < spam.Length; i++)
+        {
+            spam[i] = unchecked((int)0xDEADBEEF);
+        }
+        // Prevent the stackalloc from being elided.
+        if (spam[0] != unchecked((int)0xDEADBEEF))
+        {
+            throw new Exception();
+        }
+    }
+}
+
+public interface I1
+{
+    [UnscopedRef]
+    ref int Method1();
+}
+
+public struct Struct1 : I1
+{
+    public int Field;
+
+    [UnscopedRef]
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    public ref int Method1() => ref Field;
+}
+