Refactored optRemoveRangeCheck to handle top-level standalone range check nodes

Author: SingleAccretion

src/coreclr/jit/compiler.h

@@ -6007,7 +6007,9 @@ class Compiler
 public:
     void optInit();
 
-    void optRemoveRangeCheck(GenTree* tree, Statement* stmt);
+    GenTree* Compiler::optRemoveRangeCheck(GenTreeBoundsChk* check, GenTree* comma, Statement* stmt);
+    GenTree* Compiler::optRemoveStandaloneRangeCheck(GenTreeBoundsChk* check, Statement* stmt);
+    void Compiler::optRemoveCommaBasedRangeCheck(GenTree* comma, Statement* stmt);
     bool optIsRangeCheckRemovable(GenTree* tree);
 
 protected:
@@ -7286,7 +7288,7 @@ class Compiler
     GenTree* optAssertionProp_Call(ASSERT_VALARG_TP assertions, GenTreeCall* call, Statement* stmt);
     GenTree* optAssertionProp_RelOp(ASSERT_VALARG_TP assertions, GenTree* tree, Statement* stmt);
     GenTree* optAssertionProp_Comma(ASSERT_VALARG_TP assertions, GenTree* tree, Statement* stmt);
-    GenTree* optAssertionProp_BndsChk(ASSERT_VALARG_TP assertions, GenTree* tree);
+    GenTree* optAssertionProp_BndsChk(ASSERT_VALARG_TP assertions, GenTree* tree, Statement* stmt);
     GenTree* optAssertionPropGlobal_RelOp(ASSERT_VALARG_TP assertions, GenTree* tree, Statement* stmt);
     GenTree* optAssertionPropLocal_RelOp(ASSERT_VALARG_TP assertions, GenTree* tree, Statement* stmt);
     GenTree* optAssertionProp_Update(GenTree* newTree, GenTree* tree, Statement* stmt);

src/coreclr/jit/optimizer.cpp

