JIT: rework tail call IR validity checks

Adapt `fgValidateIRForTailCall` to use as a utility to verify that the
JIT has not added IR that would make a tail call invalid. Currently all
our cases pass this check so no tail calls are invalidated.

Remove `fgCheckStmtAfterTailCall` as it did less thorough and less
correct checking.

Contributes to dotnet#93246.

Author: AndyAyersMS

src/coreclr/jit/compiler.h

@@ -5979,7 +5979,6 @@ class Compiler
     bool fgCallArgWillPointIntoLocalFrame(GenTreeCall* call, CallArg& arg);
 
 #endif
-    bool     fgCheckStmtAfterTailCall();
     GenTree* fgMorphTailCallViaHelpers(GenTreeCall* call, CORINFO_TAILCALL_HELPERS& help);
     bool fgCanTailCallViaJitHelper(GenTreeCall* call);
     void fgMorphTailCallViaJitHelper(GenTreeCall* call);
@@ -5999,7 +5998,7 @@ class Compiler
     GenTree* getTokenHandleTree(CORINFO_RESOLVED_TOKEN* pResolvedToken, bool parent);
 
     GenTree* fgMorphPotentialTailCall(GenTreeCall* call);
-    void fgValidateIRForTailCall(GenTreeCall* call);
+    bool fgValidateIRForTailCall(GenTreeCall* call);
     GenTree* fgGetStubAddrArg(GenTreeCall* call);
     unsigned fgGetArgParameterLclNum(GenTreeCall* call, CallArg* arg);
     void fgMorphRecursiveFastTailCallIntoLoop(BasicBlock* block, GenTreeCall* recursiveTailCall);

src/coreclr/jit/morph.cpp

@@ -5894,7 +5894,7 @@ GenTree* Compiler::fgMorphPotentialTailCall(GenTreeCall* call)
         }
     }
 
-    if (!fgCheckStmtAfterTailCall())
+    if (!fgValidateIRForTailCall(call))
     {
         failTailCall("Unexpected statements after the tail call");
         return nullptr;
@@ -6088,8 +6088,6 @@ GenTree* Compiler::fgMorphPotentialTailCall(GenTreeCall* call)
 #endif
     }
 
-    fgValidateIRForTailCall(call);
-
     // If this block has a flow successor, make suitable updates.
     //
     BasicBlock* nextBlock = compCurBB->GetUniqueSucc();
@@ -6340,21 +6338,25 @@ GenTree* Compiler::fgMorphPotentialTailCall(GenTreeCall* call)
 // fgValidateIRForTailCall:
 //     Validate that the IR looks ok to perform a tailcall.
 //
+// Returns:
+//    false if IR after the tail call has non-negligble effects,
+//    in which case the tail call should be abandoned.
+//
 // Arguments:
 //     call - The call that we are dispatching as a tailcall.
 //
 // Notes:
 //   This function needs to handle somewhat complex IR that appears after
 //   tailcall candidates due to inlining.
 //
