From d77f96a9e0c86a35f5d5986058b4dcce30113cca Mon Sep 17 00:00:00 2001
From: Amir Ayupov <aaupov@fb.com>
Date: Mon, 27 Feb 2023 15:40:45 -0800
Subject: [PATCH] [BOLT][NFC] Simplify BinaryFunction::setTrapOnEntry

Reviewed By: #bolt, maksfb

Differential Revision: https://reviews.llvm.org/D144758
---
 bolt/include/bolt/Core/BinaryFunction.h |  2 --
 bolt/lib/Core/BinaryFunction.cpp        | 10 +++-------
 2 files changed, 3 insertions(+), 9 deletions(-)

diff --git a/bolt/include/bolt/Core/BinaryFunction.h b/bolt/include/bolt/Core/BinaryFunction.h
index ddff131..051a585 100644
--- a/bolt/include/bolt/Core/BinaryFunction.h
+++ b/bolt/include/bolt/Core/BinaryFunction.h
@@ -1466,8 +1466,6 @@ public:
 
   ArrayRef<uint8_t> getLSDATypeIndexTable() const { return LSDATypeIndexTable; }
 
-  const LabelsMapType &getLabels() const { return Labels; }
-
   IslandInfo &getIslandInfo() {
     assert(Islands && "function expected to have constant islands");
     return *Islands;
diff --git a/bolt/lib/Core/BinaryFunction.cpp b/bolt/lib/Core/BinaryFunction.cpp
index 13c5fd4..01e7de5 100644
--- a/bolt/lib/Core/BinaryFunction.cpp
+++ b/bolt/lib/Core/BinaryFunction.cpp
@@ -2902,16 +2902,12 @@ void BinaryFunction::clearDisasmState() {
 void BinaryFunction::setTrapOnEntry() {
   clearDisasmState();
 
-  auto addTrapAtOffset = [&](uint64_t Offset) {
+  forEachEntryPoint([&](uint64_t Offset, const MCSymbol *Label) -> bool {
     MCInst TrapInstr;
     BC.MIB->createTrap(TrapInstr);
     addInstruction(Offset, std::move(TrapInstr));
-  };
-
-  addTrapAtOffset(0);
-  for (const std::pair<const uint32_t, MCSymbol *> &KV : getLabels())
-    if (getSecondaryEntryPointSymbol(KV.second))
-      addTrapAtOffset(KV.first);
+    return true;
+  });
 
   TrapsOnEntry = true;
 }
-- 
2.7.4