Fix maintenance of genReturnBB pointer (#96935)

If the `genReturnBB` block is split, the pointer needs
to be updated.

Without this, we ended up with a situation where the
`genReturnBB` did not point to the return block, leading to
omitting the code to remove the PInvoke frame from the thread's
Frame list.

Fixes #96409

src/coreclr/jit/compiler.h

@@ -5348,6 +5348,8 @@ class Compiler
     IL_OFFSET fgFindBlockILOffset(BasicBlock* block);
     void fgFixEntryFlowForOSR();
 
+    void fgUpdateSingleReturnBlock(BasicBlock* block);
+
     BasicBlock* fgSplitBlockAtBeginning(BasicBlock* curr);
     BasicBlock* fgSplitBlockAtEnd(BasicBlock* curr);
     BasicBlock* fgSplitBlockAfterStatement(BasicBlock* curr, Statement* stmt);

src/coreclr/jit/fgbasic.cpp

@@ -4755,6 +4755,24 @@ IL_OFFSET Compiler::fgFindBlockILOffset(BasicBlock* block)
     return BAD_IL_OFFSET;
 }
 
+//------------------------------------------------------------------------------
+// fgUpdateSingleReturnBlock : A block has been split. If it was the single return
+// block, then update the single return block pointer.
+//
+// Arguments:
+//    block - The block that was split
+//
+void Compiler::fgUpdateSingleReturnBlock(BasicBlock* block)
+{
+    assert(block->KindIs(BBJ_ALWAYS));
+    if (genReturnBB == block)
+    {
+        assert(block->GetTarget()->KindIs(BBJ_RETURN));
+        JITDUMP("Updating genReturnBB from " FMT_BB " to " FMT_BB "\n", block->bbNum, block->GetTarget()->bbNum);
+        genReturnBB = block->GetTarget();
+    }
+}
+
 //------------------------------------------------------------------------------
 // fgSplitBlockAtEnd - split the given block into two blocks.
 //                   All code in the block stays in the original block.
@@ -4831,6 +4849,8 @@ BasicBlock* Compiler::fgSplitBlockAtEnd(BasicBlock* curr)
 
     fgAddRefPred(newBlock, curr);
 
+    fgUpdateSingleReturnBlock(curr);
+
     return newBlock;
 }
 

src/coreclr/jit/fgdiagnostic.cpp

@@ -2255,6 +2255,12 @@ void Compiler::fgTableDispBasicBlock(BasicBlock* block, int ibcColWidth /* = 0 *
         }
     }
 
+    // Indicate if it's the single return block
+    if (block == genReturnBB)
+    {
+        printf(" one-return");
+    }
+
     printf("\n");
 }
 
@@ -3186,6 +3192,7 @@ void Compiler::fgDebugCheckBBlist(bool checkBBNum /* = false */, bool checkBBRef
     if (genReturnBB != nullptr)
     {
         assert(genReturnBB->GetFirstLIRNode() != nullptr || genReturnBB->bbStmtList != nullptr);
+        assert(genReturnBB->KindIs(BBJ_RETURN));
     }
 
     // If this is an inlinee, we're done checking.