Merge pull request #22528 from mikedn/idx-addr

Fix genCodeForIndexAddr

2 files changed, 61 insertions, 55 deletions

diff --git a/src/jit/codegenxarch.cpp b/src/jit/codegenxarch.cpp
index f80feeb088..6a3228392e 100644
--- a/src/jit/codegenxarch.cpp
+++ b/src/jit/codegenxarch.cpp
@@ -4614,79 +4614,82 @@ void CodeGen::genCodeForIndexAddr(GenTreeIndexAddr* node)
     GenTree* const base  = node->Arr();
     GenTree* const index = node->Index();
 
-    genConsumeReg(base);
-    genConsumeReg(index);
+    const regNumber baseReg  = genConsumeReg(base);
+    regNumber       indexReg = genConsumeReg(index);
+    const regNumber dstReg   = node->gtRegNum;
 
     // NOTE: `genConsumeReg` marks the consumed register as not a GC pointer, as it assumes that the input registers
     // die at the first instruction generated by the node. This is not the case for `INDEX_ADDR`, however, as the
     // base register is multiply-used. As such, we need to mark the base register as containing a GC pointer until
     // we are finished generating the code for this node.
 
-    gcInfo.gcMarkRegPtrVal(base->gtRegNum, base->TypeGet());
-    assert(!varTypeIsGC(index->TypeGet()));
+    gcInfo.gcMarkRegPtrVal(baseReg, base->TypeGet());
+    assert(varTypeIsIntegral(index->TypeGet()));
 
     regNumber tmpReg = REG_NA;
+#ifdef _TARGET_64BIT_
+    tmpReg = node->GetSingleTempReg();
+#endif
 
     // Generate the bounds check if necessary.
     if ((node->gtFlags & GTF_INX_RNGCHK) != 0)
     {
-        // Create a GT_IND(GT_LEA)) tree for the array length access.
-        GenTreeAddrMode arrLenAddr(base->TypeGet(), base, nullptr, 0, node->gtLenOffset);
-        arrLenAddr.gtRegNum = REG_NA;
-        arrLenAddr.SetContained();
-
-        GenTreeIndir arrLen = indirForm(TYP_INT, &arrLenAddr);
-
 #ifdef _TARGET_64BIT_
         // The CLI Spec allows an array to be indexed by either an int32 or a native int.  In the case that the index
-        // is a native int on a 64-bit platform, we will need to widen the array length and the compare.
+        // is a native int on a 64-bit platform, we will need to widen the array length and then compare.
         if (index->TypeGet() == TYP_I_IMPL)
         {
-            // Load the array length into a register.
-            tmpReg          = node->GetSingleTempReg();
-            arrLen.gtRegNum = tmpReg;
-            arrLen.ClearContained();
-            getEmitter()->emitInsLoadInd(ins_Load(TYP_INT), EA_4BYTE, arrLen.gtRegNum, &arrLen);
+            getEmitter()->emitIns_R_AR(INS_mov, EA_4BYTE, tmpReg, baseReg, static_cast<int>(node->gtLenOffset));
+            getEmitter()->emitIns_R_R(INS_cmp, EA_8BYTE, indexReg, tmpReg);
         }
         else
-#endif
+#endif // _TARGET_64BIT_
         {
-            assert(varTypeIsIntegral(index->TypeGet()));
-
-            arrLen.gtRegNum = REG_NA;
-            arrLen.SetContained();
+            getEmitter()->emitIns_R_AR(INS_cmp, EA_4BYTE, indexReg, baseReg, static_cast<int>(node->gtLenOffset));
         }
 
-        // Generate the range check.
-        getEmitter()->emitInsBinary(INS_cmp, emitTypeSize(TYP_I_IMPL), index, &arrLen);
         genJumpToThrowHlpBlk(EJ_jae, SCK_RNGCHK_FAIL, node->gtIndRngFailBB);
     }
 
+#ifdef _TARGET_64BIT_
+    if (index->TypeGet() != TYP_I_IMPL)
+    {
+        // LEA needs 64-bit operands so we need to widen the index if it's TYP_INT.
+        getEmitter()->emitIns_R_R(INS_mov, EA_4BYTE, tmpReg, indexReg);
+        indexReg = tmpReg;
+    }
+#endif // _TARGET_64BIT_
+
     // Compute the address of the array element.
