Remove GT_ADDEX and replace with more generalized containment handling

src/coreclr/jit/codegen.h

@@ -1397,7 +1397,6 @@ XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
 #if defined(TARGET_ARM64)
     void genCodeForJumpCompare(GenTreeOp* tree);
     void genCodeForBfiz(GenTreeOp* tree);
-    void genCodeForAddEx(GenTreeOp* tree);
     void genCodeForCond(GenTreeOp* tree);
 #endif // TARGET_ARM64
 

src/coreclr/jit/codegenarm64.cpp

@@ -2639,6 +2639,62 @@ void CodeGen::genCodeForBinary(GenTreeOp* tree)
         genProduceReg(tree);
         return;
     }
+    else if (op2->OperIs(GT_CAST) && op2->isContained())
+    {
+        assert(varTypeIsIntegral(tree));
+
+        GenTree* a = op1;
+        GenTree* b = op2->AsCast()->CastOp();
+
+        instruction ins = genGetInsForOper(tree->OperGet(), targetType);
+        insOpts     opt = INS_OPTS_NONE;
+
+        if ((tree->gtFlags & GTF_SET_FLAGS) != 0)
+        {
+            // A subset of operations can still set flags
+
+            switch (oper)
+            {
+                case GT_ADD:
+                {
+                    ins = INS_adds;
+                    break;
+                }
+
+                case GT_SUB:
+                {
+                    ins = INS_subs;
+                    break;
+                }
+
+                default:
+                {
+                    noway_assert(!"Unexpected BinaryOp with GTF_SET_FLAGS set");
+                }
+            }
+        }
+
+        bool isZeroExtending = op2->AsCast()->IsZeroExtending();
+
+        if (varTypeIsByte(op2->CastToType()))
+        {
+            opt = isZeroExtending ? INS_OPTS_UXTB : INS_OPTS_SXTB;
+        }
+        else if (varTypeIsShort(op2->CastToType()))
+        {
+            opt = isZeroExtending ? INS_OPTS_UXTH : INS_OPTS_SXTH;
+        }
+        else
+        {
+            assert(op2->TypeIs(TYP_LONG) && genActualTypeIsInt(b));
+            opt = isZeroExtending ? INS_OPTS_UXTW : INS_OPTS_SXTW;
+        }
+
+        emit->emitIns_R_R_R(ins, emitActualTypeSize(tree), targetReg, a->GetRegNum(), b->GetRegNum(), opt);
+
+        genProduceReg(tree);
+        return;
+    }
 
     if (tree->OperIs(GT_AND) && op2->isContainedAndNotIntOrIImmed())
     {
@@ -10564,54 +10620,6 @@ void CodeGen::genCodeForBfiz(GenTreeOp* tree)
     genProduceReg(tree);
 }
 
-//------------------------------------------------------------------------
-// genCodeForAddEx: Generates the code sequence for a GenTree node that
-// represents an addition with sign or zero extended
-//
-// Arguments:
-//    tree - the add with extend node.
-//
-void CodeGen::genCodeForAddEx(GenTreeOp* tree)
-{
-    assert(tree->OperIs(GT_ADDEX));
-    genConsumeOperands(tree);
-
-    GenTree* op;
-    GenTree* containedOp;
-    if (tree->gtGetOp1()->isContained())
-    {
-        containedOp = tree->gtGetOp1();
-        op          = tree->gtGetOp2();
-    }
-    else
-    {
-        containedOp = tree->gtGetOp2();
-        op          = tree->gtGetOp1();
-    }
-    assert(containedOp->isContained() && !op->isContained());
-
-    regNumber dstReg = tree->GetRegNum();
-    regNumber op1Reg = op->GetRegNum();
-    regNumber op2Reg = containedOp->gtGetOp1()->GetRegNum();
-
-    if (containedOp->OperIs(GT_CAST))
-    {
-        GenTreeCast* cast = containedOp->AsCast();
-        assert(varTypeIsLong(cast->CastToType()));
-        insOpts opts = cast->IsUnsigned() ? INS_OPTS_UXTW : INS_OPTS_SXTW;
-        GetEmitter()->emitIns_R_R_R(tree->gtSetFlags() ? INS_adds : INS_add, emitActualTypeSize(tree), dstReg, op1Reg,
-                                    op2Reg, opts);
-    }
-    else
-    {
-        assert(containedOp->OperIs(GT_LSH));
-        ssize_t cns = containedOp->gtGetOp2()->AsIntCon()->IconValue();
-        GetEmitter()->emitIns_R_R_R_I(tree->gtSetFlags() ? INS_adds : INS_add, emitActualTypeSize(tree), dstReg, op1Reg,
-                                      op2Reg, cns, INS_OPTS_LSL);
-    }
-    genProduceReg(tree);
-}
-
 //------------------------------------------------------------------------
 // genCodeForCond: Generates the code sequence for a GenTree node that
 // represents a conditional instruction.

