From 2220c8366ed63ebf404ee48855bf69fab71fa4ab Mon Sep 17 00:00:00 2001
From: Andy Ayers <andya@microsoft.com>
Date: Fri, 24 Feb 2023 09:07:21 -0800
Subject: [PATCH] JIT: fix finally target flags after some flow graph updates

In particular those done by profile instrumentation.

Fixes #82275.
---
 src/coreclr/jit/compiler.h    |  2 ++
 src/coreclr/jit/fgbasic.cpp   |  8 ++++++++
 src/coreclr/jit/fgehopt.cpp   | 26 ++++++++++++++++++++++++++
 src/coreclr/jit/fgprofile.cpp |  4 ++++
 4 files changed, 40 insertions(+)

diff --git a/src/coreclr/jit/compiler.h b/src/coreclr/jit/compiler.h
index b7eb5c20d69df6..5df7af79c4b473 100644
--- a/src/coreclr/jit/compiler.h
+++ b/src/coreclr/jit/compiler.h
@@ -4633,6 +4633,8 @@ class Compiler
 
     void fgAddFinallyTargetFlags();
 
+    void fgFixFinallyTargetFlags(BasicBlock* pred, BasicBlock* succ, BasicBlock* newBlock);
+
 #endif // defined(FEATURE_EH_FUNCLETS) && defined(TARGET_ARM)
     PhaseStatus fgTailMergeThrows();
     void fgTailMergeThrowsFallThroughHelper(BasicBlock* predBlock,
diff --git a/src/coreclr/jit/fgbasic.cpp b/src/coreclr/jit/fgbasic.cpp
index df12d1eb12e517..581a5aa620bef8 100644
--- a/src/coreclr/jit/fgbasic.cpp
+++ b/src/coreclr/jit/fgbasic.cpp
@@ -584,6 +584,10 @@ void Compiler::fgReplaceJumpTarget(BasicBlock* block, BasicBlock* newTarget, Bas
             unreached();
             break;
     }
+
+#if defined(FEATURE_EH_FUNCLETS) && defined(TARGET_ARM)
+    fgFixFinallyTargetFlags(block, oldTarget, newTarget);
+#endif
 }
 
 //------------------------------------------------------------------------
@@ -4779,6 +4783,10 @@ BasicBlock* Compiler::fgSplitEdge(BasicBlock* curr, BasicBlock* succ)
         fgAddRefPred(newBlock, curr);
     }
 
+#if defined(FEATURE_EH_FUNCLETS) && defined(TARGET_ARM)
+    fgFixFinallyTargetFlags(curr, succ, newBlock);
+#endif
+
     // This isn't accurate, but it is complex to compute a reasonable number so just assume that we take the
     // branch 50% of the time.
     //
diff --git a/src/coreclr/jit/fgehopt.cpp b/src/coreclr/jit/fgehopt.cpp
index 4e16bb875131a3..5e1dd111a71847 100644
--- a/src/coreclr/jit/fgehopt.cpp
+++ b/src/coreclr/jit/fgehopt.cpp
@@ -1620,6 +1620,32 @@ void Compiler::fgAddFinallyTargetFlags()
     }
 }
 
+//------------------------------------------------------------------------
+// fgFixFinallyTargetFlags: Update BBF_FINALLY_TARGET bits after redirecting flow
+//
+// Arguments:
+//   pred - source of flow
+//   succ - original target of flow from pred
+//   newSucc - new target of flow from pred
+//
+void Compiler::fgFixFinallyTargetFlags(BasicBlock* pred, BasicBlock* succ, BasicBlock* newSucc)
+{
+    if (pred->isBBCallAlwaysPairTail())
+    {
+        assert(succ->bbFlags & BBF_FINALLY_TARGET);
+        newSucc->bbFlags |= BBF_FINALLY_TARGET;
+        succ->bbFlags &= ~BBF_FINALLY_TARGET;
+
+        for (BasicBlock* const pred : succ->PredBlocks())
+        {
+            if (pred->isBBCallAlwaysPairTail())
+            {
+                succ->bbFlags |= BBF_FINALLY_TARGET;
+                break;
+            }
+        }
+    }
+}
 #endif // defined(FEATURE_EH_FUNCLETS) && defined(TARGET_ARM)
 
 //------------------------------------------------------------------------
diff --git a/src/coreclr/jit/fgprofile.cpp b/src/coreclr/jit/fgprofile.cpp
index 54c1a3b26848b5..bdaf59d5e6823a 100644
--- a/src/coreclr/jit/fgprofile.cpp
+++ b/src/coreclr/jit/fgprofile.cpp
@@ -522,6 +522,10 @@ void BlockCountInstrumentor::RelocateProbes()
                 {
                     m_comp->fgRemoveRefPred(pred, block);
                     m_comp->fgAddRefPred(intermediary, block);
+
+#if defined(FEATURE_EH_FUNCLETS) && defined(TARGET_ARM)
+                    m_comp->fgFixFinallyTargetFlags(pred, block, intermediary);
+#endif
                 }
             }
         }