JIT: fix retyping of BLK ops when parent local is retyped (#116027)

Properly handle cases where we have updated the layout of a struct local (because
its fields may now refer to stack allocated objects) and then store or load from
just part of that local.

Fixes #115979.

Author: AndyAyersMS

src/coreclr/jit/objectalloc.cpp

@@ -2142,7 +2142,7 @@ void ObjectAllocator::UpdateAncestorTypes(
                     assert(parentType == TYP_BYREF);
                     parent->ChangeType(newType);
 
-                    // Propgate that upwards.
+                    // Propagate that upwards.
                     //
                     ++parentIndex;
                     keepChecking = true;
@@ -2195,20 +2195,40 @@ void ObjectAllocator::UpdateAncestorTypes(
                 else
                 {
                     assert(tree == parent->AsIndir()->Data());
-                    GenTree* const addr = parent->AsIndir()->Addr();
 
                     // If we are storing to a GC struct field, we may need to retype the store
                     //
-                    if (parent->OperIs(GT_STOREIND) && varTypeIsGC(parent->TypeGet()))
+                    if (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))
+                    else if (retypeFields && parent->OperIs(GT_STORE_BLK))
                     {
-                        parent->AsBlk()->SetLayout(newLayout);
+                        GenTreeBlk* const  block     = parent->AsBlk();
+                        ClassLayout* const oldLayout = block->GetLayout();
+
+                        if (oldLayout->HasGCPtr())
+                        {
+                            if (newLayout->GetSize() == oldLayout->GetSize())
+                            {
+                                block->SetLayout(newLayout);
+                            }
+                            else
+                            {
+                                // We must be storing just a portion of the original local
+                                //
+                                assert(newLayout->GetSize() > oldLayout->GetSize());
+
+                                if (newLayout->HasGCPtr())
+                                {
+                                    block->SetLayout(GetByrefLayout(oldLayout));
+                                }
+                                else
+                                {
+                                    block->SetLayout(GetNonGCLayout(oldLayout));
+                                }
+                            }
+                        }
                     }
                 }
                 break;
@@ -2219,19 +2239,54 @@ void ObjectAllocator::UpdateAncestorTypes(
             {
                 // If we are loading from a GC struct field, we may need to retype the load
                 //
-                if (retypeFields && (varTypeIsGC(parent->TypeGet())))
+                if (retypeFields)
                 {
-                    parent->ChangeType(newType);
+                    bool didRetype = false;
 
-                    if (parent->OperIs(GT_BLK))
+                    if (varTypeIsGC(parent->TypeGet()))
                     {
-                        parent->AsBlk()->SetLayout(newLayout);
+                        parent->ChangeType(newType);
+                        didRetype = true;
                     }
+                    else if (parent->OperIs(GT_BLK))
+                    {
+                        GenTreeBlk* const  block     = parent->AsBlk();
+                        ClassLayout* const oldLayout = block->GetLayout();
 
-                    ++parentIndex;
-                    keepChecking = true;
-                    retypeFields = false;
+                        if (oldLayout->HasGCPtr())
+                        {
+                            if (newLayout->GetSize() == oldLayout->GetSize())
+                            {
+                                block->SetLayout(newLayout);
+                            }
+                            else
+                            {
+                                // We must be loading just a portion of the original local
+                                //
+                                assert(newLayout->GetSize() > oldLayout->GetSize());
+
+                                if (newLayout->HasGCPtr())
+                                {
+                                    block->SetLayout(GetByrefLayout(oldLayout));
+                                }
+                                else
+                                {
+                                    block->SetLayout(GetNonGCLayout(oldLayout));
+                                }
+                            }
+
+                            didRetype = true;
+                        }
+                    }
+
+                    if (didRetype)
+                    {
+                        ++parentIndex;
+                        keepChecking = true;
+                        retypeFields = false;
+                    }
                 }
+
                 break;
             }
 

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

@@ -0,0 +1,68 @@
+// 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-25 16:50:09
+// Run on Arm64 MacOS
+// Seed: 18178428402533635742-vectort,vector64,vector128,armadvsimd,armadvsimdarm64,armaes,armarmbase,armarmbasearm64,armcrc32,armcrc32arm64,armdp,armrdm,armrdmarm64,armsha1,armsha256
+// Reduced from 127.3 KiB to 0.6 KiB in 00:01:07
+// 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 75; hash 0xade6b36b; FullOpts)
+// 
+//     File: /Users/runner/work/1/s/src/coreclr/jit/morphblock.cpp Line: 668
+// 
+
+using System;
+using System.Runtime.CompilerServices;
+using System.Runtime.Intrinsics;
+// using System.Numerics;
+using Xunit;
+
+public class C1
+{
+}
+
+public struct S1
+{
+    public Vector64<short> F0;
+    public long F2;
+    public C1 F6;
+    public ulong F7;
+}
+
+public struct S3
+{
+    public C1 F4;
+    public S1 F6;
+    public S3(C1 f4, S1 f6) : this()
+    {
+        F4 = f4;
+        F6 = f6;
+    }
+}
+
+public class Runtime_115979
+{
+    [Fact]
+    public static int Test()
+    {
+        int result = -1;
+        try
+        {
+            Problem();
+        }
+        catch (NullReferenceException)
+        {
+            result = 100;
+        }
+        return result;
+    }
+
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    static void Problem()
+    {
+        S1[] vr0 = default(S1[]);
+        S3 vr1 = new S3(new C1(), new S1());
+        S3 vr2 = new S3(vr0[0].F6, vr1.F6);
+        System.Console.WriteLine(vr2.F6.F0);
+    }
+}

src/tests/JIT/opt/ObjectStackAllocation/Runtime_115979.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>