JIT: update layout of BLK nodes during escape analysis (#115845)

If we retype a struct local we also need to retype any BLK operations
on those locals.

Fixes #115832.

Author: AndyAyersMS

src/coreclr/jit/gentree.h

@@ -7867,6 +7867,13 @@ struct GenTreeBlk : public GenTreeIndir
         return m_layout;
     }
 
+    void SetLayout(ClassLayout* newLayout)
+    {
+        assert(newLayout != nullptr);
+        assert(newLayout->GetSize() == m_layout->GetSize());
+        m_layout = newLayout;
+    }
+
     // The data to be stored (null for GT_BLK)
     GenTree*& Data()
     {

src/coreclr/jit/objectalloc.cpp

@@ -2020,6 +2020,7 @@ void ObjectAllocator::AnalyzeParentStack(ArrayStack<GenTree*>* parentStack, unsi
 //    tree            - Possibly-stack-pointing tree
 //    parentStack     - Parent stack of the possibly-stack-pointing tree
 //    newType         - New type of the possibly-stack-pointing tree
+//    newLayout       - Layout for a retyped local struct
 //    retypeFields    - Inspiring local is a retyped local struct; retype fields.
 //
 // Notes:
@@ -2028,10 +2029,8 @@ void ObjectAllocator::AnalyzeParentStack(ArrayStack<GenTree*>* parentStack, unsi
 //                      In addition to updating types this method may set GTF_IND_TGT_NOT_HEAP on ancestor
 //                      indirections to help codegen with write barrier selection.
 //
-void ObjectAllocator::UpdateAncestorTypes(GenTree*              tree,
-                                          ArrayStack<GenTree*>* parentStack,
-                                          var_types             newType,
-                                          bool                  retypeFields)
+void ObjectAllocator::UpdateAncestorTypes(
+    GenTree* tree, ArrayStack<GenTree*>* parentStack, var_types newType, ClassLayout* newLayout, bool retypeFields)
 {
     assert(newType == TYP_BYREF || newType == TYP_I_IMPL);
     assert(parentStack != nullptr);
@@ -2200,10 +2199,17 @@ void ObjectAllocator::UpdateAncestorTypes(GenTree*              tree,
 
                     // If we are storing to a GC struct field, we may need to retype the store
                     //
-                    if (parent->OperIs(GT_STOREIND) && addr->OperIs(GT_FIELD_ADDR) && varTypeIsGC(parent->TypeGet()))
+                    if (parent->OperIs(GT_STOREIND) && varTypeIsGC(parent->TypeGet()))
                     {
                         parent->ChangeType(newType);
                     }
+
+                    // If we are storing a struct, we may need to change the layout
+                    //
+                    if (retypeFields && parent->OperIs(GT_STORE_BLK))
+                    {
+                        parent->AsBlk()->SetLayout(newLayout);
+                    }
                 }
                 break;
             }
@@ -2216,6 +2222,12 @@ void ObjectAllocator::UpdateAncestorTypes(GenTree*              tree,
                 if (retypeFields && (varTypeIsGC(parent->TypeGet())))
                 {
                     parent->ChangeType(newType);
+
+                    if (parent->OperIs(GT_BLK))
+                    {
+                        parent->AsBlk()->SetLayout(newLayout);
+                    }
+
                     ++parentIndex;
                     keepChecking = true;
                     retypeFields = false;
@@ -2290,6 +2302,7 @@ void ObjectAllocator::RewriteUses()
 
             unsigned int newLclNum = BAD_VAR_NUM;
             var_types    newType   = lclVarDsc->TypeGet();
+            ClassLayout* newLayout = nullptr;
 
             if (m_allocator->m_HeapLocalToStackLocalMap.TryGetValue(lclNum, &newLclNum))
             {
@@ -2303,16 +2316,16 @@ void ObjectAllocator::RewriteUses()
             }
             else if (newType == TYP_STRUCT)
             {
-                ClassLayout* const layout = lclVarDsc->GetLayout();
-                newType                   = layout->HasGCPtr() ? TYP_BYREF : TYP_I_IMPL;
-                retypeFields              = true;
+                newLayout    = lclVarDsc->GetLayout();
+                newType      = newLayout->HasGCPtr() ? TYP_BYREF : TYP_I_IMPL;
+                retypeFields = true;
             }
             else
             {
                 tree->ChangeType(newType);
             }
 
-            m_allocator->UpdateAncestorTypes(tree, &m_ancestors, newType, retypeFields);
+            m_allocator->UpdateAncestorTypes(tree, &m_ancestors, newType, newLayout, retypeFields);
 
             return Compiler::fgWalkResult::WALK_CONTINUE;
         }

src/coreclr/jit/objectalloc.h

@@ -237,7 +237,8 @@ class ObjectAllocator final : public Phase
                                                Statement*           stmt);
     struct BuildConnGraphVisitorCallbackData;
     void AnalyzeParentStack(ArrayStack<GenTree*>* parentStack, unsigned int lclNum, BasicBlock* block);
-    void UpdateAncestorTypes(GenTree* tree, ArrayStack<GenTree*>* parentStack, var_types newType, bool retypeFields);
+    void UpdateAncestorTypes(
+        GenTree* tree, ArrayStack<GenTree*>* parentStack, var_types newType, ClassLayout* newLayout, bool retypeFields);
     ObjectAllocationType AllocationKind(GenTree* tree);
 
     // Conditionally escaping allocation support

src/tests/JIT/opt/ObjectStackAllocation/Runtime_115832.cs

@@ -0,0 +1,71 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+// Generated by Fuzzlyn v3.0 on 2025-05-20 22:26:35
+// Run on X64 Windows
+// Seed: 457555340882575514-vectort,vector128,vector256,x86aes,x86avx,x86avx2,x86avx512bw,x86avx512bwvl,x86avx512cd,x86avx512cdvl,x86avx512dq,x86avx512dqvl,x86avx512f,x86avx512fvl,x86avx512fx64,x86bmi1,x86bmi1x64,x86bmi2,x86bmi2x64,x86fma,x86lzcnt,x86lzcntx64,x86pclmulqdq,x86popcnt,x86popcntx64,x86sse,x86ssex64,x86sse2,x86sse2x64,x86sse3,x86sse41,x86sse41x64,x86sse42,x86sse42x64,x86ssse3,x86x86base
+// Reduced from 154.0 KiB to 0.8 KiB in 00:06:12
+// Hits JIT assert in Release:
+// Assertion failed 'm_blockLayout->CanAssignFrom(m_src->GetLayout(m_comp))' in 'Program:Main(Fuzzlyn.ExecutionServer.IRuntime)' during 'Morph - Global' (IL size 69; hash 0xade6b36b; FullOpts)
+// 
+//     File: D:\a\_work\1\s\src\coreclr\jit\morphblock.cpp Line: 668
+// 
+using System;
+using Xunit;
+
+public class C0
+{
+}
+
+public struct S0
+{
+    public int F5;
+}
+
+public struct S1
+{
+    public C0 F2;
+    public S0 F3;
+    public S1(C0 f2) : this()
+    {
+        F2 = f2;
+    }
+}
+
+public struct S2
+{
+    public S1 F5;
+    public S1 F7;
+    public S2(S1 f5) : this()
+    {
+        F5 = f5;
+    }
+}
+
+public struct S4
+{
+    public S2 F1;
+    public S4(S2 f1) : this()
+    {
+        F1 = f1;
+    }
+}
+
+public struct S5
+{
+    public S4 F5;
+    public S5(S4 f5) : this()
+    {
+        F5 = f5;
+    }
+}
+
+public class Runtime_115832
+{
+    [Fact]
+    public static void Problem()
+    {
+        S5 vr0 = new S5(new S4(new S2(new S1(new C0()))));
+        System.Console.WriteLine(vr0.F5.F1.F7.F3.F5);
+    }
+}

src/tests/JIT/opt/ObjectStackAllocation/Runtime_115832.csproj

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