Fix loop cloning of array of struct with array field (#67130)

Loop cloning needs to parse what morph creates from GT_INDEX nodes
to determine if there are array accesses with bounds checks that
could potentially be optimized. For jagged array access, this can
be a "comma chain" of bounds checks and array element address expressions.
For a case where an array of structs had a struct field, such as
`ValueTuple<int[], int>[]`, cloning was confusing the expression
`a[i].Item1[j]` for the jagged array access `a[i][j]`.

The fix here is to keep track of the type of the `GT_INDEX` node that is
being morphed, in the `GT_BOUNDS_CHECK` node that is created for it.
(This is the only thing cloning parses, to avoid the need to parse
the very complex trees morph can create.) This type is then checked
when parsing the "comma chain" trees. If a non-`TYP_REF` is found (such
as a `TYP_STRUCT` in the above example), no more levels of array indexing
are considered. (`TYP_REF` is what an array object would have, for a jagged
array.)

Fixes #66254.

Author: BruceForstall

src/coreclr/jit/compiler.h

@@ -8088,8 +8088,9 @@ class Compiler
     };
 
     bool optIsStackLocalInvariant(unsigned loopNum, unsigned lclNum);
-    bool optExtractArrIndex(GenTree* tree, ArrIndex* result, unsigned lhsNum);
-    bool optReconstructArrIndex(GenTree* tree, ArrIndex* result, unsigned lhsNum);
+    bool optExtractArrIndex(GenTree* tree, ArrIndex* result, unsigned lhsNum, bool* topLevelIsFinal);
+    bool optReconstructArrIndexHelp(GenTree* tree, ArrIndex* result, unsigned lhsNum, bool* topLevelIsFinal);
+    bool optReconstructArrIndex(GenTree* tree, ArrIndex* result);
     bool optIdentifyLoopOptInfo(unsigned loopNum, LoopCloneContext* context);
     static fgWalkPreFn optCanOptimizeByLoopCloningVisitor;
     fgWalkResult optCanOptimizeByLoopCloning(GenTree* tree, LoopCloneVisitorInfo* info);

src/coreclr/jit/compiler.hpp

@@ -3556,7 +3556,7 @@ inline bool Compiler::LoopDsc::lpArrLenLimit(Compiler* comp, ArrIndex* index) co
     // We have a[i].length, extract a[i] pattern.
     else if (limit->AsArrLen()->ArrRef()->gtOper == GT_COMMA)
     {
-        return comp->optReconstructArrIndex(limit->AsArrLen()->ArrRef(), index, BAD_VAR_NUM);
+        return comp->optReconstructArrIndex(limit->AsArrLen()->ArrRef(), index);
     }
     return false;
 }

src/coreclr/jit/gentree.cpp

