From f9d00d418b3e095bf0b551f7c88cd2fd0a41d4aa Mon Sep 17 00:00:00 2001
From: Amir Ayupov <aaupov@fb.com>
Date: Tue, 17 Nov 2020 18:40:19 -0800
Subject: [PATCH] [BOLT] Handle insertion of updated CFI at the first basic
 block

Summary:
Fix corner case of insertion of updated CFI with unset `PrevBB`.
Handle it in the same way as inserting past hot-cold split point.

(cherry picked from FBD24943911)
---
 bolt/src/Passes/ShrinkWrapping.cpp       |  4 +--
 bolt/test/X86/shrinkwrapping-insertcfi.s | 41 ++++++++++++++++++++++++
 2 files changed, 43 insertions(+), 2 deletions(-)
 create mode 100644 bolt/test/X86/shrinkwrapping-insertcfi.s

diff --git a/bolt/src/Passes/ShrinkWrapping.cpp b/bolt/src/Passes/ShrinkWrapping.cpp
index a8e99014f515..94f3dcd2b6c1 100644
--- a/bolt/src/Passes/ShrinkWrapping.cpp
+++ b/bolt/src/Passes/ShrinkWrapping.cpp
@@ -1576,8 +1576,8 @@ void ShrinkWrapping::insertUpdatedCFI(unsigned CSR, int SPValPush,
         }
       }
     }
-    // Are we at the hot-cold split point?
-    if (BF.isSplit() && PrevBB && BB->isCold() != PrevBB->isCold()) {
+    // Are we at the first basic block or hot-cold split point?
+    if (!PrevBB || (BF.isSplit() && BB->isCold() != PrevBB->isCold())) {
       if (InAffectedZoneAtBegin) {
         insertCFIsForPushOrPop(*BB, BB->begin(), CSR, true, 0, SPValPush);
       }
diff --git a/bolt/test/X86/shrinkwrapping-insertcfi.s b/bolt/test/X86/shrinkwrapping-insertcfi.s
new file mode 100644
index 000000000000..83eb2eb869f7
--- /dev/null
+++ b/bolt/test/X86/shrinkwrapping-insertcfi.s
@@ -0,0 +1,41 @@
+# This test reproduces the issue with inserting updated CFI in shrink wrapping
+# into the first basic block.
+
+# REQUIRES: system-linux
+
+# RUN: llvm-mc -filetype=obj -triple x86_64-unknown-unknown %s -o %t.o
+# RUN: link_fdata %s %t.o %t.fdata
+# RUN: strip --strip-unneeded %t.o
+# RUN: %host_cc %t.o -o %t.exe -Wl,-q -nostdlib
+# RUN: llvm-bolt %t.exe -o %t.out -data %t.fdata -frame-opt=all -lite=0 \
+# RUN:           -print-fop 2>&1 | FileCheck %s
+
+# Check shrink wrapping results:
+# CHECK: BOLT-INFO: Shrink wrapping moved 0 spills inserting load/stores and 1 spills inserting push/pops
+
+# Check that CFI is successfully inserted into the first basic block:
+# CHECK:      Binary Function "_start" after frame-optimizer
+# CHECK:      .LBB0 (2 instructions, align : 1)
+# CHECK-NEXT: Entry Point
+# CHECK:      00000000:   !CFI {{.*}}
+# CHECK-NEXT: 00000000:   je  .Ltmp{{.*}}
+
+  .globl _start
+_start:
+    .cfi_startproc
+# FDATA: 0 [unknown] 0 1 _start 0 0 6
+# !CFI OpOffset for reg3/rbx is inserted into this block.
+    je	a
+b:  jne _start
+# FDATA: 1 _start #b# 1 _start #c# 0 3
+
+c:
+    push	%rbx
+    .cfi_offset 3, 4
+    pop	%rbx
+
+# This basic block is treated as having 0 execution count.
+# push and pop will be sinked into this block.
+a:
+    ud2
+    .cfi_endproc