JIT: array allocation fixes (#112676)

Fix some issues with array stack allocation that are currently hard to observe:
* use correct method table for shared array types
* pad array layout out to a multiple of TARGET_POINTER_SIZE

It's currently not easy to observe the type of a stack allocated array as
accessing the type causes escape. But this becomes possible when supporting
span capture or once we have a version of the covariant store check that can
handle stack allocated arrays.

Non-padded layouts may mean elements just off the end of an odd-length
short or byte array won't be zeroed, which some unsafe uses may expect.

Added some test cases.

Author: AndyAyersMS

src/coreclr/jit/helperexpansion.cpp

@@ -2822,6 +2822,7 @@ bool Compiler::fgExpandStackArrayAllocation(BasicBlock* block, Statement* stmt,
 
     const CorInfoHelpFunc helper         = eeGetHelperNum(call->gtCallMethHnd);
     int                   lengthArgIndex = -1;
+    int                   typeArgIndex   = -1;
 
     switch (helper)
     {
@@ -2830,10 +2831,7 @@ bool Compiler::fgExpandStackArrayAllocation(BasicBlock* block, Statement* stmt,
         case CORINFO_HELP_NEWARR_1_OBJ:
         case CORINFO_HELP_NEWARR_1_ALIGN8:
             lengthArgIndex = 1;
-            break;
-
-        case CORINFO_HELP_READYTORUN_NEWARR_1:
-            lengthArgIndex = 0;
+            typeArgIndex   = 0;
             break;
 
         default:
@@ -2871,9 +2869,7 @@ bool Compiler::fgExpandStackArrayAllocation(BasicBlock* block, Statement* stmt,
 
     // Initialize the array method table pointer.
     //
-    CORINFO_CLASS_HANDLE arrayHnd = (CORINFO_CLASS_HANDLE)call->compileTimeHelperArgumentHandle;
-
-    GenTree* const   mt      = gtNewIconEmbClsHndNode(arrayHnd);
+    GenTree* const   mt      = call->gtArgs.GetArgByIndex(typeArgIndex)->GetNode();
     GenTree* const   mtStore = gtNewStoreValueNode(TYP_I_IMPL, stackLocalAddress, mt);
     Statement* const mtStmt  = fgNewStmtFromTree(mtStore);
 

src/coreclr/jit/layout.cpp

@@ -778,6 +778,7 @@ ClassLayoutBuilder ClassLayoutBuilder::BuildArray(Compiler* compiler, CORINFO_CL
 
     ClrSafeInt<unsigned> totalSize(elementSize);
     totalSize *= static_cast<unsigned>(length);
+    totalSize.AlignUp(TARGET_POINTER_SIZE);
     totalSize += static_cast<unsigned>(OFFSETOF__CORINFO_Array__data);
     assert(!totalSize.IsOverflow());
 

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

@@ -156,7 +156,11 @@ public static int TestEntryPoint()
             // Stack allocation of boxed structs is now enabled
             CallTestAndVerifyAllocation(BoxSimpleStructAndAddFields, 12, expectedAllocationKind);
 
+            // Fixed-sized stack array cases
             CallTestAndVerifyAllocation(AllocateArrayWithNonGCElements, 84, expectedAllocationKind);
+            CallTestAndVerifyAllocation(AllocateArrayWithGCElements, 84, expectedAllocationKind);
+            CallTestAndVerifyAllocation(AllocateArrayT<int>, 84, expectedAllocationKind);
+            CallTestAndVerifyAllocation(AllocateArrayT<string>, 84, expectedAllocationKind);
 
             // The remaining tests currently never allocate on the stack
             if (expectedAllocationKind == AllocationKind.Stack) {
@@ -170,6 +174,7 @@ public static int TestEntryPoint()
             CallTestAndVerifyAllocation(AllocateSimpleClassAndCast, 7, expectedAllocationKind);
 
             CallTestAndVerifyAllocation(AllocateArrayWithNonGCElementsEscape, 42, expectedAllocationKind);
+            CallTestAndVerifyAllocation(AllocateArrayWithGCElementsEscape, 42, expectedAllocationKind);
 
             // This test calls CORINFO_HELP_OVERFLOW
             CallTestAndVerifyAllocation(AllocateArrayWithNonGCElementsOutOfRangeLeft, 0, expectedAllocationKind, true);
@@ -375,6 +380,29 @@ static int AllocateArrayWithNonGCElements()
             return array[24] + array.Length;
         }
 
+        [MethodImpl(MethodImplOptions.NoInlining)]
+        static int AllocateArrayWithGCElements()
+        {
+            string[] array = new string[42];
+            array[24] = "42";
+            GC.Collect();
+            return array[24].Length * 21 + array.Length;
+        }
+
+        [MethodImpl(MethodImplOptions.NoInlining)]
+        static int AllocateArrayT<T>()
+        {
+            T[] array = new T[42];
+            T t = array[24];
+            GC.Collect();
+
+            // Todo -- validate array type (currently causes escape for shared)
+            // Todo -- store to array (currently causes escape for shared)
+
+            Consume(t);
+            return array.Length + 42;
+        }
+
         [MethodImpl(MethodImplOptions.NoInlining)]
         static int AllocateArrayWithNonGCElementsEscape()
         {
@@ -384,6 +412,15 @@ static int AllocateArrayWithNonGCElementsEscape()
             return array[24];
         }
 
+        [MethodImpl(MethodImplOptions.NoInlining)]
+        static int AllocateArrayWithGCElementsEscape()
+        {
+            string[] array = new string[42];
+            Use(ref array[24]);
+            GC.Collect();
+            return array[24].Length * 21;
+        }
+
         [MethodImpl(MethodImplOptions.NoInlining)]
         static int AllocateArrayWithNonGCElementsOutOfRangeRight()
         {
@@ -424,6 +461,12 @@ static void Use(ref int v)
             v = 42;
         }
 
+        [MethodImpl(MethodImplOptions.NoInlining)]
+        static void Use(ref string s)
+        {
+            s = "42";
+        }
+
         [MethodImpl(MethodImplOptions.NoInlining)]
         private static void ZeroAllocTest()
         {