-    switch (node->gtElemSize)
+    unsigned scale = node->gtElemSize;
+
+    switch (scale)
     {
         case 1:
         case 2:
         case 4:
         case 8:
-            getEmitter()->emitIns_R_ARX(INS_lea, emitTypeSize(node), node->gtRegNum, base->gtRegNum, index->gtRegNum,
-                                        node->gtElemSize, static_cast<int>(node->gtElemOffset));
+            tmpReg = indexReg;
             break;
 
         default:
-        {
-            // Multiply the index by the element size.
-            //
-            // TODO-CQ: this should really just use `imul index, index, #gtElemSize`
-            tmpReg = (tmpReg == REG_NA) ? node->GetSingleTempReg() : tmpReg;
-            CodeGen::genSetRegToIcon(tmpReg, (ssize_t)node->gtElemSize, TYP_INT);
-            inst_RV_RV(INS_imul, tmpReg, index->gtRegNum);
-            getEmitter()->emitIns_R_ARX(INS_lea, emitTypeSize(node), node->gtRegNum, base->gtRegNum, tmpReg, 1,
-                                        static_cast<int>(node->gtElemOffset));
+#ifdef _TARGET_64BIT_
+            // IMUL treats its immediate operand as signed so scale can't be larger than INT32_MAX.
+            // The VM doesn't allow such large array elements but let's be sure.
+            noway_assert(scale <= INT32_MAX);
+#else  // !_TARGET_64BIT_
+            tmpReg = node->GetSingleTempReg();
+#endif // !_TARGET_64BIT_
+
+            getEmitter()->emitIns_R_I(emitter::inst3opImulForReg(tmpReg), EA_PTRSIZE, indexReg,
+                                      static_cast<ssize_t>(scale));
+            scale = 1;
             break;
-        }
     }
 
+    getEmitter()->emitIns_R_ARX(INS_lea, emitTypeSize(node->TypeGet()), dstReg, baseReg, tmpReg, scale,
+                                static_cast<int>(node->gtElemOffset));
+
     gcInfo.gcMarkRegSetNpt(base->gtGetRegMask());
 
     genProduceReg(node);
diff --git a/src/jit/lsraxarch.cpp b/src/jit/lsraxarch.cpp
index 67fe539a2f..364a4b7103 100644
--- a/src/jit/lsraxarch.cpp
+++ b/src/jit/lsraxarch.cpp
@@ -666,25 +666,28 @@ int LinearScan::BuildNode(GenTree* tree)
         {
             assert(dstCount == 1);
             RefPosition* internalDef = nullptr;
-            if (tree->AsIndexAddr()->Index()->TypeGet() == TYP_I_IMPL)
-            {
-                internalDef = buildInternalIntRegisterDefForNode(tree);
-            }
-            else
-            {
-                switch (tree->AsIndexAddr()->gtElemSize)
-                {
-                    case 1:
-                    case 2:
-                    case 4:
-                    case 8:
-                        break;
-
-                    default:
-                        internalDef = buildInternalIntRegisterDefForNode(tree);
-                        break;
-                }
+#ifdef _TARGET_64BIT_
+            // On 64-bit we always need a temporary register:
+            //   - if the index is `native int` then we need to load the array
+            //     length into a register to widen it to `native int`
+            //   - if the index is `int` (or smaller) then we need to widen
+            //     it to `long` to peform the address calculation
+            internalDef = buildInternalIntRegisterDefForNode(tree);
+#else  // !_TARGET_64BIT_
+            assert(!varTypeIsLong(tree->AsIndexAddr()->Index()->TypeGet()));
+            switch (tree->AsIndexAddr()->gtElemSize)
+            {
+                case 1:
+                case 2:
+                case 4:
+                case 8:
+                    break;
+
+                default:
+                    internalDef = buildInternalIntRegisterDefForNode(tree);
+                    break;
             }
+#endif // !_TARGET_64BIT_
             srcCount = BuildBinaryUses(tree->AsOp());
             if (internalDef != nullptr)
             {