From bd606afe26f258d081fbe025b20a71b277c1edde Mon Sep 17 00:00:00 2001
From: Zhiyao Ma <zhiyao.ma.98@gmail.com>
Date: Tue, 3 May 2022 12:29:48 +0100
Subject: [PATCH] [ARM] Only update the successor edges for immediate
 predecessors of PrologueMBB

When adjusting the function prologue for segmented stacks, only update
the successor edges of the immediate predecessors of the original
prologue.

Differential Revision: https://reviews.llvm.org/D122959
---
 llvm/lib/Target/ARM/ARMFrameLowering.cpp     |  5 +++--
 llvm/test/CodeGen/ARM/segmented-stacks.ll    | 20 ++++++++++++++++++++
 llvm/test/CodeGen/Thumb/segmented-stacks.ll  | 20 ++++++++++++++++++++
 llvm/test/CodeGen/Thumb2/segmented-stacks.ll | 20 ++++++++++++++++++++
 4 files changed, 63 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Target/ARM/ARMFrameLowering.cpp b/llvm/lib/Target/ARM/ARMFrameLowering.cpp
index c2cc302..b2765b2 100644
--- a/llvm/lib/Target/ARM/ARMFrameLowering.cpp
+++ b/llvm/lib/Target/ARM/ARMFrameLowering.cpp
@@ -2562,8 +2562,9 @@ void ARMFrameLowering::adjustForSegmentedStacks(
     // Make sure the LiveIns are still sorted and unique.
     MBB->sortUniqueLiveIns();
     // Replace the edges to PrologueMBB by edges to the sequences
-    // we are about to add.
-    MBB->ReplaceUsesOfBlockWith(&PrologueMBB, AddedBlocks[0]);
+    // we are about to add, but only update for immediate predecessors.
+    if (MBB->isSuccessor(&PrologueMBB))
+      MBB->ReplaceUsesOfBlockWith(&PrologueMBB, AddedBlocks[0]);
   }
 
   // The required stack size that is aligned to ARM constant criterion.
diff --git a/llvm/test/CodeGen/ARM/segmented-stacks.ll b/llvm/test/CodeGen/ARM/segmented-stacks.ll
index 10ce065..3861712 100644
--- a/llvm/test/CodeGen/ARM/segmented-stacks.ll
+++ b/llvm/test/CodeGen/ARM/segmented-stacks.ll
@@ -280,4 +280,24 @@ define i32 @test_sibling_call_empty_frame(i32 %x) #0 {
 
 }
 
+
+declare void @panic() unnamed_addr
+
+; We used to crash while compiling the following function.
+; ARM-linux-LABEL: build_should_not_segfault:
+; ARM-android-LABEL: build_should_not_segfault:
+define void @build_should_not_segfault(i8 %x) unnamed_addr #0 {
+start:
+  %_0 = icmp ult i8 %x, 16
+  %or.cond = select i1 undef, i1 true, i1 %_0
+  br i1 %or.cond, label %bb1, label %bb2
+
+bb1:
+  ret void
+
+bb2:
+  call void @panic()
+  unreachable
+}
+
 attributes #0 = { "split-stack" }
diff --git a/llvm/test/CodeGen/Thumb/segmented-stacks.ll b/llvm/test/CodeGen/Thumb/segmented-stacks.ll
index 6ad8cf9..daebafe 100644
--- a/llvm/test/CodeGen/Thumb/segmented-stacks.ll
+++ b/llvm/test/CodeGen/Thumb/segmented-stacks.ll
@@ -275,4 +275,24 @@ define void @test_nostack() #0 {
 ; Thumb-linux-NOT:   bl __morestack
 }
 
+
+declare void @panic() unnamed_addr
+
+; We used to crash while compiling the following function.
+; Thumb-linux-LABEL: build_should_not_segfault:
+; Thumb-android-LABEL: build_should_not_segfault:
+define void @build_should_not_segfault(i8 %x) unnamed_addr #0 {
+start:
+  %_0 = icmp ult i8 %x, 16
+  %or.cond = select i1 undef, i1 true, i1 %_0
+  br i1 %or.cond, label %bb1, label %bb2
+
+bb1:
+  ret void
+
+bb2:
+  call void @panic()
+  unreachable
+}
+
 attributes #0 = { "split-stack" }
diff --git a/llvm/test/CodeGen/Thumb2/segmented-stacks.ll b/llvm/test/CodeGen/Thumb2/segmented-stacks.ll
index 0f34e9c..70892073 100644
--- a/llvm/test/CodeGen/Thumb2/segmented-stacks.ll
+++ b/llvm/test/CodeGen/Thumb2/segmented-stacks.ll
@@ -171,4 +171,24 @@ define fastcc void @test_fastcc_large() #0 {
 ; ARM-NEXT:    .long   40192
 }
 
+
+declare void @panic() unnamed_addr
+
+; We used to crash while compiling the following function.
+; THUMB-LABEL: build_should_not_segfault:
+; ARM-LABEL: build_should_not_segfault:
+define void @build_should_not_segfault(i8 %x) unnamed_addr #0 {
+start:
+  %_0 = icmp ult i8 %x, 16
+  %or.cond = select i1 undef, i1 true, i1 %_0
+  br i1 %or.cond, label %bb1, label %bb2
+
+bb1:
+  ret void
+
+bb2:
+  call void @panic()
+  unreachable
+}
+
 attributes #0 = { "split-stack" }
-- 
2.7.4