Move JCMP transform to LowerJTrue

Unlike many other relop transforms we do this one is only triggerred by the presence of a conditional branch (JTRUE) so it makes more sense to do it when lowering JTRUE nodes, avoids unnecessary calls to TryGetUse.

src/jit/lower.cpp

@@ -178,12 +178,10 @@ GenTree* Lowering::LowerNode(GenTree* node)
         case GT_TEST_EQ:
         case GT_TEST_NE:
         case GT_CMP:
-        case GT_JCMP:
             return LowerCompare(node);
 
         case GT_JTRUE:
-            ContainCheckJTrue(node->AsOp());
-            break;
+            return LowerJTrue(node->AsOp());
 
         case GT_JMP:
             LowerJmpMethod(node);
@@ -2696,6 +2694,66 @@ GenTree* Lowering::LowerCompare(GenTree* cmp)
     return cmp->gtNext;
 }
 
+//------------------------------------------------------------------------
+// Lowering::LowerJTrue: Lowers a JTRUE node.
+//
+// Arguments:
+//    jtrue - the JTRUE node
+//
+// Return Value:
+//    The next node to lower (usually nullptr).
+//
+// Notes:
+//    On ARM64 this may remove the JTRUE node and transform its associated
+//    relop into a JCMP node.
+//
+GenTree* Lowering::LowerJTrue(GenTreeOp* jtrue)
+{
+#ifdef _TARGET_ARM64_
+    GenTree* relop    = jtrue->gtGetOp1();
+    GenTree* relopOp2 = relop->gtOp.gtGetOp2();
+
+    if ((relop->gtNext == jtrue) && relopOp2->IsCnsIntOrI())
+    {
+        bool     useJCMP = false;
+        unsigned flags   = 0;
+
+        if (relop->OperIs(GT_EQ, GT_NE) && relopOp2->IsIntegralConst(0))
+        {
+            // Codegen will use cbz or cbnz in codegen which do not affect the flag register
+            flags   = relop->OperIs(GT_EQ) ? GTF_JCMP_EQ : 0;
+            useJCMP = true;
+        }
+        else if (relop->OperIs(GT_TEST_EQ, GT_TEST_NE) && isPow2(relopOp2->AsIntCon()->IconValue()))
+        {
+            // Codegen will use tbz or tbnz in codegen which do not affect the flag register
+            flags   = GTF_JCMP_TST | (relop->OperIs(GT_TEST_EQ) ? GTF_JCMP_EQ : 0);
+            useJCMP = true;
+        }
+
+        if (useJCMP)
+        {
+            relop->SetOper(GT_JCMP);
+            relop->gtFlags &= ~(GTF_JCMP_TST | GTF_JCMP_EQ);
+            relop->gtFlags |= flags;
+            relop->gtLsraInfo.isNoRegCompare = true;
+
+            relopOp2->SetContained();
+
+            BlockRange().Remove(jtrue);
+
+            assert(relop->gtNext == nullptr);
+            return nullptr;
+        }
+    }
+#endif // _TARGET_ARM64_
+
+    ContainCheckJTrue(jtrue);
+
+    assert(jtrue->gtNext == nullptr);
+    return nullptr;
+}
+
 // Lower "jmp <method>" tail call to insert PInvoke method epilog if required.
 void Lowering::LowerJmpMethod(GenTree* jmp)
 {

src/jit/lower.h

@@ -138,6 +138,7 @@ class Lowering : public Phase
     // ------------------------------
     void LowerCall(GenTree* call);
     GenTree* LowerCompare(GenTree* tree);
+    GenTree* LowerJTrue(GenTreeOp* jtrue);
     void LowerJmpMethod(GenTree* jmp);
     void LowerRet(GenTree* ret);
     GenTree* LowerDelegateInvoke(GenTreeCall* call);

src/jit/lowerarmarch.cpp

@@ -694,46 +694,7 @@ void Lowering::ContainCheckCast(GenTreeCast* node)
 //
 void Lowering::ContainCheckCompare(GenTreeOp* cmp)
 {
-    if (CheckImmedAndMakeContained(cmp, cmp->gtOp2))
-    {
-#ifdef _TARGET_ARM64_
-        GenTreePtr op1 = cmp->gtOp.gtOp1;
-        GenTreePtr op2 = cmp->gtOp.gtOp2;
-
-        if (!varTypeIsFloating(cmp) && op2->IsCnsIntOrI())
-        {
-            LIR::Use cmpUse;
-            bool     useJCMP = false;
-            uint64_t flags   = 0;
-
-            if (cmp->OperIs(GT_EQ, GT_NE) && op2->IsIntegralConst(0) && BlockRange().TryGetUse(cmp, &cmpUse) &&
-                cmpUse.User()->OperIs(GT_JTRUE))
-            {
-                // Codegen will use cbz or cbnz in codegen which do not affect the flag register
-                flags   = cmp->OperIs(GT_EQ) ? GTF_JCMP_EQ : 0;
-                useJCMP = true;
-            }
-            else if (cmp->OperIs(GT_TEST_EQ, GT_TEST_NE) && isPow2(op2->gtIntCon.IconValue()) &&
-                     BlockRange().TryGetUse(cmp, &cmpUse) && cmpUse.User()->OperIs(GT_JTRUE))
-            {
-                // Codegen will use tbz or tbnz in codegen which do not affect the flag register
-                flags   = GTF_JCMP_TST | (cmp->OperIs(GT_TEST_EQ) ? GTF_JCMP_EQ : 0);
-                useJCMP = true;
-            }
-
-            if (useJCMP)
-            {
-                cmp->gtLsraInfo.isNoRegCompare = true;
-                cmp->SetOper(GT_JCMP);
-
-                cmp->gtFlags &= ~(GTF_JCMP_TST | GTF_JCMP_EQ);
-                cmp->gtFlags |= flags;
-
-                BlockRange().Remove(cmpUse.User());
-            }
-        }
-#endif // _TARGET_ARM64_
-    }
+    CheckImmedAndMakeContained(cmp, cmp->gtOp2);
 }
 
 //------------------------------------------------------------------------