Undo struct promotion on "RetInd" code path (#58582)

* TOREVERT: small repro case

* Fix the undone struct promotion for return

* Display lclVar number

* Revert "TOREVERT: small repro case"

This reverts commit 9ef80ba.

* Add test case

* jit format

* Combine the condition

* Add an assert to catch the missing places

* Remove unneeded variable setting from csproj

Author: kunalspathak

src/coreclr/jit/clrjit.natvis

@@ -65,8 +65,8 @@ Documentation for VS debugger format specifiers: https://docs.microsoft.com/en-u
   </Type>
 
   <Type Name="LclVarDsc">
-    <DisplayString Condition="lvReason==0">[{lvType,en}]</DisplayString>
-    <DisplayString>[{lvType,en}-{lvReason,s}]</DisplayString>
+    <DisplayString Condition="lvReason==0">[V{lvSlotNum,d}: {lvType,en}]</DisplayString>
+    <DisplayString>[V{lvSlotNum,d}: {lvType,en}-{lvReason,s}]</DisplayString>
   </Type>
 
   <Type Name="GenTreeLclVar" Inheritable="false">

src/coreclr/jit/compiler.h

@@ -502,6 +502,9 @@ class LclVarDsc
     unsigned char lvUnusedStruct : 1; // All references to this promoted struct are through its field locals.
                                       // I.e. there is no longer any reference to the struct directly.
                                       // In this case we can simply remove this struct local.
+
+    unsigned char lvUndoneStructPromotion : 1; // The struct promotion was undone and hence there should be no
+                                               // reference to the fields of this struct.
 #endif
 
     unsigned char lvLRACandidate : 1; // Tracked for linear scan register allocation purposes

src/coreclr/jit/lclvars.cpp

@@ -4025,6 +4025,16 @@ void Compiler::lvaMarkLclRefs(GenTree* tree, BasicBlock* block, Statement* stmt,
 
     varDsc->incRefCnts(weight, this);
 
+#ifdef DEBUG
+    if (varDsc->lvIsStructField)
+    {
+        // If ref count was increased for struct field, ensure that the
+        // parent struct is still promoted.
+        LclVarDsc* parentStruct = &lvaTable[varDsc->lvParentLcl];
+        assert(!parentStruct->lvUndoneStructPromotion);
+    }
+#endif
+
     if (!isRecompute)
     {
         if (lvaVarAddrExposed(lclNum))

src/coreclr/jit/morph.cpp

@@ -13880,6 +13880,12 @@ GenTree* Compiler::fgMorphRetInd(GenTreeUnOp* ret)
 
     if (addr->OperIs(GT_ADDR) && addr->gtGetOp1()->OperIs(GT_LCL_VAR))
     {
+        // If struct promotion was undone, adjust the annotations
+        if (fgGlobalMorph && fgMorphImplicitByRefArgs(addr))
+        {
+            return ind;
+        }
+
         // If `return` retypes LCL_VAR as a smaller struct it should not set `doNotEnregister` on that
         // LclVar.
         // Example: in `Vector128:AsVector2` we have RETURN SIMD8(OBJ SIMD8(ADDR byref(LCL_VAR SIMD16))).
@@ -17531,6 +17537,8 @@ void Compiler::fgRetypeImplicitByRefArgs()
 
 void Compiler::fgMarkDemotedImplicitByRefArgs()
 {
+    JITDUMP("\n*************** In fgMarkDemotedImplicitByRefArgs()\n");
+
 #if (defined(TARGET_AMD64) && !defined(UNIX_AMD64_ABI)) || defined(TARGET_ARM64)
 
     for (unsigned lclNum = 0; lclNum < info.compArgsCount; lclNum++)
@@ -17539,6 +17547,8 @@ void Compiler::fgMarkDemotedImplicitByRefArgs()
 
         if (lvaIsImplicitByRefLocal(lclNum))
         {
+            JITDUMP("Clearing annotation for V%02d\n", lclNum);
+
             if (varDsc->lvPromoted)
             {
                 // The parameter is simply a pointer now, so clear lvPromoted.  It was left set
@@ -17565,7 +17575,8 @@ void Compiler::fgMarkDemotedImplicitByRefArgs()
                 LclVarDsc* structVarDsc     = &lvaTable[structLclNum];
                 structVarDsc->lvAddrExposed = false;
 #ifdef DEBUG
-                structVarDsc->lvUnusedStruct = true;
+                structVarDsc->lvUnusedStruct          = true;
+                structVarDsc->lvUndoneStructPromotion = true;
 #endif // DEBUG
 
                 unsigned fieldLclStart = structVarDsc->lvFieldLclStart;
@@ -17574,6 +17585,8 @@ void Compiler::fgMarkDemotedImplicitByRefArgs()
 
                 for (unsigned fieldLclNum = fieldLclStart; fieldLclNum < fieldLclStop; ++fieldLclNum)
                 {
+                    JITDUMP("Fixing pointer for field V%02d from V%02d to V%02d\n", fieldLclNum, lclNum, structLclNum);
+
                     // Fix the pointer to the parent local.
                     LclVarDsc* fieldVarDsc = &lvaTable[fieldLclNum];
                     assert(fieldVarDsc->lvParentLcl == lclNum);

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

@@ -0,0 +1,43 @@
+// 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.Runtime.CompilerServices;
+using System.Runtime.InteropServices;
+
+[StructLayout(LayoutKind.Sequential)]
+internal struct AA
+{
+    public short tmp1;
+    public short q;
+
+    public ushort tmp2;
+    public int tmp3;
+
+    public AA(short qq)
+    {
+        tmp1 = 106;
+        tmp2 = 107;
+        tmp3 = 108;
+        q = qq;
+    }
+    
+    // The test verifies that we accurately update the byref variable that is a field of struct.
+    public static short call_target_ref(ref short arg) { arg = 100; return arg; }
+}
+
+   
+public class Runtime_57912
+{
+
+    public static int Main()
+    {
+        return (int)test_0_17(100, new AA(100), new AA(0));
+    }
+    
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    static short test_0_17(int num, AA init, AA zero)
+    {
+        return AA.call_target_ref(ref init.q);
+    }
+}

src/tests/JIT/Regression/JitBlue/Runtime_57912/Runtime_57912.csproj

@@ -0,0 +1,10 @@
+<Project Sdk="Microsoft.NET.Sdk">
+  <PropertyGroup>
+    <OutputType>Exe</OutputType>
+    <Optimize>True</Optimize>
+    <DebugType>None</DebugType>
+  </PropertyGroup>
+  <ItemGroup>
+    <Compile Include="$(MSBuildProjectName).cs" />
+  </ItemGroup>
+</Project>