@@ -5068,7 +5068,7 @@ void Compiler::optPerformStaticOptimizations(unsigned loopNum, LoopCloneContext*
             {
                 LcJaggedArrayOptInfo* arrIndexInfo = optInfo->AsLcJaggedArrayOptInfo();
                 compCurBB                          = arrIndexInfo->arrIndex.useBlock;
-                optRemoveRangeCheck(arrIndexInfo->arrIndex.bndsChks[arrIndexInfo->dim], arrIndexInfo->stmt);
+                optRemoveCommaBasedRangeCheck(arrIndexInfo->arrIndex.bndsChks[arrIndexInfo->dim], arrIndexInfo->stmt);
                 DBEXEC(dynamicPath, optDebugLogLoopCloning(arrIndexInfo->arrIndex.useBlock, arrIndexInfo->stmt));
             }
             break;
@@ -8211,24 +8211,34 @@ void Compiler::AddModifiedElemTypeAllContainingLoops(unsigned lnum, CORINFO_CLAS
 }
 
 //------------------------------------------------------------------------------
-// optRemoveRangeCheck : Given an array index node, mark it as not needing a range check.
+// optRemoveRangeCheck : Given an indexing node, mark it as not needing a range check.
 //
 // Arguments:
-//    tree   -  Range check tree
-//    stmt   -  Statement the tree belongs to
-
-void Compiler::optRemoveRangeCheck(GenTree* tree, Statement* stmt)
+//    check  -  Range check tree, the raw CHECK node (ARRAY, SIMD or HWINTRINSIC).
+//    comma  -  GT_COMMA to which the "check" belongs, "nullptr" if the check is a standalone one.
+//    stmt   -  Statement the indexing nodes belongs to.
+//
+// Return Value:
+//    Rewritten "check" - no-op if it has no side effects or the tree that contains them.
+//
+// Assumptions:
+//    This method is capable of removing checks of two kinds: COMMA-based and standalone top-level ones.
+//    In case of a COMMA-based check, "check" must be a non-null first operand of a non-null COMMA.
+//    If case of a standalone check, "comma" must be null and "check" - "stmt"'s root.
+//
+GenTree* Compiler::optRemoveRangeCheck(GenTreeBoundsChk* check, GenTree* comma, Statement* stmt)
 {
 #if !REARRANGE_ADDS
     noway_assert(!"can't remove range checks without REARRANGE_ADDS right now");
 #endif
 
-    noway_assert(tree->gtOper == GT_COMMA);
-
-    GenTree* bndsChkTree = tree->AsOp()->gtOp1;
-
-    noway_assert(bndsChkTree->OperIsBoundsCheck());
+    noway_assert(stmt != nullptr);
+    noway_assert((comma != nullptr && comma->OperIs(GT_COMMA) && comma->gtGetOp1() == check) ||
+                 (check != nullptr && check->OperIsBoundsCheck() && comma == nullptr));
+    noway_assert(check->OperIsBoundsCheck());
 
+    GenTree* tree = comma != nullptr ? comma : check;
+    
 #ifdef DEBUG
     if (verbose)
     {
@@ -8239,19 +8249,40 @@ void Compiler::optRemoveRangeCheck(GenTree* tree, Statement* stmt)
 
     // Extract side effects
     GenTree* sideEffList = nullptr;
-    gtExtractSideEffList(bndsChkTree, &sideEffList, GTF_ASG);
+    gtExtractSideEffList(check, &sideEffList, GTF_ASG);
+
+    if (sideEffList != nullptr)
+    {
+        // We've got some side effects.
+        if (tree->OperIs(GT_COMMA))
+        {
+            // Make the comma handle them.
+            tree->AsOp()->gtOp1 = sideEffList;
+        }
+        else
+        {
+            // Make the statement execute them instead of the check.
+            stmt->SetRootNode(sideEffList);
+            tree = sideEffList;
+        }
+    }
+    else
+    {
+        check->gtBashToNOP();
+    }
 
-    // Just replace the bndsChk with a NOP as an operand to the GT_COMMA, if there are no side effects.
-    tree->AsOp()->gtOp1 = (sideEffList != nullptr) ? sideEffList : gtNewNothingNode();
-    // TODO-CQ: We should also remove the GT_COMMA, but in any case we can no longer CSE the GT_COMMA.
-    tree->gtFlags |= GTF_DONT_CSE;
+    if (tree->OperIs(GT_COMMA))
+    {
+        // TODO-CQ: We should also remove the GT_COMMA, but in any case we can no longer CSE the GT_COMMA.
+        tree->gtFlags |= GTF_DONT_CSE;
+    }
 
     gtUpdateSideEffects(stmt, tree);
 
-    /* Recalculate the GetCostSz(), etc... */
+    // Recalculate the GetCostSz(), etc...
     gtSetStmtInfo(stmt);
 
-    /* Re-thread the nodes if necessary */
+    // Re-thread the nodes if necessary
     if (fgStmtListThreaded)
     {
         fgSetStmtSeq(stmt);
@@ -8264,6 +8295,32 @@ void Compiler::optRemoveRangeCheck(GenTree* tree, Statement* stmt)
         gtDispTree(tree);
     }
 #endif
+
+    return check;
+}
+
+//------------------------------------------------------------------------------
+// optRemoveStandaloneRangeCheck : A thin wrapper over optRemoveRangeCheck that removes standalone checks.
+//
+GenTree* Compiler::optRemoveStandaloneRangeCheck(GenTreeBoundsChk* check, Statement* stmt)
+{
+    assert(check != nullptr);
+    assert(stmt != nullptr);
+    assert(check == stmt->GetRootNode());
+
+    return optRemoveRangeCheck(check, nullptr, stmt);
+}
+
+//------------------------------------------------------------------------------
+// optRemoveCommaBasedRangeCheck : A thin wrapper over optRemoveRangeCheck that removes COMMA-based checks.
+//
+void Compiler::optRemoveCommaBasedRangeCheck(GenTree* comma, Statement* stmt)
+{
+    assert(comma != nullptr);
+    assert(stmt != nullptr);
+    assert(comma->gtGetOp1()->OperIsBoundsCheck());
+
+    optRemoveRangeCheck(comma->gtGetOp1()->AsBoundsChk(), comma, stmt);
 }
 
 /*****************************************************************************