From eafcb9ada30e711c09aa0b817b7b5b0c739eb0ee Mon Sep 17 00:00:00 2001
From: Bruce Forstall <brucefo@microsoft.com>
Date: Thu, 18 Jan 2024 09:22:27 -0800
Subject: [PATCH] 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       |  2 ++
 src/coreclr/jit/fgbasic.cpp      | 20 ++++++++++++++++++++
 src/coreclr/jit/fgdiagnostic.cpp |  7 +++++++
 3 files changed, 29 insertions(+)

diff --git a/src/coreclr/jit/compiler.h b/src/coreclr/jit/compiler.h
index 2d15aab8a96560..80e0dc70809c78 100644
--- a/src/coreclr/jit/compiler.h
+++ b/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);
diff --git a/src/coreclr/jit/fgbasic.cpp b/src/coreclr/jit/fgbasic.cpp
index d88ed59a8f1963..a7a7cc0b5d980d 100644
--- a/src/coreclr/jit/fgbasic.cpp
+++ b/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;
 }
 
diff --git a/src/coreclr/jit/fgdiagnostic.cpp b/src/coreclr/jit/fgdiagnostic.cpp
index 2fbf905d131acf..1b8ed455ae5a67 100644
--- a/src/coreclr/jit/fgdiagnostic.cpp
+++ b/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.