From de284bc510b8d6d8a7ebb122f3ac6373bfbad093 Mon Sep 17 00:00:00 2001
From: Rafael Auler <rafaelauler@fb.com>
Date: Wed, 18 Dec 2019 12:14:42 -0800
Subject: [PATCH] [BOLT] Fix symbol table entries for secondary entries

Summary:
Commit "Support full instrumentation" changed the map
SymbolToFunction in BinaryContext to map secondary entries of functions
too. This introduced unexpected behavior in our symbol table rewriting
logic, which caused it to mistakenly write them with the address of the
original function. Fix the behavior of getBinaryFunctionAtAddress to
correct this. Also fix other users of SymbolToFunction to ensure they
are not accidentally using secondary entries when they shouldn't.

(cherry picked from FBD19168319)
---
 bolt/src/BinaryContext.cpp               |  3 ++-
 bolt/src/BinaryFunction.cpp              |  2 +-
 bolt/src/BinaryFunction.h                | 12 +++++++++++-
 bolt/src/Passes/IdenticalCodeFolding.cpp |  4 ++++
 bolt/src/Passes/Inliner.cpp              |  4 ++++
 bolt/src/Passes/LongJmp.cpp              |  3 ++-
 6 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/bolt/src/BinaryContext.cpp b/bolt/src/BinaryContext.cpp
index 18277d7d2518..b13eb97833a0 100644
--- a/bolt/src/BinaryContext.cpp
+++ b/bolt/src/BinaryContext.cpp
@@ -1910,7 +1910,8 @@ BinaryContext::getBinaryFunctionContainingAddress(uint64_t Address,
 BinaryFunction *
 BinaryContext::getBinaryFunctionAtAddress(uint64_t Address, bool Shallow) {
   if (const auto *BD = getBinaryDataAtAddress(Address)) {
-    if (auto *BF = getFunctionForSymbol(BD->getSymbol())) {
+    auto *BF = getFunctionForSymbol(BD->getSymbol());
+    if (BF && BF->getAddress() == Address) {
       while (BF->getParentFunction() && !Shallow) {
         BF = BF->getParentFunction();
       }
diff --git a/bolt/src/BinaryFunction.cpp b/bolt/src/BinaryFunction.cpp
index cff497543032..48f860908d1f 100644
--- a/bolt/src/BinaryFunction.cpp
+++ b/bolt/src/BinaryFunction.cpp
@@ -3391,7 +3391,7 @@ uint64_t BinaryFunction::getEntryForSymbol(const MCSymbol *EntrySymbol) const {
     ++NumEntries;
   }
 
-  llvm_unreachable("no entry for symbol");
+  return std::numeric_limits<uint64_t>::max();
 }
 
 BinaryFunction::BasicBlockOrderType BinaryFunction::dfs() const {
diff --git a/bolt/src/BinaryFunction.h b/bolt/src/BinaryFunction.h
index 3c17969786cf..9ce2cd15cc71 100644
--- a/bolt/src/BinaryFunction.h
+++ b/bolt/src/BinaryFunction.h
@@ -1061,9 +1061,19 @@ public:
   /// functions.
   const MCSymbol *getSymbolForEntry(uint64_t EntryNum) const;
 
-  /// Return an entry ID corresponding to a symbol.
+  /// Return an entry ID corresponding to a symbol. Return
+  /// numeric_limits<T>::max() if entry does not exist.
   uint64_t getEntryForSymbol(const MCSymbol *EntrySymbol) const;
 
+  /// Return true if \p is a valid secondary entry point in this function, false
+  /// otherwise.
+  bool isSecondaryEntryPoint(const MCSymbol *Label) const {
+    uint64_t EntryNo = getEntryForSymbol(Label);
+    if (EntryNo == std::numeric_limits<uint64_t>::max())
+      return false;
+    return EntryNo != 0;
+  }
+
   MCSymbol *getColdSymbol() {
     if (ColdSymbol)
       return ColdSymbol;
diff --git a/bolt/src/Passes/IdenticalCodeFolding.cpp b/bolt/src/Passes/IdenticalCodeFolding.cpp
index 25609a32bba1..71c9c8423bc6 100644
--- a/bolt/src/Passes/IdenticalCodeFolding.cpp
+++ b/bolt/src/Passes/IdenticalCodeFolding.cpp
@@ -225,6 +225,10 @@ bool isIdenticalWith(const BinaryFunction &A, const BinaryFunction &B,
           // Compare symbols as functions.
           const auto *FunctionA = BC.getFunctionForSymbol(SymbolA);
           const auto *FunctionB = BC.getFunctionForSymbol(SymbolB);
+          if (FunctionA && FunctionA->isSecondaryEntryPoint(SymbolA))
+            FunctionA = nullptr;
+          if (FunctionB && FunctionB->isSecondaryEntryPoint(SymbolB))
+            FunctionB = nullptr;
           if (FunctionA && FunctionB) {
             // Self-referencing functions and recursive calls.
             if (FunctionA == &A && FunctionB == &B)
diff --git a/bolt/src/Passes/Inliner.cpp b/bolt/src/Passes/Inliner.cpp
index 05283a759f02..bcb8166525d2 100644
--- a/bolt/src/Passes/Inliner.cpp
+++ b/bolt/src/Passes/Inliner.cpp
@@ -461,6 +461,10 @@ bool Inliner::inlineCallsInFunction(BinaryFunction &Function) {
         continue;
       }
 
+      // Don't inline calls to a secondary entry point in a target function
+      if (TargetFunction->isSecondaryEntryPoint(TargetSymbol))
+        continue;
+
       // Don't do recursive inlining.
       if (TargetFunction == &Function) {
         ++InstIt;
diff --git a/bolt/src/Passes/LongJmp.cpp b/bolt/src/Passes/LongJmp.cpp
index a7a8bc0eb513..a9cd6af5819e 100644
--- a/bolt/src/Passes/LongJmp.cpp
+++ b/bolt/src/Passes/LongJmp.cpp
@@ -424,7 +424,8 @@ uint64_t LongJmpPass::getSymbolAddress(const BinaryContext &BC,
   }
   auto *TargetFunc = BC.getFunctionForSymbol(Target);
   auto Iter = HotAddresses.find(TargetFunc);
-  if (Iter == HotAddresses.end()) {
+  if (Iter == HotAddresses.end() ||
+      (TargetFunc && TargetFunc->isSecondaryEntryPoint(Target))) {
     // Look at BinaryContext's resolution for this symbol - this is a symbol not
     // mapped to a BinaryFunction
     auto ValueOrError = BC.getSymbolValue(*Target);