From 0b7e8baf83befdb9cd7c282d943492581f2dd17c Mon Sep 17 00:00:00 2001
From: Denis Revunov <revunov.denis@huawei-partners.com>
Date: Wed, 8 Jun 2022 15:08:31 -0700
Subject: [PATCH] [BOLT][AArch64] Handle data at the beginning of a function
 when disassembling and building CFG.

This patch adds getFirstInstructionOffset method for BinaryFunction
which is used to properly handle cases where data is at zero offset in
a function. The main change is that we add basic block at first
instruction offset when disassembling, which prevents assertion
failures in buildCFG.

Reviewed By: yota9, rafauler

Differential Revision: https://reviews.llvm.org/D127111
---
 bolt/include/bolt/Core/BinaryFunction.h |  9 +++++++++
 bolt/lib/Core/BinaryFunction.cpp        | 11 +++++++----
 bolt/lib/Core/Exceptions.cpp            |  2 +-
 bolt/test/AArch64/data-at-0-offset.c    | 17 +++++++++++++++++
 4 files changed, 34 insertions(+), 5 deletions(-)
 create mode 100644 bolt/test/AArch64/data-at-0-offset.c

diff --git a/bolt/include/bolt/Core/BinaryFunction.h b/bolt/include/bolt/Core/BinaryFunction.h
index 7cb2dbc..58df5b9 100644
--- a/bolt/include/bolt/Core/BinaryFunction.h
+++ b/bolt/include/bolt/Core/BinaryFunction.h
@@ -982,6 +982,15 @@ public:
     return const_cast<BinaryFunction *>(this)->getInstructionAtOffset(Offset);
   }
 
+  /// Return offset for the first instruction. If there is data at the
+  /// beginning of a function then offset of the first instruction could
+  /// be different from 0
+  uint64_t getFirstInstructionOffset() const {
+    if (Instructions.empty())
+      return 0;
+    return Instructions.begin()->first;
+  }
+
   /// Return jump table that covers a given \p Address in memory.
   JumpTable *getJumpTableContainingAddress(uint64_t Address) {
     auto JTI = JumpTables.upper_bound(Address);
diff --git a/bolt/lib/Core/BinaryFunction.cpp b/bolt/lib/Core/BinaryFunction.cpp
index 6784763..8ff29b6 100644
--- a/bolt/lib/Core/BinaryFunction.cpp
+++ b/bolt/lib/Core/BinaryFunction.cpp
@@ -1383,6 +1383,9 @@ add_instruction:
   // Reset symbolizer for the disassembler.
   BC.SymbolicDisAsm->setSymbolizer(nullptr);
 
+  if (uint64_t Offset = getFirstInstructionOffset())
+    Labels[Offset] = BC.Ctx->createNamedTempSymbol();
+
   clearList(Relocations);
 
   if (!IsSimple) {
@@ -1895,7 +1898,7 @@ bool BinaryFunction::buildCFG(MCPlusBuilder::AllocatorIdTy AllocatorId) {
     return false;
 
   assert(BasicBlocks.empty() && "basic block list should be empty");
-  assert((Labels.find(0) != Labels.end()) &&
+  assert((Labels.find(getFirstInstructionOffset()) != Labels.end()) &&
          "first instruction should always have a label");
 
   // Create basic blocks in the original layout order:
@@ -1999,9 +2002,9 @@ bool BinaryFunction::buildCFG(MCPlusBuilder::AllocatorIdTy AllocatorId) {
         updateOffset(LastInstrOffset);
       }
     }
-    if (Offset == 0) {
-      // Add associated CFI pseudos in the first offset (0)
-      addCFIPlaceholders(0, InsertBB);
+    if (Offset == getFirstInstructionOffset()) {
+      // Add associated CFI pseudos in the first offset
+      addCFIPlaceholders(Offset, InsertBB);
     }
 
     const bool IsBlockEnd = MIB->isTerminator(Instr);
diff --git a/bolt/lib/Core/Exceptions.cpp b/bolt/lib/Core/Exceptions.cpp
index f5217de..b3ee89b 100644
--- a/bolt/lib/Core/Exceptions.cpp
+++ b/bolt/lib/Core/Exceptions.cpp
@@ -493,7 +493,7 @@ bool CFIReaderWriter::fillCFIInfoFor(BinaryFunction &Function) const {
   Optional<uint64_t> LSDA = CurFDE.getLSDAAddress();
   Function.setLSDAAddress(LSDA ? *LSDA : 0);
 
-  uint64_t Offset = 0;
+  uint64_t Offset = Function.getFirstInstructionOffset();
   uint64_t CodeAlignment = CurFDE.getLinkedCIE()->getCodeAlignmentFactor();
   uint64_t DataAlignment = CurFDE.getLinkedCIE()->getDataAlignmentFactor();
   if (CurFDE.getLinkedCIE()->getPersonalityAddress()) {
diff --git a/bolt/test/AArch64/data-at-0-offset.c b/bolt/test/AArch64/data-at-0-offset.c
new file mode 100644
index 0000000..e0c689a
--- /dev/null
+++ b/bolt/test/AArch64/data-at-0-offset.c
@@ -0,0 +1,17 @@
+// RUN: %clang %cflags -O2 -fPIE -Wl,-q -pie  %s -o %t.exe
+// RUN: llvm-bolt %t.exe -o %t.bolt 2>&1 | FileCheck %s
+// CHECK-NOT: BOLT-WARNING: unable to disassemble instruction at offset
+
+void extra_space() {
+  asm volatile(".rept 256\n"
+               "    .byte 0xff\n"
+               ".endr\n");
+  return;
+}
+
+int main(int argc, char **argv) {
+  void (*fn)(void);
+  fn = extra_space + 256;
+  fn();
+  return 0;
+}
-- 
2.7.4