src/coreclr/jit/codegenarmarch.cpp

@@ -315,10 +315,6 @@ void CodeGen::genCodeForTreeNode(GenTree* treeNode)
             genCodeForSwap(treeNode->AsOp());
             break;
 
-        case GT_ADDEX:
-            genCodeForAddEx(treeNode->AsOp());
-            break;
-
         case GT_BFIZ:
             genCodeForBfiz(treeNode->AsOp());
             break;

src/coreclr/jit/codegencommon.cpp

@@ -1030,14 +1030,7 @@ bool CodeGen::genCreateAddrMode(
 
     if (!addr->OperIs(GT_ADD))
     {
-#if TARGET_ARM64
-        if (!addr->OperIs(GT_ADDEX))
-        {
-            return false;
-        }
-#else
         return false;
-#endif
     }
 
     GenTree* rv1 = nullptr;
@@ -1064,23 +1057,6 @@ bool CodeGen::genCreateAddrMode(
         op2 = addr->AsOp()->gtOp2;
     }
 
-#if TARGET_ARM64
-    if (addr->OperIs(GT_ADDEX))
-    {
-        if (op2->isContained() && op2->OperIs(GT_CAST))
-        {
-            *rv1Ptr = op1;
-            *rv2Ptr = op2;
-            *mulPtr = 1;
-            *cnsPtr = 0;
-            *revPtr = false; // op2 is never a gc type
-            assert(!varTypeIsGC(op2));
-            return true;
-        }
-        return false;
-    }
-#endif
-
     // Can't use indirect addressing mode as we need to check for overflow.
     // Also, can't use 'lea' as it doesn't set the flags.
 

src/coreclr/jit/gtlist.h

@@ -220,7 +220,6 @@ GTNODE(MUL_LONG         , GenTreeOp          ,1,GTK_BINOP|DBK_NOTHIR)
 GTNODE(AND_NOT          , GenTreeOp          ,0,GTK_BINOP|DBK_NOTHIR)
 
 #ifdef TARGET_ARM64
-GTNODE(ADDEX,             GenTreeOp          ,0,GTK_BINOP|DBK_NOTHIR) // Add with sign/zero extension.
 GTNODE(BFIZ             , GenTreeOp          ,0,GTK_BINOP|DBK_NOTHIR) // Bitfield Insert in Zero.
 GTNODE(CSNEG_MI         , GenTreeOp          ,0,GTK_BINOP|DBK_NOTHIR) // Conditional select, negate, minus result
 GTNODE(CNEG_LT          , GenTreeOp          ,0,GTK_UNOP|DBK_NOTHIR)  // Conditional, negate, signed less than result

src/coreclr/jit/lower.cpp

@@ -5213,14 +5213,7 @@ bool Lowering::TryCreateAddrMode(GenTree* addr, bool isContainable, GenTree* par
 {
     if (!addr->OperIs(GT_ADD) || addr->gtOverflow())
     {
-#ifdef TARGET_ARM64
-        if (!addr->OperIs(GT_ADDEX))
-        {
-            return false;
-        }
-#else
         return false;
-#endif
     }
 
 #ifdef TARGET_ARM64

src/coreclr/jit/lowerarmarch.cpp

@@ -244,17 +244,12 @@ bool Lowering::IsContainableBinaryOp(GenTree* parentNode, GenTree* childNode) co
             return false;
         }
 
-        ssize_t shiftAmount = shiftAmountNode->AsIntCon()->IconValue();
+        const ssize_t shiftAmount = shiftAmountNode->AsIntCon()->IconValue();
+        const ssize_t maxShift    = (static_cast<ssize_t>(genTypeSize(parentNode)) * BITS_IN_BYTE) - 1;
 
-        if ((shiftAmount < 0x01) || (shiftAmount > 0x3F))
+        if ((shiftAmount < 0x01) || (shiftAmount > maxShift))
         {
-            // Cannot contain if the shift amount is less than 1 or greater than 63
-            return false;
-        }
-
-        if (!varTypeIsLong(childNode) && (shiftAmount > 0x1F))
-        {
-            // Cannot contain if the shift amount is greater than 31
+            // Cannot contain if the shift amount is less than 1 or greater than maxShift
             return false;
         }
 
@@ -275,7 +270,7 @@ bool Lowering::IsContainableBinaryOp(GenTree* parentNode, GenTree* childNode) co
             return false;
         }
 
-        if (parentNode->OperIs(GT_CMP, GT_AND, GT_OR, GT_XOR))
+        if (parentNode->OperIs(GT_CMP, GT_OR, GT_XOR))
         {
             if (IsSafeToContainMem(parentNode, childNode))
             {
@@ -288,6 +283,64 @@ bool Lowering::IsContainableBinaryOp(GenTree* parentNode, GenTree* childNode) co
         return false;
     }
 
+    if (childNode->OperIs(GT_CAST))
+    {
+        // Find "a op cast(b)"
+        GenTree* castOp = childNode->AsCast()->CastOp();
+
+        // We want to prefer the combined op here over containment of the cast op
+        castOp->ClearContained();
+
+        bool isSupportedCast = false;
+
+        if (varTypeIsSmall(childNode->CastToType()))
+        {
+            // The JIT doesn't track upcasts from small types, instead most types
+            // are tracked as TYP_INT and then we get explicit downcasts to the
+            // desired small type instead.
+
+            assert(!varTypeIsFloating(castOp));
+            isSupportedCast = true;
+        }
+        else if (childNode->TypeIs(TYP_LONG) && genActualTypeIsInt(castOp))
+        {
+            // We can handle "INT -> LONG", "INT -> ULONG", "UINT -> LONG", and "UINT -> ULONG"
+            isSupportedCast = true;
+        }
+
+        if (!isSupportedCast)
+        {
+            return false;
+        }
+
+        if (parentNode->OperIs(GT_ADD, GT_SUB))
+        {
+            // These operations can still report flags
+
+            if (IsSafeToContainMem(parentNode, childNode))
+            {
+                return true;
+            }
+        }
+
+        if ((parentNode->gtFlags & GTF_SET_FLAGS) != 0)
+        {
+            // Cannot contain if the parent operation needs to set flags
+            return false;
+        }
+
+        if (parentNode->OperIs(GT_CMP))
+        {
+            if (IsSafeToContainMem(parentNode, childNode))
+            {
+                return true;
+            }
+        }
+
+        // TODO: Handle CMN
+        return false;
+    }
+
     return false;
 }
 #endif // TARGET_ARM64
@@ -1968,44 +2021,6 @@ void Lowering::ContainCheckBinary(GenTreeOp* node)
             return;
         }
     }