-void Compiler::fgValidateIRForTailCall(GenTreeCall* call)
+bool Compiler::fgValidateIRForTailCall(GenTreeCall* call)
 {
-#ifdef DEBUG
     class TailCallIRValidatorVisitor final : public GenTreeVisitor<TailCallIRValidatorVisitor>
     {
         GenTreeCall* m_tailcall;
         unsigned     m_lclNum;
         bool         m_active;
+        bool         m_isValid;
 
     public:
         enum
@@ -6363,14 +6365,23 @@ void Compiler::fgValidateIRForTailCall(GenTreeCall* call)
             UseExecutionOrder = true,
         };
 
+        bool IsValid() const
+        {
+            return m_isValid;
+        }
+        void SetIsNotValid()
+        {
+            m_isValid = false;
+        }
+
         TailCallIRValidatorVisitor(Compiler* comp, GenTreeCall* tailcall)
-            : GenTreeVisitor(comp), m_tailcall(tailcall), m_lclNum(BAD_VAR_NUM), m_active(false)
+            : GenTreeVisitor(comp), m_tailcall(tailcall), m_lclNum(BAD_VAR_NUM), m_active(false), m_isValid(true)
         {
         }
 
         fgWalkResult PostOrderVisit(GenTree** use, GenTree* user)
         {
-            GenTree* tree = *use;
+            GenTree* const tree = *use;
 
             // Wait until we get to the actual call...
             if (!m_active)
@@ -6385,8 +6396,13 @@ void Compiler::fgValidateIRForTailCall(GenTreeCall* call)
 
             if (tree->OperIs(GT_RETURN))
             {
-                assert((tree->TypeIs(TYP_VOID) || ValidateUse(tree->gtGetOp1())) &&
-                       "Expected return to be result of tailcall");
+                if (!tree->TypeIs(TYP_VOID) && !ValidateUse(tree->gtGetOp1()))
+                {
+                    JITDUMP("Tail call validation failed: expected return to be result of tailcall\n");
+                    DISPTREE(tree);
+                    m_isValid = false;
+                }
+
                 return WALK_ABORT;
             }
 
@@ -6409,17 +6425,32 @@ void Compiler::fgValidateIRForTailCall(GenTreeCall* call)
             //
             else if (tree->OperIs(GT_STORE_LCL_VAR))
             {
-                assert(ValidateUse(tree->AsLclVar()->Data()) && "Expected value of store to be result of tailcall");
+                if (!ValidateUse(tree->AsLclVar()->Data()))
+                {
+                    JITDUMP("Tail call validation failed: [%06u] is not use of V%02u\n", m_compiler->dspTreeID(tree),
+                            m_lclNum);
+                    m_isValid = false;
+                    return WALK_ABORT;
+                }
+
                 m_lclNum = tree->AsLclVar()->GetLclNum();
             }
             else if (tree->OperIs(GT_LCL_VAR))
             {
-                assert(ValidateUse(tree) && "Expected use of local to be tailcall value");
+                if (!ValidateUse(tree))
+                {
+                    JITDUMP("Tail call validation failed [%06u] is not use of V%02u\n", m_compiler->dspTreeID(tree),
+                            m_lclNum);
+                    m_isValid = false;
+                    return WALK_ABORT;
+                }
             }
             else
             {
+                JITDUMP("Tail call validation failed: unexpected tree\n");
                 DISPTREE(tree);
-                assert(!"Unexpected tree op after call marked as tailcall");
+                m_isValid = false;
+                return WALK_ABORT;
             }
 
             return WALK_CONTINUE;
@@ -6452,24 +6483,39 @@ void Compiler::fgValidateIRForTailCall(GenTreeCall* call)
         }
     };
 
+    JITDUMP("Validating IR for tail call candidate [%06u] in " FMT_STMT "\n", dspTreeID(call), compCurStmt->GetID());
+
     TailCallIRValidatorVisitor visitor(this, call);
-    for (Statement* stmt = compCurStmt; stmt != nullptr; stmt = stmt->GetNextStmt())
+    for (Statement* stmt = compCurStmt; visitor.IsValid() && (stmt != nullptr); stmt = stmt->GetNextStmt())
     {
         visitor.WalkTree(stmt->GetRootNodePointer(), nullptr);
     }
 
     BasicBlock* bb = compCurBB;
-    while (!bb->KindIs(BBJ_RETURN))
+    while (!bb->KindIs(BBJ_RETURN) && visitor.IsValid())
     {
-        bb = bb->GetUniqueSucc();
-        assert((bb != nullptr) && "Expected straight flow after tailcall");
+        BasicBlock* const succBB = bb->GetUniqueSucc();
 
-        for (Statement* stmt : bb->Statements())
+        if (succBB == nullptr)
+        {
+            JITDUMP("Tail call validation failed: " FMT_BB " does not have linear flow\n", bb->bbNum);
+            visitor.SetIsNotValid();
+            break;
+        }
+
+        for (Statement* stmt : succBB->Statements())
         {
             visitor.WalkTree(stmt->GetRootNodePointer(), nullptr);
+            if (!visitor.IsValid())
+            {
+                break;
+            }
         }
+
+        bb = succBB;
     }
-#endif
+
+    return visitor.IsValid();
 }
 
 //------------------------------------------------------------------------
@@ -15367,108 +15413,6 @@ void Compiler::fgMarkDemotedImplicitByRefArgs()
 #endif // FEATURE_IMPLICIT_BYREFS
 }
 
