From 517dc51c488739e45b9a6ad67c31f85778f43459 Mon Sep 17 00:00:00 2001
From: Florian Hahn <florian.hahn@arm.com>
Date: Tue, 6 Mar 2018 14:00:58 +0000
Subject: [PATCH] [CallSiteSplitting] Do not crash when BB's terminator
 changes.

Change doCallSiteSplitting to iterate until we reach the terminator instruction.
tryToSplitCallSite can replace BB's terminator in case BB is a successor of
itself. Then IE will be invalidated and we also have to check the current
terminator.

Reviewers: junbuml, davidxl, davide, fhahn

Reviewed By: fhahn, junbuml

Differential Revision: https://reviews.llvm.org/D43824

llvm-svn: 326793
---
 llvm/lib/Transforms/Scalar/CallSiteSplitting.cpp   |  8 +-
 .../Transforms/CallSiteSplitting/split-loop.ll     | 90 ++++++++++++++++++++++
 2 files changed, 97 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/Transforms/CallSiteSplitting/split-loop.ll

diff --git a/llvm/lib/Transforms/Scalar/CallSiteSplitting.cpp b/llvm/lib/Transforms/Scalar/CallSiteSplitting.cpp
index dcd21e8..ddebaf6 100644
--- a/llvm/lib/Transforms/Scalar/CallSiteSplitting.cpp
+++ b/llvm/lib/Transforms/Scalar/CallSiteSplitting.cpp
@@ -469,7 +469,13 @@ static bool doCallSiteSplitting(Function &F, TargetLibraryInfo &TLI,
   bool Changed = false;
   for (Function::iterator BI = F.begin(), BE = F.end(); BI != BE;) {
     BasicBlock &BB = *BI++;
-    for (BasicBlock::iterator II = BB.begin(), IE = BB.end(); II != IE;) {
+    auto II = BB.getFirstNonPHIOrDbg()->getIterator();
+    auto IE = BB.getTerminator()->getIterator();
+    // Iterate until we reach the terminator instruction. tryToSplitCallSite
+    // can replace BB's terminator in case BB is a successor of itself. In that
+    // case, IE will be invalidated and we also have to check the current
+    // terminator.
+    while (II != IE && &*II != BB.getTerminator()) {
       Instruction *I = &*II++;
       CallSite CS(cast<Value>(I));
       if (!CS || isa<IntrinsicInst>(I) || isInstructionTriviallyDead(I, &TLI))
diff --git a/llvm/test/Transforms/CallSiteSplitting/split-loop.ll b/llvm/test/Transforms/CallSiteSplitting/split-loop.ll
new file mode 100644
index 0000000..3e49a73
--- /dev/null
+++ b/llvm/test/Transforms/CallSiteSplitting/split-loop.ll
@@ -0,0 +1,90 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -S -callsite-splitting -simplifycfg < %s | FileCheck %s
+
+define i16 @test1() {
+; CHECK-LABEL: @test1(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[SPEC_SELECT:%.*]] = select i1 undef, i16 1, i16 0
+; CHECK-NEXT:    call void @callee(i16 0)
+; CHECK-NEXT:    br label [[FOR_COND12:%.*]]
+; CHECK:       for.cond12:
+; CHECK-NEXT:    call void @callee(i16 [[SPEC_SELECT]])
+; CHECK-NEXT:    br label [[FOR_COND12]]
+;
+entry:
+  %spec.select = select i1 undef, i16 1, i16 0
+  %tobool18 = icmp ne i16 %spec.select, 0
+  br i1 %tobool18, label %for.cond12.us, label %for.cond12
+
+for.cond12.us:
+  unreachable
+
+for.cond12:
+  call void @callee(i16 %spec.select)
+  br label %for.cond12
+}
+
+define i16 @test2() {
+; CHECK-LABEL: @test2(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[S:%.*]] = select i1 undef, i16 1, i16 0
+; CHECK-NEXT:    call void @callee(i16 0)
+; CHECK-NEXT:    br label [[FOR_COND12:%.*]]
+; CHECK:       for.cond12:
+; CHECK-NEXT:    [[ADD:%.*]] = add i16 [[S]], 10
+; CHECK-NEXT:    [[ADD2:%.*]] = add i16 [[S]], 10
+; CHECK-NEXT:    call void @callee(i16 [[S]])
+; CHECK-NEXT:    br label [[FOR_COND12]]
+;
+entry:
+  %s= select i1 undef, i16 1, i16 0
+  %tobool18 = icmp ne i16 %s, 0
+  br i1 %tobool18, label %for.cond12.us, label %for.cond12
+
+for.cond12.us:
+  unreachable
+
+for.cond12:
+  call void @callee(i16 %s)
+  %add = add i16 %s, 10
+  %add2 = add i16 %s, 10
+  br label %for.cond12
+}
+
+define i16 @test3() {
+; CHECK-LABEL: @test3(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[S:%.*]] = select i1 undef, i16 1, i16 0
+; CHECK-NEXT:    call void @callee(i16 0)
+; CHECK-NEXT:    br label [[FOR_COND12:%.*]]
+; CHECK:       for.cond12:
+; CHECK-NEXT:    [[ADD:%.*]] = add i16 [[S]], 10
+; CHECK-NEXT:    [[ADD2:%.*]] = add i16 [[ADD]], 10
+; CHECK-NEXT:    br i1 undef, label [[FOR_COND12_SPLIT:%.*]], label [[EXIT:%.*]]
+; CHECK:       for.cond12.split:
+; CHECK-NEXT:    call void @callee(i16 [[S]])
+; CHECK-NEXT:    br label [[FOR_COND12]]
+; CHECK:       exit:
+; CHECK-NEXT:    ret i16 [[ADD2]]
+;
+entry:
+  %s= select i1 undef, i16 1, i16 0
+  %tobool18 = icmp ne i16 %s, 0
+  br i1 %tobool18, label %for.cond12.us, label %for.cond12
+
+for.cond12.us:
+  unreachable
+
+for.cond12:
+  call void @callee(i16 %s)
+  %add = add i16 %s, 10
+  %add2 = add i16 %add, 10
+  br i1 undef, label %for.cond12, label %exit
+
+exit:
+  ret i16 %add2
+}
+
+define internal void @callee(i16 %flag) {
+  ret void
+}
-- 
2.7.4