-
-    // Change ADD TO ADDEX for ADD(X, CAST(Y)) or ADD(CAST(X), Y) where CAST is int->long
-    // or for ADD(LSH(X, CNS), X) or ADD(X, LSH(X, CNS)) where CNS is in the (0..typeWidth) range
-    if (node->OperIs(GT_ADD) && !op1->isContained() && !op2->isContained() && varTypeIsIntegral(node) &&
-        !node->gtOverflow())
-    {
-        assert(!node->isContained());
-
-        if (op1->OperIs(GT_CAST) || op2->OperIs(GT_CAST))
-        {
-            GenTree* cast = op1->OperIs(GT_CAST) ? op1 : op2;
-            if (cast->gtGetOp1()->TypeIs(TYP_INT) && cast->TypeIs(TYP_LONG) && !cast->gtOverflow())
-            {
-                node->ChangeOper(GT_ADDEX);
-                cast->AsCast()->CastOp()->ClearContained(); // Uncontain any memory operands.
-                MakeSrcContained(node, cast);
-            }
-        }
-        else if (op1->OperIs(GT_LSH) || op2->OperIs(GT_LSH))
-        {
-            GenTree* lsh     = op1->OperIs(GT_LSH) ? op1 : op2;
-            GenTree* shiftBy = lsh->gtGetOp2();
-
-            if (shiftBy->IsCnsIntOrI())
-            {
-                const ssize_t shiftByCns = shiftBy->AsIntCon()->IconValue();
-                const ssize_t maxShift   = (ssize_t)genTypeSize(node) * BITS_IN_BYTE;
-
-                if ((shiftByCns > 0) && (shiftByCns < maxShift))
-                {
-                    // shiftBy is small so it has to be contained at this point.
-                    assert(shiftBy->isContained());
-                    node->ChangeOper(GT_ADDEX);
-                    MakeSrcContained(node, lsh);
-                }
-            }
-        }
-    }
 #endif
 }
 

src/coreclr/jit/lsraarm64.cpp

@@ -288,7 +288,6 @@ int LinearScan::BuildNode(GenTree* tree)
             }
             FALLTHROUGH;
 
-        case GT_ADDEX:
         case GT_AND:
         case GT_AND_NOT:
         case GT_OR: