Add unroller support to CloneBlockState

This helper is used by the loop cloner to copy a block's attributes and
statements.  Expand it so it can also be used by the loop unroller:
 - Accept var/val parameters to pass through to `gtCloneExpr` to
   perform in-place substitutions
 - When presented with an expression that `gtCloneExpr` can't clone,
   return false rather than dereferencing null

3 files changed, 39 insertions, 7 deletions

diff --git a/src/jit/block.cpp b/src/jit/block.cpp
index cf2a107479..ed8a103f53 100644
--- a/src/jit/block.cpp
+++ b/src/jit/block.cpp
@@ -568,7 +568,25 @@ void* BasicBlock::HeapPhiArg::operator new(size_t sz, Compiler* comp)
     return comp->compGetMem(sz, CMK_HeapPhiArg);
 }
 
-void BasicBlock::CloneBlockState(Compiler* compiler, BasicBlock* to, const BasicBlock* from)
+//------------------------------------------------------------------------
+// CloneBlockState: Try to populate `to` block with a copy of `from` block's statements, replacing
+//                  uses of local `varNum` with IntCns `varVal`.
+//
+// Arguments:
+//    compiler - Jit compiler instance
+//    to - New/empty block to copy statements into
+//    from - Block to copy statements from
+//    varNum - lclVar uses with lclNum `varNum` will be replaced; can be ~0 to indicate no replacement.
+//    varVal - If replacing uses of `varNum`, replace them with int constants with value `varVal`.
+//
+// Return Value:
+//    Cloning may fail because this routine uses `gtCloneExpr` for cloning and it can't handle all
+//    IR nodes.  If cloning of any statement fails, `false` will be returned and block `to` may be
+//    partially populated.  If cloning of all statements succeeds, `true` will be returned and
+//    block `to` will be fully populated.
+
+bool BasicBlock::CloneBlockState(
+    Compiler* compiler, BasicBlock* to, const BasicBlock* from, unsigned varNum, int varVal)
 {
     assert(to->bbTreeList == nullptr);
 
@@ -595,9 +613,17 @@ void BasicBlock::CloneBlockState(Compiler* compiler, BasicBlock* to, const Basic
 
     for (GenTreePtr fromStmt = from->bbTreeList; fromStmt != nullptr; fromStmt = fromStmt->gtNext)
     {
-        compiler->fgInsertStmtAtEnd(to,
-                                    compiler->fgNewStmtFromTree(compiler->gtCloneExpr(fromStmt->gtStmt.gtStmtExpr)));
+        auto newExpr = compiler->gtCloneExpr(fromStmt->gtStmt.gtStmtExpr, 0, varNum, varVal);
+        if (!newExpr)
+        {
+            // gtCloneExpr doesn't handle all opcodes, so may fail to clone a statement.
+            // When that happens, it returns nullptr; abandon the rest of this block and
+            // return `false` to the caller to indicate that cloning was unsuccessful.
+            return false;
+        }
+        compiler->fgInsertStmtAtEnd(to, compiler->fgNewStmtFromTree(newExpr));
     }
+    return true;
 }
 
 // LIR helpers
diff --git a/src/jit/block.h b/src/jit/block.h
index 280b60298e..20d4952ad6 100644
--- a/src/jit/block.h
+++ b/src/jit/block.h
@@ -1088,9 +1088,11 @@ public:
         return AllSuccs(comp, this);
     }
 
-    // Clone block state and statements from 'from' block to 'to' block.
-    // Assumes that "to" is an empty block.
-    static void CloneBlockState(Compiler* compiler, BasicBlock* to, const BasicBlock* from);
+    // Try to clone block state and statements from `from` block to `to` block (which must be new/empty),
+    // optionally replacing uses of local `varNum` with IntCns `varVal`.  Return true if all statements
+    // in the block are cloned successfully, false (with partially-populated `to` block) if one fails.
+    static bool CloneBlockState(
+        Compiler* compiler, BasicBlock* to, const BasicBlock* from, unsigned varNum = (unsigned)-1, int varVal = 0);
 
     void MakeLIR(GenTree* firstNode, GenTree* lastNode);
     bool IsLIR();
diff --git a/src/jit/optimizer.cpp b/src/jit/optimizer.cpp
index dc6c9976b0..fdcae0eeee 100644
--- a/src/jit/optimizer.cpp
+++ b/src/jit/optimizer.cpp
@@ -4648,7 +4648,11 @@ void Compiler::optCloneLoop(unsigned loopInd, LoopCloneContext* context)
         BasicBlock* newBlk = fgNewBBafter(blk->bbJumpKind, newPred,
                                           /*extendRegion*/ true);
 
-        BasicBlock::CloneBlockState(this, newBlk, blk);
+        // Call CloneBlockState to make a copy of the block's statements (and attributes), and assert that it
+        // has a return value indicating success, because optCanOptimizeByLoopCloningVisitor has already
+        // checked them to guarantee they are clonable.
+        bool cloneOk = BasicBlock::CloneBlockState(this, newBlk, blk);
+        noway_assert(cloneOk);
         // TODO-Cleanup: The above clones the bbNatLoopNum, which is incorrect.  Eventually, we should probably insert
         // the cloned loop in the loop table.  For now, however, we'll just make these blocks be part of the surrounding
         // loop, if one exists -- the parent of the loop we're cloning.