From dfa5c00474ba985b8db52dee29ca33bfae647f43 Mon Sep 17 00:00:00 2001
From: Jakob Botsch Nielsen <jakob.botsch.nielsen@gmail.com>
Date: Tue, 30 Jan 2024 15:19:57 +0100
Subject: [PATCH] JIT: Fix some cases using `BasicBlock::bbFallsThrough`

`bbFallsThrough` still returns true for `BBJ_COND`; we have a couple of
places using it as a "control flows from prev block" check, which is
wrong after #97488.
---
 src/coreclr/jit/block.cpp      | 12 +++++++-----
 src/coreclr/jit/morph.cpp      |  6 +++++-
 src/coreclr/jit/optimizer.cpp  | 11 +++++++----
 src/coreclr/jit/rangecheck.cpp |  2 +-
 4 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/src/coreclr/jit/block.cpp b/src/coreclr/jit/block.cpp
index b8356189bd13f..542edf3518737 100644
--- a/src/coreclr/jit/block.cpp
+++ b/src/coreclr/jit/block.cpp
@@ -1123,11 +1123,13 @@ Statement* BasicBlock::FirstNonPhiDefOrCatchArgStore() const
     return stmt;
 }
 
-/*****************************************************************************
- *
- *  Can a BasicBlock be inserted after this without altering the flowgraph
- */
-
+//------------------------------------------------------------------------
+// bbFallsThrough: Check if inserting a BasicBlock after this one will alter
+// the flowgraph.
+//
+// Returns:
+//    True if so.
+//
 bool BasicBlock::bbFallsThrough() const
 {
     switch (bbKind)
diff --git a/src/coreclr/jit/morph.cpp b/src/coreclr/jit/morph.cpp
index 1ce1b840de778..582f7a22569b9 100644
--- a/src/coreclr/jit/morph.cpp
+++ b/src/coreclr/jit/morph.cpp
@@ -13869,7 +13869,7 @@ void Compiler::fgMorphBlock(BasicBlock* block)
                     // Yes, pred assertions are available.
                     // If the pred is (a non-degenerate) BBJ_COND, fetch the appropriate out set.
                     //
-                    ASSERT_TP  assertionsOut     = pred->bbAssertionOut;
+                    ASSERT_TP  assertionsOut;
                     const bool useCondAssertions = pred->KindIs(BBJ_COND) && (pred->NumSucc() == 2);
 
                     if (useCondAssertions)
@@ -13886,6 +13886,10 @@ void Compiler::fgMorphBlock(BasicBlock* block)
                             assertionsOut = pred->bbAssertionOutIfFalse;
                         }
                     }
+                    else
+                    {
+                        assertionsOut = pred->bbAssertionOut;
+                    }
 
                     // If this is the first pred, copy (or share, when block is the only successor).
                     // If this is a subsequent pred, intersect.
diff --git a/src/coreclr/jit/optimizer.cpp b/src/coreclr/jit/optimizer.cpp
index 6c3e9e83e6e58..867cdc965870e 100644
--- a/src/coreclr/jit/optimizer.cpp
+++ b/src/coreclr/jit/optimizer.cpp
@@ -463,11 +463,14 @@ bool Compiler::optExtractInitTestIncr(
         // If we are rebuilding the loops, we would already have the pre-header block introduced
         // the first time, which might be empty if no hoisting has yet occurred. In this case, look a
         // little harder for the possible loop initialization statement.
-        if (initBlock->KindIs(BBJ_ALWAYS) && initBlock->TargetIs(header) && (initBlock->countOfInEdges() == 1) &&
-            !initBlock->IsFirst() && initBlock->Prev()->bbFallsThrough())
+        if (initBlock->KindIs(BBJ_ALWAYS) && initBlock->TargetIs(header))
         {
-            initBlock = initBlock->Prev();
-            phdrStmt  = initBlock->firstStmt();
+            BasicBlock* uniquePred = initBlock->GetUniquePred(this);
+            if (uniquePred != nullptr)
+            {
+                initBlock = uniquePred;
+                phdrStmt  = initBlock->firstStmt();
+            }
         }
     }
 
diff --git a/src/coreclr/jit/rangecheck.cpp b/src/coreclr/jit/rangecheck.cpp
index 3b03a96ca145a..6ac1c1c01646d 100644
--- a/src/coreclr/jit/rangecheck.cpp
+++ b/src/coreclr/jit/rangecheck.cpp
@@ -935,7 +935,7 @@ void RangeCheck::MergeAssertion(BasicBlock* block, GenTree* op, Range* pRange DE
     {
         GenTreePhiArg* arg  = (GenTreePhiArg*)op;
         BasicBlock*    pred = arg->gtPredBB;
-        if (pred->bbFallsThrough() && pred->NextIs(block))
+        if (pred->KindIs(BBJ_COND) && pred->FalseTargetIs(block))
         {
             assertions = pred->bbAssertionOut;
             JITDUMP("Merge assertions from pred " FMT_BB " edge: ", pred->bbNum);