4 files changed, 45 insertions, 21 deletions

diff --git a/src/jit/assertionprop.cpp b/src/jit/assertionprop.cpp
index cd1b90ffa2..5b62e88303 100644
--- a/src/jit/assertionprop.cpp
+++ b/src/jit/assertionprop.cpp
@@ -3901,11 +3901,22 @@ GenTree* Compiler::optAssertionProp_Update(GenTree* newTree, GenTree* tree, GenT
         // locate our parent node and update it so that it points to newTree.
         if (newTree != tree)
         {
-            GenTree** link = gtFindLink(stmt, tree);
-            noway_assert(link != nullptr);
+            FindLinkData linkData = gtFindLink(stmt, tree);
+            GenTree**    useEdge  = linkData.result;
+            GenTree*     parent   = linkData.parent;
+            noway_assert(useEdge != nullptr);
 
-            // Replace the old operand with the newTree
-            *link = newTree;
+            if (parent != nullptr)
+            {
+                parent->ReplaceOperand(useEdge, newTree);
+            }
+            else
+            {
+                // If there's no parent, the tree being replaced is the root of the
+                // statement.
+                assert((stmt->gtStmtExpr == tree) && (&stmt->gtStmtExpr == useEdge));
+                stmt->gtStmtExpr = newTree;
+            }
 
             // We only need to ensure that the gtNext field is set as it is used to traverse
             // to the next node in the tree. We will re-morph this entire statement in
diff --git a/src/jit/compiler.h b/src/jit/compiler.h
index b091c929c9..fed7003faa 100644
--- a/src/jit/compiler.h
+++ b/src/jit/compiler.h
@@ -2994,7 +2994,14 @@ public:
     static fgWalkPreFn gtMarkColonCond;
     static fgWalkPreFn gtClearColonCond;
 
-    GenTree** gtFindLink(GenTreeStmt* stmt, GenTree* node);
+    struct FindLinkData
+    {
+        GenTree*  nodeToFind;
+        GenTree** result;
+        GenTree*  parent;
+    };
+
+    FindLinkData gtFindLink(GenTreeStmt* stmt, GenTree* node);
     bool gtHasCatchArg(GenTree* tree);
 
     typedef ArrayStack<GenTree*> GenTreeStack;
diff --git a/src/jit/gentree.cpp b/src/jit/gentree.cpp
index f0dfbc7ee5..df6b324a06 100644
--- a/src/jit/gentree.cpp
+++ b/src/jit/gentree.cpp
@@ -7931,23 +7931,32 @@ bool Compiler::gtCompareTree(GenTree* op1, GenTree* op2)
 
 GenTree* Compiler::gtGetThisArg(GenTreeCall* call)
 {
-    if (call->gtCallObjp != nullptr)
+    GenTree* thisArg = call->gtCallObjp;
+    if (thisArg != nullptr)
     {
-        if (call->gtCallObjp->gtOper != GT_NOP && call->gtCallObjp->gtOper != GT_ASG)
+        if (thisArg->OperIs(GT_NOP, GT_ASG) == false)
         {
-            if (!(call->gtCallObjp->gtFlags & GTF_LATE_ARG))
+            if ((thisArg->gtFlags & GTF_LATE_ARG) == 0)
             {
                 return call->gtCallObjp;
             }
         }
 
-        if (call->gtCallLateArgs)
+        if (call->gtCallLateArgs != nullptr)
         {
             unsigned       argNum          = 0;
             fgArgTabEntry* thisArgTabEntry = gtArgEntryByArgNum(call, argNum);
             GenTree*       result          = thisArgTabEntry->node;
 
+            // Assert if we used DEBUG_DESTROY_NODE.
+            assert(result->gtOper != GT_COUNT);
+
 #if !FEATURE_FIXED_OUT_ARGS && defined(DEBUG)
+            // Check that call->fgArgInfo used in gtArgEntryByArgNum was not
+            // left outdated by assertion propogation updates.
+            // There is no information about registers of late args for platforms
+            // with FEATURE_FIXED_OUT_ARGS that is why this debug check is under
+            // !FEATURE_FIXED_OUT_ARGS.
             regNumber thisReg  = REG_ARG_0;
             GenTree*  lateArgs = call->gtCallLateArgs;
             regList   list     = call->regArgList;
@@ -7966,6 +7975,7 @@ GenTree* Compiler::gtGetThisArg(GenTreeCall* call)
                 index++;
             }
 #endif // !FEATURE_FIXED_OUT_ARGS && defined(DEBUG)
+
             return result;
         }
     }
@@ -15156,42 +15166,37 @@ Compiler::fgWalkResult Compiler::gtClearColonCond(GenTree** pTree, fgWalkData* d
     return WALK_CONTINUE;
 }
 
-struct FindLinkData
-{
-    GenTree*  nodeToFind;
-    GenTree** result;
-};
-
 /*****************************************************************************
  *
  *  Callback used by the tree walker to implement fgFindLink()
  */
 static Compiler::fgWalkResult gtFindLinkCB(GenTree** pTree, Compiler::fgWalkData* cbData)
 {
-    FindLinkData* data = (FindLinkData*)cbData->pCallbackData;
+    Compiler::FindLinkData* data = (Compiler::FindLinkData*)cbData->pCallbackData;
     if (*pTree == data->nodeToFind)
     {
         data->result = pTree;
+        data->parent = cbData->parent;
         return Compiler::WALK_ABORT;
     }
 
     return Compiler::WALK_CONTINUE;
 }
 
-GenTree** Compiler::gtFindLink(GenTreeStmt* stmt, GenTree* node)
+Compiler::FindLinkData Compiler::gtFindLink(GenTreeStmt* stmt, GenTree* node)
 {
-    FindLinkData data = {node, nullptr};
+    FindLinkData data = {node, nullptr, nullptr};
 
     fgWalkResult result = fgWalkTreePre(&stmt->gtStmtExpr, gtFindLinkCB, &data);
 
     if (result == WALK_ABORT)
     {
         assert(data.nodeToFind == *data.result);
-        return data.result;
+        return data;
     }
     else
     {
-        return nullptr;
+        return {node, nullptr, nullptr};
     }
 }
 
diff --git a/src/jit/optcse.cpp b/src/jit/optcse.cpp
index 9dc421b3fd..a25e17c4fa 100644
--- a/src/jit/optcse.cpp
+++ b/src/jit/optcse.cpp
@@ -2364,7 +2364,8 @@ public:
             // Walk the statement 'stmt' and find the pointer
             // in the tree is pointing to 'exp'
             //
-            GenTree** link = m_pCompiler->gtFindLink(stmt, exp);
+            Compiler::FindLinkData linkData = m_pCompiler->gtFindLink(stmt, exp);
+            GenTree**              link     = linkData.result;
 
 #ifdef DEBUG
             if (link == nullptr)