@@ -7664,6 +7664,7 @@ GenTree* Compiler::gtCloneExpr(
                     GenTreeBoundsChk(tree->AsBoundsChk()->GetIndex(), tree->AsBoundsChk()->GetArrayLength(),
                                      tree->AsBoundsChk()->gtThrowKind);
                 copy->AsBoundsChk()->gtIndRngFailBB = tree->AsBoundsChk()->gtIndRngFailBB;
+                copy->AsBoundsChk()->gtInxType      = tree->AsBoundsChk()->gtInxType;
                 break;
 
             case GT_LEA:

src/coreclr/jit/gentree.h

@@ -5876,8 +5876,16 @@ struct GenTreeBoundsChk : public GenTreeOp
     BasicBlock*     gtIndRngFailBB; // Basic block to jump to for index-out-of-range
     SpecialCodeKind gtThrowKind;    // Kind of throw block to branch to on failure
 
+    // Store some information about the array element type that was in the GT_INDEX node before morphing.
+    // Note that this information is also stored in the m_arrayInfoMap of the morphed IND node (that
+    // is marked with GTF_IND_ARR_INDEX), but that can be hard to find.
+    var_types gtInxType; // Save the GT_INDEX type
+
     GenTreeBoundsChk(GenTree* index, GenTree* length, SpecialCodeKind kind)
-        : GenTreeOp(GT_BOUNDS_CHECK, TYP_VOID, index, length), gtIndRngFailBB(nullptr), gtThrowKind(kind)
+        : GenTreeOp(GT_BOUNDS_CHECK, TYP_VOID, index, length)
+        , gtIndRngFailBB(nullptr)
+        , gtThrowKind(kind)
+        , gtInxType(TYP_UNKNOWN)
     {
         gtFlags |= GTF_EXCEPT;
     }

src/coreclr/jit/loopcloning.cpp

@@ -2140,9 +2140,10 @@ bool Compiler::optIsStackLocalInvariant(unsigned loopNum, unsigned lclNum)
 //  optExtractArrIndex: Try to extract the array index from "tree".
 //
 //  Arguments:
-//      tree        the tree to be checked if it is the array [] operation.
-//      result      the extracted GT_INDEX information is updated in result.
-//      lhsNum      for the root level (function is recursive) callers should pass BAD_VAR_NUM.
+//      tree             the tree to be checked if it is the array [] operation.
+//      result           the extracted GT_INDEX information is updated in result.
+//      lhsNum           for the root level (function is recursive) callers should pass BAD_VAR_NUM.
+//      topLevelIsFinal  OUT: set to `true` if see a non-TYP_REF element type array.
 //
 //  Return Value:
 //      Returns true if array index can be extracted, else, return false. See assumption about
@@ -2203,7 +2204,7 @@ bool Compiler::optIsStackLocalInvariant(unsigned loopNum, unsigned lclNum)
 // used as an index expression, or array base var is used as the array base. This saves us from parsing
 // all the forms that morph can create, especially for arrays of structs.
 //
-bool Compiler::optExtractArrIndex(GenTree* tree, ArrIndex* result, unsigned lhsNum)
+bool Compiler::optExtractArrIndex(GenTree* tree, ArrIndex* result, unsigned lhsNum, bool* topLevelIsFinal)
 {
     if (tree->gtOper != GT_COMMA)
     {
@@ -2247,37 +2248,31 @@ bool Compiler::optExtractArrIndex(GenTree* tree, ArrIndex* result, unsigned lhsN
     result->useBlock = compCurBB;
     result->rank++;
 
+    // If the array element type (saved from the GT_INDEX node during morphing) is anything but
+    // TYP_REF, then it must the the final level of jagged array.
+    assert(arrBndsChk->gtInxType != TYP_VOID);
+    *topLevelIsFinal = (arrBndsChk->gtInxType != TYP_REF);
+
     return true;
 }
 
 //---------------------------------------------------------------------------------------------------------------
-//  optReconstructArrIndex: Reconstruct array index.
+//  optReconstructArrIndexHelp: Helper function for optReconstructArrIndex. See that function for more details.
 //
 //  Arguments:
-//      tree        the tree to be checked if it is an array [][][] operation.
-//      result      OUT: the extracted GT_INDEX information.
-//      lhsNum      for the root level (function is recursive) callers should pass BAD_VAR_NUM.
+//      tree             the tree to be checked if it is an array [][][] operation.
+//      result           OUT: the extracted GT_INDEX information.
+//      lhsNum           var number of array object we're looking for.
+//      topLevelIsFinal  OUT: set to `true` if we reached a non-TYP_REF element type array.
 //
 //  Return Value:
 //      Returns true if array index can be extracted, else, return false. "rank" field in
-//      "result" contains the array access depth. The "indLcls" fields contain the indices.
-//
-//  Operation:
-//      Recursively look for a list of array indices. For example, if the tree is
-//          V03 = (V05 = V00[V01]), V05[V02]
-//      that corresponds to access of V00[V01][V02]. The return value would then be:
-//      ArrIndex result { arrLcl: V00, indLcls: [V01, V02], rank: 2 }
-//
-//      Note that the array expression is implied by the array bounds check under the COMMA, and the array bounds
-//      checks is what is parsed from the morphed tree; the array addressing expression is not parsed.
-//
-//  Assumption:
-//      The method extracts only if the array base and indices are GT_LCL_VAR.
+//      "result" contains the array access depth. The "indLcls" field contains the indices.
 //
-bool Compiler::optReconstructArrIndex(GenTree* tree, ArrIndex* result, unsigned lhsNum)
+bool Compiler::optReconstructArrIndexHelp(GenTree* tree, ArrIndex* result, unsigned lhsNum, bool* topLevelIsFinal)
 {
     // If we can extract "tree" (which is a top level comma) return.
-    if (optExtractArrIndex(tree, result, lhsNum))
+    if (optExtractArrIndex(tree, result, lhsNum, topLevelIsFinal))
     {
         return true;
     }
@@ -2294,18 +2289,152 @@ bool Compiler::optReconstructArrIndex(GenTree* tree, ArrIndex* result, unsigned
         GenTree* rhs = before->gtGetOp2();
 
         // "rhs" should contain an GT_INDEX
-        if (!lhs->IsLocal() || !optReconstructArrIndex(rhs, result, lhsNum))
+        if (!lhs->IsLocal() || !optReconstructArrIndexHelp(rhs, result, lhsNum, topLevelIsFinal))
         {
             return false;
         }
+
+        // If rhs represents an array of elements other than arrays (e.g., an array of structs),
+        // then we can't go any farther.
+        if (*topLevelIsFinal)
+        {
+            return false;
+        }
+
         unsigned lhsNum = lhs->AsLclVarCommon()->GetLclNum();
         GenTree* after  = tree->gtGetOp2();
         // Pass the "lhsNum", so we can verify if indeed it is used as the array base.
-        return optExtractArrIndex(after, result, lhsNum);
+        return optExtractArrIndex(after, result, lhsNum, topLevelIsFinal);
     }
     return false;
 }
 
+//---------------------------------------------------------------------------------------------------------------
+//  optReconstructArrIndex: Reconstruct array index from a post-morph tree.
+//
+//  Arguments:
+//      tree        the tree to be checked if it is an array [][][] operation.
+//      result      OUT: the extracted GT_INDEX information.
+//
+//  Return Value:
+//      Returns true if array index can be extracted, else, return false. "rank" field in
+//      "result" contains the array access depth. The "indLcls" field contains the indices.
+//
+//  Operation:
+//      Recursively look for a list of array indices. For example, if the tree is
+//          V03 = (V05 = V00[V01]), V05[V02]
+//      that corresponds to access of V00[V01][V02]. The return value would then be:
+//      ArrIndex result { arrLcl: V00, indLcls: [V01, V02], rank: 2 }
+//
+//      Note that the array expression is implied by the array bounds check under the COMMA, and the array bounds
+//      checks is what is parsed from the morphed tree; the array addressing expression is not parsed.
+//      However, the array bounds checks are not quite sufficient because of the way "morph" alters the trees.
+//      Specifically, we normally see a COMMA node with a LHS of the morphed array INDEX expression and RHS
+//      of the bounds check. E.g., for int[][], a[i][j] we have a pre-morph tree:
+//
+// \--*  INDEX     int
+//   +--*  INDEX     ref
+//   |  +--*  LCL_VAR   ref    V00 loc0
+//   |  \--*  LCL_VAR   int    V02 loc2
+//   \--*  LCL_VAR   int    V03 loc3
+//
+//      and post-morph tree:
+//
+// \--*  COMMA     int
+//    +--*  ASG       ref
+//    |  +--*  LCL_VAR   ref    V19 tmp12
+//    |  \--*  COMMA     ref
+//    |     +--*  BOUNDS_CHECK_Rng void
+//    |     |  +--*  LCL_VAR   int    V02 loc2
+//    |     |  \--*  ARR_LENGTH int
+//    |     |     \--*  LCL_VAR   ref    V00 loc0
+//    |     \--*  IND       ref
+//    |        \--*  ADD       byref
+//    |           +--*  LCL_VAR   ref    V00 loc0
+//    |           \--*  ADD       long
+//    |              +--*  LSH       long
+//    |              |  +--*  CAST      long <- uint
+//    |              |  |  \--*  LCL_VAR   int    V02 loc2
+//    |              |  \--*  CNS_INT   long   3
+//    |              \--*  CNS_INT   long   16 Fseq[#FirstElem]
+//    \--*  COMMA     int
+//       +--*  BOUNDS_CHECK_Rng void
+//       |  +--*  LCL_VAR   int    V03 loc3
+//       |  \--*  ARR_LENGTH int
+//       |     \--*  LCL_VAR   ref    V19 tmp12
+//       \--*  IND       int
+//          \--*  ADD       byref
+//             +--*  LCL_VAR   ref    V19 tmp12
+//             \--*  ADD       long
+//                +--*  LSH       long
+//                |  +--*  CAST      long <- uint
+//                |  |  \--*  LCL_VAR   int    V03 loc3
+//                |  \--*  CNS_INT   long   2
+//                \--*  CNS_INT   long   16 Fseq[#FirstElem]
+//
+//      However, for an array of structs that contains an array field, e.g. ValueTuple<int[], int>[], expression
+//      a[i].Item1[j],
+//
+// \--*  INDEX     int
+//    +--*  FIELD     ref    Item1
+//       |  \--*  ADDR      byref
+//       |     \--*  INDEX     struct<System.ValueTuple`2[System.Int32[],System.Int32], 16>
+//       |        +--*  LCL_VAR   ref    V01 loc1
+//       |        \--*  LCL_VAR   int    V04 loc4
+//       \--*  LCL_VAR   int    V06 loc6
+//
+//      Morph "hoists" the bounds check above the struct field access:
+//
+// \--*  COMMA     int
+//    +--*  ASG       ref
+//    |  +--*  LCL_VAR   ref    V23 tmp16
+//    |  \--*  COMMA     ref
+//    |     +--*  BOUNDS_CHECK_Rng void
+//    |     |  +--*  LCL_VAR   int    V04 loc4
+//    |     |  \--*  ARR_LENGTH int
+//    |     |     \--*  LCL_VAR   ref    V01 loc1
+//    |     \--*  IND       ref
+//    |        \--*  ADDR      byref  Zero Fseq[Item1]
+//    |           \--*  IND       struct<System.ValueTuple`2[System.Int32[],System.Int32], 16>
+//    |              \--*  ADD       byref
+//    |                 +--*  LCL_VAR   ref    V01 loc1
+//    |                 \--*  ADD       long
+//    |                    +--*  LSH       long
+//    |                    |  +--*  CAST      long <- uint
+//    |                    |  |  \--*  LCL_VAR   int    V04 loc4
+//    |                    |  \--*  CNS_INT   long   4
+//    |                    \--*  CNS_INT   long   16 Fseq[#FirstElem]
+//    \--*  COMMA     int
+//       +--*  BOUNDS_CHECK_Rng void
+//       |  +--*  LCL_VAR   int    V06 loc6
+//       |  \--*  ARR_LENGTH int
+//       |     \--*  LCL_VAR   ref    V23 tmp16
+//       \--*  IND       int
+//          \--*  ADD       byref
+//             +--*  LCL_VAR   ref    V23 tmp16
+//             \--*  ADD       long
+//                +--*  LSH       long
+//                |  +--*  CAST      long <- uint
+//                |  |  \--*  LCL_VAR   int    V06 loc6
+//                |  \--*  CNS_INT   long   2
+//                \--*  CNS_INT   long   16 Fseq[#FirstElem]
+//
+//      This should not be parsed as a jagged array (e.g., a[i][j]). To ensure that it is not, the type of the
+//      GT_INDEX node is stashed in the GT_BOUNDS_CHECK node during morph. If we see a bounds check node where
+//      the GT_INDEX was not TYP_REF, then it must be the outermost jagged array level. E.g., if it is
+//      TYP_STRUCT, then we have an array of structs, and any further bounds checks must be of one of its fields.
+//
+//      It would be much better if we didn't need to parse these trees at all, and did all this work pre-morph.
+//
+//  Assumption:
+//      The method extracts only if the array base and indices are GT_LCL_VAR.
+//
+bool Compiler::optReconstructArrIndex(GenTree* tree, ArrIndex* result)
+{
+    bool topLevelIsFinal = false;
+    return optReconstructArrIndexHelp(tree, result, BAD_VAR_NUM, &topLevelIsFinal);
+}
+
 //----------------------------------------------------------------------------------------------
 //  optCanOptimizeByLoopCloning: Check if the tree can be optimized by loop cloning and if so,
 //      identify as potential candidate and update the loop context.
@@ -2329,7 +2458,7 @@ Compiler::fgWalkResult Compiler::optCanOptimizeByLoopCloning(GenTree* tree, Loop
     ArrIndex arrIndex(getAllocator(CMK_LoopClone));
 
     // Check if array index can be optimized.
-    if (optReconstructArrIndex(tree, &arrIndex, BAD_VAR_NUM))
+    if (optReconstructArrIndex(tree, &arrIndex))
     {
         assert(tree->gtOper == GT_COMMA);
 

src/coreclr/jit/morph.cpp

@@ -5436,6 +5436,7 @@ GenTree* Compiler::fgMorphArrayIndex(GenTree* tree)
         }
 
         GenTreeBoundsChk* arrBndsChk = new (this, GT_BOUNDS_CHECK) GenTreeBoundsChk(index, arrLen, SCK_RNGCHK_FAIL);
+        arrBndsChk->gtInxType        = elemTyp;
 
         bndsChk = arrBndsChk;
 

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

@@ -0,0 +1,104 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+// Test that loop cloning won't consider a[i].struct_field[j] to be
+// a jagged array a[i][j].
+
+using System;
+
+public class Runtime_66254
+{
+    public static void t1()
+    {
+        var a = new ValueTuple<int[], int>[]
+        {
+            new (new[] { 0 }, 1),
+            new (new[] { 2 }, 3)
+        };
+
+        for (var i1 = 0; i1 < a.Length; i1++)
+        {
+            for (var i2 = 0; i2 < a[i1].Item1.Length; i2++)
+            {
+                var elem = a[i1].Item1[i2];
+                Console.WriteLine(elem);
+            }
+        }
+    }
+
+    public static void t2()
+    {
+        var a = new ValueTuple<int[], int>[]
+        {
+            new (new[] { 0 }, 1),
+            new (new[] { 2 }, 3)
+        };
+
+        for (var i1 = 0; i1 < a.Length; i1++)
+        {
+            var length = a[i1].Item1.Length;
+            for (var i2 = 0; i2 < length; i2++)
+            {
+                var elem = a[i1].Item1[i2];
+                Console.WriteLine(elem);
+            }
+        }
+    }
+
+    public static void t3()
+    {
+        var a = new ValueTuple<int, int[]>[]
+        {
+            new (1, new[] { 2 }),
+            new (3, new[] { 4 })
+        };
+
+        for (var i1 = 0; i1 < a.Length; i1++)
+        {
+            var length = a[i1].Item2.Length;
+            for (var i2 = 0; i2 < length; i2++)
+            {
+                var elem = a[i1].Item2[i2];
+                Console.WriteLine(elem);
+            }
+        }
+    }
+
+    public static int Main()
+    {
+        int result = 100;
+
+        try
+        {
+            t1();
+        }
+        catch (Exception e)
+        {
+            Console.WriteLine($"t1 failed: {e}");
+            result = 101;
+        }
+
+        try
+        {
+            t2();
+        }
+        catch (Exception e)
+        {
+            Console.WriteLine($"t2 failed: {e}");
+            result = 101;
+        }
+
+        try
+        {
+            t3();
+        }
+        catch (Exception e)
+        {
+            Console.WriteLine($"t3 failed: {e}");
+            result = 101;
+        }
+
+        Console.WriteLine((result == 100) ? "PASS" : "FAIL");
+        return result;
+    }
+}