-//------------------------------------------------------------------------
-// fgCheckStmtAfterTailCall: check that statements after the tail call stmt
-// candidate are in one of expected forms, that are desctibed below.
-//
-// Return Value:
-//    'true' if stmts are in the expected form, else 'false'.
-//
-bool Compiler::fgCheckStmtAfterTailCall()
-{
-
-    // For void calls, we would have created a GT_CALL in the stmt list.
-    // For non-void calls, we would have created a GT_RETURN(GT_CAST(GT_CALL)).
-    // For calls returning structs, we would have a void call, followed by a void return.
-    // For debuggable code, it would be an assignment of the call to a temp
-    // We want to get rid of any of this extra trees, and just leave
-    // the call.
-    Statement* callStmt = fgMorphStmt;
-
-    Statement* nextMorphStmt = callStmt->GetNextStmt();
-
-    // Check that the rest stmts in the block are in one of the following pattern:
-    //  1) ret(void)
-    //  2) ret(cast*(callResultLclVar))
-    //  3) lclVar = callResultLclVar, the actual ret(lclVar) in another block
-    //  4) nop
-    if (nextMorphStmt != nullptr)
-    {
-        GenTree* callExpr = callStmt->GetRootNode();
-        if (!callExpr->OperIs(GT_STORE_LCL_VAR))
-        {
-            // The next stmt can be GT_RETURN(TYP_VOID) or GT_RETURN(lclVar),
-            // where lclVar was return buffer in the call for structs or simd.
-            Statement* retStmt = nextMorphStmt;
-            GenTree*   retExpr = retStmt->GetRootNode();
-            noway_assert(retExpr->gtOper == GT_RETURN);
-
-            nextMorphStmt = retStmt->GetNextStmt();
-        }
-        else
-        {
-            noway_assert(callExpr->OperIs(GT_STORE_LCL_VAR));
-            unsigned callResultLclNumber = callExpr->AsLclVar()->GetLclNum();
-
-#if FEATURE_TAILCALL_OPT_SHARED_RETURN
-
-            // We can have a chain of assignments from the call result to
-            // various inline return spill temps. These are ok as long
-            // as the last one ultimately provides the return value or is ignored.
-            //
-            // And if we're returning a small type we may see a cast
-            // on the source side.
-            while ((nextMorphStmt != nullptr) && (nextMorphStmt->GetRootNode()->OperIs(GT_STORE_LCL_VAR, GT_NOP)))
-            {
-                if (nextMorphStmt->GetRootNode()->OperIs(GT_NOP))
-                {
-                    nextMorphStmt = nextMorphStmt->GetNextStmt();
-                    continue;
-                }
-                Statement* moveStmt = nextMorphStmt;
-                GenTree*   moveExpr = nextMorphStmt->GetRootNode();
-
-                // Tunnel through any casts on the source side.
-                GenTree* moveSource = moveExpr->AsLclVar()->Data();
-                while (moveSource->OperIs(GT_CAST))
-                {
-                    noway_assert(!moveSource->gtOverflow());
-                    moveSource = moveSource->gtGetOp1();
-                }
-                noway_assert(moveSource->OperIsLocal());
-
-                // Verify we're just passing the value from one local to another
-                // along the chain.
-                const unsigned srcLclNum = moveSource->AsLclVarCommon()->GetLclNum();
-                noway_assert(srcLclNum == callResultLclNumber);
-                const unsigned dstLclNum = moveExpr->AsLclVar()->GetLclNum();
-                callResultLclNumber      = dstLclNum;
-
-                nextMorphStmt = moveStmt->GetNextStmt();
-            }
-            if (nextMorphStmt != nullptr)
-#endif
-            {
-                Statement* retStmt = nextMorphStmt;
-                GenTree*   retExpr = nextMorphStmt->GetRootNode();
-                noway_assert(retExpr->gtOper == GT_RETURN);
-
-                GenTree* treeWithLcl = retExpr->gtGetOp1();
-                while (treeWithLcl->gtOper == GT_CAST)
-                {
-                    noway_assert(!treeWithLcl->gtOverflow());
-                    treeWithLcl = treeWithLcl->gtGetOp1();
-                }
-
-                noway_assert(callResultLclNumber == treeWithLcl->AsLclVarCommon()->GetLclNum());
-
-                nextMorphStmt = retStmt->GetNextStmt();
-            }
-        }
-    }
-    return nextMorphStmt == nullptr;
-}
-
 //------------------------------------------------------------------------
 // fgCanTailCallViaJitHelper: check whether we can use the faster tailcall
 // JIT helper on x86.