[llvm] [BasicBlockUtils] Fix loopinfo header in split block before (PR #179408)

[Yevgeny Rouban via llvm-commits]

https://github.com/yrouban updated https://github.com/llvm/llvm-project/pull/179408

>From c0e312713d7e99bb44a0eb75a27cd096decc934f Mon Sep 17 00:00:00 2001
From: Yevgeny Rouban <yrouban at gmail.com>
Date: Tue, 3 Feb 2026 09:02:27 +0300
Subject: [PATCH 1/2] [BasicBlockUtilsTests] SplitBlock breaks LoopInfo. NFC

Splitting a basic block BB into a pair of blocks NewBB->BB
SplitBlock(... DT, &LI,... /* Before */ true) does not change
NewBB to be the new head of the loop if BB was the loop head
before the split.

This change just introduces a new unit test that demonstrates
the bug.
---
 .../Transforms/Utils/BasicBlockUtilsTest.cpp  | 34 +++++++++++++++++++
 1 file changed, 34 insertions(+)

diff --git a/llvm/unittests/Transforms/Utils/BasicBlockUtilsTest.cpp b/llvm/unittests/Transforms/Utils/BasicBlockUtilsTest.cpp
index 00d9e9ff81e05..37b20d00d18d9 100644
--- a/llvm/unittests/Transforms/Utils/BasicBlockUtilsTest.cpp
+++ b/llvm/unittests/Transforms/Utils/BasicBlockUtilsTest.cpp
@@ -26,6 +26,8 @@
 #include "llvm/Support/SourceMgr.h"
 #include "gtest/gtest.h"
 
+#define DEBUG_TYPE "basic-block-utils-tests"
+
 using namespace llvm;
 
 static std::unique_ptr<Module> parseIR(LLVMContext &C, const char *IR) {
@@ -354,6 +356,38 @@ define void @foo() {
 }
 #endif
 
+#ifndef NDEBUG
+TEST(BasicBlockUtils, SplitBlockBefore) {
+  LLVMContext C;
+  std::unique_ptr<Module> M = parseIR(C, R"IR(
+define void @split-block-before-test(i1 %flag) {
+entry:
+  br label %loop
+
+loop:
+  br i1 %flag, label %loop, label %exit
+
+exit:
+  ret void
+}
+)IR");
+  Function *F = M->getFunction("split-block-before-test");
+  DominatorTree DT(*F);
+  LoopInfo LI(DT);
+
+  EXPECT_TRUE(DT.verify());
+  LI.verify(DT);
+  auto *LoopBB = getBasicBlockByName(*F, "loop");
+  SplitBlock(LoopBB, LoopBB->getFirstInsertionPt(), &DT, &LI,
+              /* MemorySSAUpdater */ nullptr, LoopBB->getName() + ".split",
+              /* Before */ true);
+
+  EXPECT_TRUE(DT.verify());
+  LLVM_DEBUG(F->dump(); LI.print(dbgs()));
+  EXPECT_DEATH({LI.verify(DT);}, "Assertion.*Loop is unreachable!");
+}
+#endif
+
 TEST(BasicBlockUtils, NoUnreachableBlocksToEliminate) {
   LLVMContext C;
   std::unique_ptr<Module> M = parseIR(C, R"IR(

>From ef0bf05db95cb097cba0e94f0c732cd1019f718a Mon Sep 17 00:00:00 2001
From: Yevgeny Rouban <yrouban at gmail.com>
Date: Tue, 3 Feb 2026 10:38:06 +0300
Subject: [PATCH 2/2] [BasicBlockUtils] Fix loop header in splitBlockBefore()

If BB was the loop header before NewBB = splitBlockBefore(BB, ...)
then the NewBB must become the new header of the loop.
---
 llvm/lib/Transforms/Utils/BasicBlockUtils.cpp          |  5 ++++-
 .../unittests/Transforms/Utils/BasicBlockUtilsTest.cpp | 10 ++++++----
 2 files changed, 10 insertions(+), 5 deletions(-)

diff --git a/llvm/lib/Transforms/Utils/BasicBlockUtils.cpp b/llvm/lib/Transforms/Utils/BasicBlockUtils.cpp
index b0c04086f5e84..381cb7170c0f6 100644
--- a/llvm/lib/Transforms/Utils/BasicBlockUtils.cpp
+++ b/llvm/lib/Transforms/Utils/BasicBlockUtils.cpp
@@ -1080,8 +1080,11 @@ BasicBlock *llvm::splitBlockBefore(BasicBlock *Old, BasicBlock::iterator SplitPt
   // The new block lives in whichever loop the old one did. This preserves
   // LCSSA as well, because we force the split point to be after any PHI nodes.
   if (LI)
-    if (Loop *L = LI->getLoopFor(Old))
+    if (Loop *L = LI->getLoopFor(Old)) {
       L->addBasicBlockToLoop(New, *LI);
+      if (L->getHeader() == Old)
+        L->moveToHeader(New);
+    }
 
   if (DTU) {
     SmallVector<DominatorTree::UpdateType, 8> DTUpdates;
diff --git a/llvm/unittests/Transforms/Utils/BasicBlockUtilsTest.cpp b/llvm/unittests/Transforms/Utils/BasicBlockUtilsTest.cpp
index 37b20d00d18d9..198518e130551 100644
--- a/llvm/unittests/Transforms/Utils/BasicBlockUtilsTest.cpp
+++ b/llvm/unittests/Transforms/Utils/BasicBlockUtilsTest.cpp
@@ -378,13 +378,15 @@ define void @split-block-before-test(i1 %flag) {
   EXPECT_TRUE(DT.verify());
   LI.verify(DT);
   auto *LoopBB = getBasicBlockByName(*F, "loop");
-  SplitBlock(LoopBB, LoopBB->getFirstInsertionPt(), &DT, &LI,
-              /* MemorySSAUpdater */ nullptr, LoopBB->getName() + ".split",
-              /* Before */ true);
+  auto *New =
+      SplitBlock(LoopBB, LoopBB->getFirstInsertionPt(), &DT, &LI,
+                 /* MemorySSAUpdater */ nullptr, LoopBB->getName() + ".split",
+                 /* Before */ true);
 
   EXPECT_TRUE(DT.verify());
   LLVM_DEBUG(F->dump(); LI.print(dbgs()));
-  EXPECT_DEATH({LI.verify(DT);}, "Assertion.*Loop is unreachable!");
+  LI.verify(DT);
+  EXPECT_EQ(LI.getLoopFor(New)->getHeader(), New);
 }
 #endif