diff --git a/bolt/include/bolt/Core/MCPlusBuilder.h b/bolt/include/bolt/Core/MCPlusBuilder.h
@@ -19,6 +19,7 @@
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/BitVector.h"
 #include "llvm/ADT/StringMap.h"
+#include "llvm/CodeGen/TargetOpcodes.h"
 #include "llvm/MC/MCAsmBackend.h"
 #include "llvm/MC/MCDisassembler/MCSymbolizer.h"
 #include "llvm/MC/MCExpr.h"
@@ -27,6 +28,7 @@
 #include "llvm/MC/MCInstrDesc.h"
 #include "llvm/MC/MCInstrInfo.h"
 #include "llvm/Support/Allocator.h"
+#include "llvm/Support/Casting.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/ErrorOr.h"
 #include "llvm/Support/RWMutex.h"
@@ -533,9 +535,7 @@ class MCPlusBuilder {
     return Analysis->isReturn(Inst);
   }
 
-  virtual bool isTerminator(const MCInst &Inst) const {
-    return Analysis->isTerminator(Inst);
-  }
+  virtual bool isTerminator(const MCInst &Inst) const;
 
   virtual bool isNoop(const MCInst &Inst) const {
     llvm_unreachable("not implemented");

diff --git a/bolt/include/bolt/Passes/BinaryPasses.h b/bolt/include/bolt/Passes/BinaryPasses.h
@@ -18,7 +18,6 @@
 #include "bolt/Core/DynoStats.h"
 #include "llvm/Support/CommandLine.h"
 #include <atomic>
-#include <map>
 #include <set>
 #include <string>
 #include <unordered_set>

diff --git a/bolt/include/bolt/Passes/CacheMetrics.h b/bolt/include/bolt/Passes/CacheMetrics.h
@@ -13,7 +13,6 @@
 #ifndef BOLT_PASSES_CACHEMETRICS_H
 #define BOLT_PASSES_CACHEMETRICS_H
 
-#include <cstdint>
 #include <vector>
 
 namespace llvm {

diff --git a/bolt/include/bolt/Passes/DominatorAnalysis.h b/bolt/include/bolt/Passes/DominatorAnalysis.h
@@ -11,7 +11,6 @@
 
 #include "bolt/Passes/DataflowAnalysis.h"
 #include "llvm/Support/CommandLine.h"
-#include "llvm/Support/Timer.h"
 
 namespace opts {
 extern llvm::cl::opt<bool> TimeOpts;

diff --git a/bolt/include/bolt/Passes/ReachingDefOrUse.h b/bolt/include/bolt/Passes/ReachingDefOrUse.h
@@ -11,9 +11,7 @@
 
 #include "bolt/Passes/DataflowAnalysis.h"
 #include "bolt/Passes/RegAnalysis.h"
-#include "llvm/MC/MCRegisterInfo.h"
 #include "llvm/Support/CommandLine.h"
-#include "llvm/Support/Timer.h"
 #include <optional>
 
 namespace opts {

diff --git a/bolt/include/bolt/Passes/ReachingInsns.h b/bolt/include/bolt/Passes/ReachingInsns.h
@@ -11,7 +11,6 @@
 
 #include "bolt/Passes/DataflowAnalysis.h"
 #include "llvm/Support/CommandLine.h"
-#include "llvm/Support/Timer.h"
 
 namespace opts {
 extern llvm::cl::opt<bool> TimeOpts;

diff --git a/bolt/include/bolt/Passes/ReorderUtils.h b/bolt/include/bolt/Passes/ReorderUtils.h
@@ -14,7 +14,6 @@
 #ifndef BOLT_PASSES_REORDER_UTILS_H
 #define BOLT_PASSES_REORDER_UTILS_H
 
-#include <memory>
 #include <vector>
 
 #include "llvm/ADT/BitVector.h"

diff --git a/bolt/include/bolt/Profile/BoltAddressTranslation.h b/bolt/include/bolt/Profile/BoltAddressTranslation.h
@@ -19,6 +19,7 @@
 #include <unordered_map>
 
 namespace llvm {
+class MCSymbol;
 class raw_ostream;
 
 namespace object {
@@ -118,10 +119,12 @@ class BoltAddressTranslation {
   /// True if a given \p Address is a function with translation table entry.
   bool isBATFunction(uint64_t Address) const { return Maps.count(Address); }
 
-  /// Returns branch offsets grouped by containing basic block in a given
-  /// function.
-  std::unordered_map<uint32_t, std::vector<uint32_t>>
-  getBFBranches(uint64_t FuncOutputAddress) const;
+  /// For a given \p Symbol in the output binary and known \p InputOffset
+  /// return a corresponding pair of parent BinaryFunction and secondary entry
+  /// point in it.
+  std::pair<const BinaryFunction *, unsigned>
+  translateSymbol(const BinaryContext &BC, const MCSymbol &Symbol,
+                  uint32_t InputOffset) const;
 
 private:
   /// Helper to update \p Map by inserting one or more BAT entries reflecting
@@ -158,6 +161,10 @@ class BoltAddressTranslation {
   /// Map a function to its secondary entry points vector
   std::unordered_map<uint64_t, std::vector<uint32_t>> SecondaryEntryPointsMap;
 
+  /// Return a secondary entry point ID for a function located at \p Address and
+  /// \p Offset within that function.
+  unsigned getSecondaryEntryPointId(uint64_t Address, uint32_t Offset) const;
+
   /// Links outlined cold bocks to their original function
   std::map<uint64_t, uint64_t> ColdPartSource;
 
@@ -181,7 +188,7 @@ class BoltAddressTranslation {
       EntryTy(unsigned Index, size_t Hash) : Index(Index), Hash(Hash) {}
     };
 
-    std::unordered_map<uint32_t, EntryTy> Map;
+    std::map<uint32_t, EntryTy> Map;
     const EntryTy &getEntry(uint32_t BBInputOffset) const {
       auto It = Map.find(BBInputOffset);
       assert(It != Map.end());
@@ -206,6 +213,10 @@ class BoltAddressTranslation {
     }
 
     size_t getNumBasicBlocks() const { return Map.size(); }
+
+    auto begin() const { return Map.begin(); }
+    auto end() const { return Map.end(); }
+    auto upper_bound(uint32_t Offset) const { return Map.upper_bound(Offset); }
   };
 
   /// Map function output address to its hash and basic blocks hash map.

diff --git a/bolt/include/bolt/Profile/DataAggregator.h b/bolt/include/bolt/Profile/DataAggregator.h
@@ -225,6 +225,10 @@ class DataAggregator : public DataReader {
   /// Aggregation statistics
   uint64_t NumInvalidTraces{0};
   uint64_t NumLongRangeTraces{0};
+  /// Specifies how many samples were recorded in cold areas if we are dealing
+  /// with profiling data collected in a bolted binary. For LBRs, incremented
+  /// for the source of the branch to avoid counting cold activity twice (one
+  /// for source and another for destination).
   uint64_t NumColdSamples{0};
 
   /// Looks into system PATH for Linux Perf and set up the aggregator to use it
@@ -245,14 +249,12 @@ class DataAggregator : public DataReader {
   /// disassembled BinaryFunctions
   BinaryFunction *getBinaryFunctionContainingAddress(uint64_t Address) const;
 
+  /// Perform BAT translation for a given \p Func and return the parent
+  /// BinaryFunction or nullptr.
+  BinaryFunction *getBATParentFunction(const BinaryFunction &Func) const;
+
   /// Retrieve the location name to be used for samples recorded in \p Func.
-  /// If doing BAT translation, link cold parts to the hot part  names (used by
-  /// the original binary).  \p Count specifies how many samples were recorded
-  /// at that location, so we can tally total activity in cold areas if we are
-  /// dealing with profiling data collected in a bolted binary. For LBRs,
-  /// \p Count should only be used for the source of the branch to avoid
-  /// counting cold activity twice (one for source and another for destination).
-  StringRef getLocationName(BinaryFunction &Func, uint64_t Count);
+  StringRef getLocationName(const BinaryFunction &Func) const;
 
   /// Semantic actions - parser hooks to interpret parsed perf samples
   /// Register a sample (non-LBR mode), i.e. a new hit at \p Address
@@ -467,9 +469,6 @@ class DataAggregator : public DataReader {
   std::error_code writeBATYAML(BinaryContext &BC,
                                StringRef OutputFilename) const;
 
-  /// Fixup profile collected on BOLTed binary, namely handle split functions.
-  void fixupBATProfile(BinaryContext &BC);
-
   /// Filter out binaries based on PID
   void filterBinaryMMapInfo();
 

diff --git a/bolt/include/bolt/Profile/ProfileReaderBase.h b/bolt/include/bolt/Profile/ProfileReaderBase.h
@@ -65,7 +65,7 @@ class ProfileReaderBase {
   /// Return true if the function \p BF may have a profile available.
   /// The result is based on the name(s) of the function alone and the profile
   /// match is not guaranteed.
-  virtual bool mayHaveProfileData(const BinaryFunction &BF);
+  virtual bool mayHaveProfileData(const BinaryFunction &BF) { return true; }
 
   /// Return true if the profile contains an entry for a local object
   /// that has an associated file name.

diff --git a/bolt/include/bolt/Profile/ProfileYAMLMapping.h b/bolt/include/bolt/Profile/ProfileYAMLMapping.h
@@ -14,7 +14,6 @@
 #define BOLT_PROFILE_PROFILEYAMLMAPPING_H
 
 #include "bolt/Core/BinaryFunction.h"
-#include "llvm/ADT/StringRef.h"
 #include "llvm/Support/YAMLTraits.h"
 #include <vector>
 

diff --git a/bolt/include/bolt/Profile/YAMLProfileWriter.h b/bolt/include/bolt/Profile/YAMLProfileWriter.h
@@ -15,6 +15,7 @@
 
 namespace llvm {
 namespace bolt {
+class BoltAddressTranslation;
 class RewriteInstance;
 
 class YAMLProfileWriter {
@@ -31,8 +32,16 @@ class YAMLProfileWriter {
   /// Save execution profile for that instance.
   std::error_code writeProfile(const RewriteInstance &RI);
 
-  static yaml::bolt::BinaryFunctionProfile convert(const BinaryFunction &BF,
-                                                   bool UseDFS);
+  static yaml::bolt::BinaryFunctionProfile
+  convert(const BinaryFunction &BF, bool UseDFS,
+          const BoltAddressTranslation *BAT = nullptr);
+
+  /// Set CallSiteInfo destination fields from \p Symbol and return a target
+  /// BinaryFunction for that symbol.
+  static const BinaryFunction *
+  setCSIDestination(const BinaryContext &BC, yaml::bolt::CallSiteInfo &CSI,
+                    const MCSymbol *Symbol, const BoltAddressTranslation *BAT,
+                    uint32_t Offset = 0);
 };
 
 } // namespace bolt

diff --git a/bolt/include/bolt/Rewrite/DWARFRewriter.h b/bolt/include/bolt/Rewrite/DWARFRewriter.h
@@ -22,9 +22,7 @@
 #include <memory>
 #include <mutex>
 #include <optional>
-#include <set>
 #include <unordered_map>
-#include <unordered_set>
 #include <vector>
 
 namespace llvm {

diff --git a/bolt/include/bolt/Rewrite/MetadataManager.h b/bolt/include/bolt/Rewrite/MetadataManager.h
@@ -11,7 +11,6 @@
 
 #include "bolt/Rewrite/MetadataRewriter.h"
 #include "llvm/ADT/SmallVector.h"
-#include "llvm/Support/Error.h"
 
 namespace llvm {
 namespace bolt {

diff --git a/bolt/include/bolt/Rewrite/RewriteInstance.h b/bolt/include/bolt/Rewrite/RewriteInstance.h
@@ -17,7 +17,6 @@
 #include "bolt/Core/Linker.h"
 #include "bolt/Rewrite/MetadataManager.h"
 #include "bolt/Utils/NameResolver.h"
-#include "llvm/ADT/ArrayRef.h"
 #include "llvm/MC/StringTableBuilder.h"
 #include "llvm/Object/ELFObjectFile.h"
 #include "llvm/Object/ObjectFile.h"

diff --git a/bolt/include/bolt/RuntimeLibs/RuntimeLibrary.h b/bolt/include/bolt/RuntimeLibs/RuntimeLibrary.h
@@ -17,7 +17,6 @@
 
 #include "bolt/Core/Linker.h"
 #include "llvm/ADT/StringRef.h"
-#include <functional>
 #include <vector>
 
 namespace llvm {

diff --git a/bolt/include/bolt/Utils/NameShortener.h b/bolt/include/bolt/Utils/NameShortener.h
@@ -14,7 +14,6 @@
 #define BOLT_UTILS_NAME_SHORTENER_H
 
 #include "llvm/ADT/StringMap.h"
-#include "llvm/ADT/Twine.h"
 
 namespace llvm {
 namespace bolt {

diff --git a/bolt/lib/Core/BinaryContext.cpp b/bolt/lib/Core/BinaryContext.cpp
@@ -14,7 +14,6 @@
 #include "bolt/Core/BinaryEmitter.h"
 #include "bolt/Core/BinaryFunction.h"
 #include "bolt/Utils/CommandLineOpts.h"
-#include "bolt/Utils/NameResolver.h"
 #include "bolt/Utils/Utils.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/Twine.h"
@@ -39,7 +38,6 @@
 #include <algorithm>
 #include <functional>
 #include <iterator>
-#include <numeric>
 #include <unordered_set>
 
 using namespace llvm;
@@ -162,28 +160,30 @@ BinaryContext::~BinaryContext() {
 
 /// Create BinaryContext for a given architecture \p ArchName and
 /// triple \p TripleName.
-Expected<std::unique_ptr<BinaryContext>>
-BinaryContext::createBinaryContext(const ObjectFile *File, bool IsPIC,
-                                   std::unique_ptr<DWARFContext> DwCtx,
-                                   JournalingStreams Logger) {
+Expected<std::unique_ptr<BinaryContext>> BinaryContext::createBinaryContext(
+    Triple TheTriple, StringRef InputFileName, SubtargetFeatures *Features,
+    bool IsPIC, std::unique_ptr<DWARFContext> DwCtx, JournalingStreams Logger) {
   StringRef ArchName = "";
   std::string FeaturesStr = "";
-  switch (File->getArch()) {
+  switch (TheTriple.getArch()) {
   case llvm::Triple::x86_64:
+    if (Features)
+      return createFatalBOLTError(
+          "x86_64 target does not use SubtargetFeatures");
     ArchName = "x86-64";
     FeaturesStr = "+nopl";
     break;
   case llvm::Triple::aarch64:
+    if (Features)
+      return createFatalBOLTError(
+          "AArch64 target does not use SubtargetFeatures");
     ArchName = "aarch64";
     FeaturesStr = "+all";
     break;
   case llvm::Triple::riscv64: {
     ArchName = "riscv64";
-    Expected<SubtargetFeatures> Features = File->getFeatures();
-
-    if (auto E = Features.takeError())
-      return std::move(E);
-
+    if (!Features)
+      return createFatalBOLTError("RISCV target needs SubtargetFeatures");
     // We rely on relaxation for some transformations (e.g., promoting all calls
     // to PseudoCALL and then making JITLink relax them). Since the relax
     // feature is not stored in the object file, we manually enable it.
@@ -196,12 +196,11 @@ BinaryContext::createBinaryContext(const ObjectFile *File, bool IsPIC,
                              "BOLT-ERROR: Unrecognized machine in ELF file");
   }
 
-  auto TheTriple = std::make_unique<Triple>(File->makeTriple());
-  const std::string TripleName = TheTriple->str();
+  const std::string TripleName = TheTriple.str();
 
   std::string Error;
   const Target *TheTarget =
-      TargetRegistry::lookupTarget(std::string(ArchName), *TheTriple, Error);
+      TargetRegistry::lookupTarget(std::string(ArchName), TheTriple, Error);
   if (!TheTarget)
     return createStringError(make_error_code(std::errc::not_supported),
                              Twine("BOLT-ERROR: ", Error));
@@ -240,13 +239,13 @@ BinaryContext::createBinaryContext(const ObjectFile *File, bool IsPIC,
         Twine("BOLT-ERROR: no instruction info for target ", TripleName));
 
   std::unique_ptr<MCContext> Ctx(
-      new MCContext(*TheTriple, AsmInfo.get(), MRI.get(), STI.get()));
+      new MCContext(TheTriple, AsmInfo.get(), MRI.get(), STI.get()));
   std::unique_ptr<MCObjectFileInfo> MOFI(
       TheTarget->createMCObjectFileInfo(*Ctx, IsPIC));
   Ctx->setObjectFileInfo(MOFI.get());
   // We do not support X86 Large code model. Change this in the future.
   bool Large = false;
-  if (TheTriple->getArch() == llvm::Triple::aarch64)
+  if (TheTriple.getArch() == llvm::Triple::aarch64)
     Large = true;
   unsigned LSDAEncoding =
       Large ? dwarf::DW_EH_PE_absptr : dwarf::DW_EH_PE_udata4;
@@ -273,7 +272,7 @@ BinaryContext::createBinaryContext(const ObjectFile *File, bool IsPIC,
 
   int AsmPrinterVariant = AsmInfo->getAssemblerDialect();
   std::unique_ptr<MCInstPrinter> InstructionPrinter(
-      TheTarget->createMCInstPrinter(*TheTriple, AsmPrinterVariant, *AsmInfo,
+      TheTarget->createMCInstPrinter(TheTriple, AsmPrinterVariant, *AsmInfo,
                                      *MII, *MRI));
   if (!InstructionPrinter)
     return createStringError(
@@ -285,8 +284,8 @@ BinaryContext::createBinaryContext(const ObjectFile *File, bool IsPIC,
       TheTarget->createMCCodeEmitter(*MII, *Ctx));
 
   auto BC = std::make_unique<BinaryContext>(
-      std::move(Ctx), std::move(DwCtx), std::move(TheTriple), TheTarget,
-      std::string(TripleName), std::move(MCE), std::move(MOFI),
+      std::move(Ctx), std::move(DwCtx), std::make_unique<Triple>(TheTriple),
+      TheTarget, std::string(TripleName), std::move(MCE), std::move(MOFI),
       std::move(AsmInfo), std::move(MII), std::move(STI),
       std::move(InstructionPrinter), std::move(MIA), nullptr, std::move(MRI),
       std::move(DisAsm), Logger);
@@ -296,7 +295,7 @@ BinaryContext::createBinaryContext(const ObjectFile *File, bool IsPIC,
   BC->MAB = std::unique_ptr<MCAsmBackend>(
       BC->TheTarget->createMCAsmBackend(*BC->STI, *BC->MRI, MCTargetOptions()));
 
-  BC->setFilename(File->getFileName());
+  BC->setFilename(InputFileName);
 
   BC->HasFixedLoadAddress = !IsPIC;
 

diff --git a/bolt/lib/Core/BinaryFunction.cpp b/bolt/lib/Core/BinaryFunction.cpp
@@ -12,7 +12,6 @@
 
 #include "bolt/Core/BinaryFunction.h"
 #include "bolt/Core/BinaryBasicBlock.h"
-#include "bolt/Core/BinaryDomTree.h"
 #include "bolt/Core/DynoStats.h"
 #include "bolt/Core/HashUtilities.h"
 #include "bolt/Core/MCPlusBuilder.h"
@@ -35,6 +34,8 @@
 #include "llvm/Object/ObjectFile.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
+#include "llvm/Support/GenericDomTreeConstruction.h"
+#include "llvm/Support/GenericLoopInfoImpl.h"
 #include "llvm/Support/GraphWriter.h"
 #include "llvm/Support/LEB128.h"
 #include "llvm/Support/Regex.h"
@@ -3547,7 +3548,7 @@ MCSymbol *BinaryFunction::getSymbolForEntryID(uint64_t EntryID) {
   if (!isMultiEntry())
     return nullptr;
 
-  uint64_t NumEntries = 0;
+  uint64_t NumEntries = 1;
   if (hasCFG()) {
     for (BinaryBasicBlock *BB : BasicBlocks) {
       MCSymbol *EntrySymbol = getSecondaryEntryPointSymbol(*BB);
@@ -3580,7 +3581,7 @@ uint64_t BinaryFunction::getEntryIDForSymbol(const MCSymbol *Symbol) const {
       return 0;
 
   // Check all secondary entries available as either basic blocks or lables.
-  uint64_t NumEntries = 0;
+  uint64_t NumEntries = 1;
   for (const BinaryBasicBlock *BB : BasicBlocks) {
     MCSymbol *EntrySymbol = getSecondaryEntryPointSymbol(*BB);
     if (!EntrySymbol)
@@ -3589,7 +3590,7 @@ uint64_t BinaryFunction::getEntryIDForSymbol(const MCSymbol *Symbol) const {
       return NumEntries;
     ++NumEntries;
   }
-  NumEntries = 0;
+  NumEntries = 1;
   for (const std::pair<const uint32_t, MCSymbol *> &KV : Labels) {
     MCSymbol *EntrySymbol = getSecondaryEntryPointSymbol(KV.second);
     if (!EntrySymbol)
@@ -4076,12 +4077,17 @@ BinaryFunction::~BinaryFunction() {
     delete BB;
 }
 
+void BinaryFunction::constructDomTree() {
+  BDT.reset(new BinaryDominatorTree);
+  BDT->recalculate(*this);
+}
+
 void BinaryFunction::calculateLoopInfo() {
+  if (!hasDomTree())
+    constructDomTree();
   // Discover loops.
-  BinaryDominatorTree DomTree;
-  DomTree.recalculate(*this);
   BLI.reset(new BinaryLoopInfo());
-  BLI->analyze(DomTree);
+  BLI->analyze(getDomTree());
 
   // Traverse discovered loops and add depth and profile information.
   std::stack<BinaryLoop *> St;

diff --git a/bolt/lib/Core/DIEBuilder.cpp b/bolt/lib/Core/DIEBuilder.cpp
@@ -22,7 +22,6 @@
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ErrorHandling.h"
-#include "llvm/Support/Format.h"
 #include "llvm/Support/LEB128.h"
 
 #include <algorithm>

diff --git a/bolt/lib/Core/DebugData.cpp b/bolt/lib/Core/DebugData.cpp
@@ -13,7 +13,6 @@
 #include "bolt/Core/DebugData.h"
 #include "bolt/Core/BinaryContext.h"
 #include "bolt/Core/DIEBuilder.h"
-#include "bolt/Rewrite/RewriteInstance.h"
 #include "bolt/Utils/Utils.h"
 #include "llvm/BinaryFormat/Dwarf.h"
 #include "llvm/CodeGen/DIE.h"
@@ -23,7 +22,6 @@
 #include "llvm/MC/MCAssembler.h"
 #include "llvm/MC/MCContext.h"
 #include "llvm/MC/MCObjectStreamer.h"
-#include "llvm/Support/Allocator.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/EndianStream.h"
 #include "llvm/Support/LEB128.h"
@@ -32,7 +30,6 @@
 #include <cassert>
 #include <cstdint>
 #include <functional>
-#include <limits>
 #include <memory>
 #include <unordered_map>
 #include <vector>

diff --git a/bolt/lib/Core/FunctionLayout.cpp b/bolt/lib/Core/FunctionLayout.cpp
@@ -1,12 +1,17 @@
+//===- bolt/Core/FunctionLayout.cpp - Fragmented Function Layout -*- C++ -*-==//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
 #include "bolt/Core/FunctionLayout.h"
-#include "bolt/Core/BinaryFunction.h"
+#include "bolt/Core/BinaryBasicBlock.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/edit_distance.h"
 #include <algorithm>
-#include <cstddef>
-#include <functional>
 #include <iterator>
-#include <memory>
 
 using namespace llvm;
 using namespace bolt;

diff --git a/bolt/lib/Core/HashUtilities.cpp b/bolt/lib/Core/HashUtilities.cpp
@@ -12,7 +12,6 @@
 
 #include "bolt/Core/HashUtilities.h"
 #include "bolt/Core/BinaryContext.h"
-#include "bolt/Core/BinaryFunction.h"
 #include "llvm/MC/MCInstPrinter.h"
 
 namespace llvm {

diff --git a/bolt/lib/Core/MCPlusBuilder.cpp b/bolt/lib/Core/MCPlusBuilder.cpp
@@ -12,22 +12,30 @@
 
 #include "bolt/Core/MCPlusBuilder.h"
 #include "bolt/Core/MCPlus.h"
+#include "bolt/Utils/CommandLineOpts.h"
 #include "llvm/MC/MCContext.h"
 #include "llvm/MC/MCInst.h"
 #include "llvm/MC/MCInstrAnalysis.h"
 #include "llvm/MC/MCInstrDesc.h"
 #include "llvm/MC/MCInstrInfo.h"
 #include "llvm/MC/MCRegisterInfo.h"
+#include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include <cstdint>
-#include <queue>
 
 #define DEBUG_TYPE "mcplus"
 
 using namespace llvm;
 using namespace bolt;
 using namespace MCPlus;
 
+namespace opts {
+cl::opt<bool>
+    TerminalTrap("terminal-trap",
+                 cl::desc("Assume that execution stops at trap instruction"),
+                 cl::init(true), cl::Hidden, cl::cat(BoltCategory));
+}
+
 bool MCPlusBuilder::equals(const MCInst &A, const MCInst &B,
                            CompFuncTy Comp) const {
   if (A.getOpcode() != B.getOpcode())
@@ -121,6 +129,11 @@ bool MCPlusBuilder::equals(const MCTargetExpr &A, const MCTargetExpr &B,
   llvm_unreachable("target-specific expressions are unsupported");
 }
 
+bool MCPlusBuilder::isTerminator(const MCInst &Inst) const {
+  return Analysis->isTerminator(Inst) ||
+         (opts::TerminalTrap && Info->get(Inst.getOpcode()).isTrap());
+}
+
 void MCPlusBuilder::setTailCall(MCInst &Inst) const {
   assert(!hasAnnotation(Inst, MCAnnotation::kTailCall));
   setAnnotationOpValue(Inst, MCAnnotation::kTailCall, true);

diff --git a/bolt/lib/Core/Relocation.cpp b/bolt/lib/Core/Relocation.cpp
@@ -774,60 +774,95 @@ static bool isPCRelativeRISCV(uint64_t Type) {
 }
 
 bool Relocation::isSupported(uint64_t Type) {
-  if (Arch == Triple::aarch64)
+  switch (Arch) {
+  default:
+    return false;
+  case Triple::aarch64:
     return isSupportedAArch64(Type);
-  if (Arch == Triple::riscv64)
+  case Triple::riscv64:
     return isSupportedRISCV(Type);
-  return isSupportedX86(Type);
+  case Triple::x86_64:
+    return isSupportedX86(Type);
+  }
 }
 
 size_t Relocation::getSizeForType(uint64_t Type) {
-  if (Arch == Triple::aarch64)
+  switch (Arch) {
+  default:
+    llvm_unreachable("Unsupported architecture");
+  case Triple::aarch64:
     return getSizeForTypeAArch64(Type);
-  if (Arch == Triple::riscv64)
+  case Triple::riscv64:
     return getSizeForTypeRISCV(Type);
-  return getSizeForTypeX86(Type);
+  case Triple::x86_64:
+    return getSizeForTypeX86(Type);
+  }
 }
 
 bool Relocation::skipRelocationType(uint64_t Type) {
-  if (Arch == Triple::aarch64)
+  switch (Arch) {
+  default:
+    llvm_unreachable("Unsupported architecture");
+  case Triple::aarch64:
     return skipRelocationTypeAArch64(Type);
-  if (Arch == Triple::riscv64)
+  case Triple::riscv64:
     return skipRelocationTypeRISCV(Type);
-  return skipRelocationTypeX86(Type);
+  case Triple::x86_64:
+    return skipRelocationTypeX86(Type);
+  }
 }
 
 bool Relocation::skipRelocationProcess(uint64_t &Type, uint64_t Contents) {
-  if (Arch == Triple::aarch64)
+  switch (Arch) {
+  default:
+    llvm_unreachable("Unsupported architecture");
+  case Triple::aarch64:
     return skipRelocationProcessAArch64(Type, Contents);
-  if (Arch == Triple::riscv64)
-    skipRelocationProcessRISCV(Type, Contents);
-  return skipRelocationProcessX86(Type, Contents);
+  case Triple::riscv64:
+    return skipRelocationProcessRISCV(Type, Contents);
+  case Triple::x86_64:
+    return skipRelocationProcessX86(Type, Contents);
+  }
 }
 
 uint64_t Relocation::encodeValue(uint64_t Type, uint64_t Value, uint64_t PC) {
-  if (Arch == Triple::aarch64)
+  switch (Arch) {
+  default:
+    llvm_unreachable("Unsupported architecture");
+  case Triple::aarch64:
     return encodeValueAArch64(Type, Value, PC);
-  if (Arch == Triple::riscv64)
+  case Triple::riscv64:
     return encodeValueRISCV(Type, Value, PC);
-  return encodeValueX86(Type, Value, PC);
+  case Triple::x86_64:
+    return encodeValueX86(Type, Value, PC);
+  }
 }
 
 uint64_t Relocation::extractValue(uint64_t Type, uint64_t Contents,
                                   uint64_t PC) {
-  if (Arch == Triple::aarch64)
+  switch (Arch) {
+  default:
+    llvm_unreachable("Unsupported architecture");
+  case Triple::aarch64:
     return extractValueAArch64(Type, Contents, PC);
-  if (Arch == Triple::riscv64)
+  case Triple::riscv64:
     return extractValueRISCV(Type, Contents, PC);
-  return extractValueX86(Type, Contents, PC);
+  case Triple::x86_64:
+    return extractValueX86(Type, Contents, PC);
+  }
 }
 
 bool Relocation::isGOT(uint64_t Type) {
-  if (Arch == Triple::aarch64)
+  switch (Arch) {
+  default:
+    llvm_unreachable("Unsupported architecture");
+  case Triple::aarch64:
     return isGOTAArch64(Type);
-  if (Arch == Triple::riscv64)
+  case Triple::riscv64:
     return isGOTRISCV(Type);
-  return isGOTX86(Type);
+  case Triple::x86_64:
+    return isGOTX86(Type);
+  }
 }
 
 bool Relocation::isX86GOTPCRELX(uint64_t Type) {
@@ -845,27 +880,42 @@ bool Relocation::isX86GOTPC64(uint64_t Type) {
 bool Relocation::isNone(uint64_t Type) { return Type == getNone(); }
 
 bool Relocation::isRelative(uint64_t Type) {
-  if (Arch == Triple::aarch64)
+  switch (Arch) {
+  default:
+    llvm_unreachable("Unsupported architecture");
+  case Triple::aarch64:
     return Type == ELF::R_AARCH64_RELATIVE;
-  if (Arch == Triple::riscv64)
+  case Triple::riscv64:
     return Type == ELF::R_RISCV_RELATIVE;
-  return Type == ELF::R_X86_64_RELATIVE;
+  case Triple::x86_64:
+    return Type == ELF::R_X86_64_RELATIVE;
+  }
 }
 
 bool Relocation::isIRelative(uint64_t Type) {
-  if (Arch == Triple::aarch64)
+  switch (Arch) {
+  default:
+    llvm_unreachable("Unsupported architecture");
+  case Triple::aarch64:
     return Type == ELF::R_AARCH64_IRELATIVE;
-  if (Arch == Triple::riscv64)
+  case Triple::riscv64:
     llvm_unreachable("not implemented");
-  return Type == ELF::R_X86_64_IRELATIVE;
+  case Triple::x86_64:
+    return Type == ELF::R_X86_64_IRELATIVE;
+  }
 }
 
 bool Relocation::isTLS(uint64_t Type) {
-  if (Arch == Triple::aarch64)
+  switch (Arch) {
+  default:
+    llvm_unreachable("Unsupported architecture");
+  case Triple::aarch64:
     return isTLSAArch64(Type);
-  if (Arch == Triple::riscv64)
+  case Triple::riscv64:
     return isTLSRISCV(Type);
-  return isTLSX86(Type);
+  case Triple::x86_64:
+    return isTLSX86(Type);
+  }
 }
 
 bool Relocation::isInstructionReference(uint64_t Type) {
@@ -882,49 +932,81 @@ bool Relocation::isInstructionReference(uint64_t Type) {
 }
 
 uint64_t Relocation::getNone() {
-  if (Arch == Triple::aarch64)
+  switch (Arch) {
+  default:
+    llvm_unreachable("Unsupported architecture");
+  case Triple::aarch64:
     return ELF::R_AARCH64_NONE;
-  if (Arch == Triple::riscv64)
+  case Triple::riscv64:
     return ELF::R_RISCV_NONE;
-  return ELF::R_X86_64_NONE;
+  case Triple::x86_64:
+    return ELF::R_X86_64_NONE;
+  }
 }
 
 uint64_t Relocation::getPC32() {
-  if (Arch == Triple::aarch64)
+  switch (Arch) {
+  default:
+    llvm_unreachable("Unsupported architecture");
+  case Triple::aarch64:
     return ELF::R_AARCH64_PREL32;
-  if (Arch == Triple::riscv64)
+  case Triple::riscv64:
     return ELF::R_RISCV_32_PCREL;
-  return ELF::R_X86_64_PC32;
+  case Triple::x86_64:
+    return ELF::R_X86_64_PC32;
+  }
 }
 
 uint64_t Relocation::getPC64() {
-  if (Arch == Triple::aarch64)
+  switch (Arch) {
+  default:
+    llvm_unreachable("Unsupported architecture");
+  case Triple::aarch64:
     return ELF::R_AARCH64_PREL64;
-  if (Arch == Triple::riscv64)
+  case Triple::riscv64:
     llvm_unreachable("not implemented");
-  return ELF::R_X86_64_PC64;
+  case Triple::x86_64:
+    return ELF::R_X86_64_PC64;
+  }
 }
 
 bool Relocation::isPCRelative(uint64_t Type) {
-  if (Arch == Triple::aarch64)
+  switch (Arch) {
+  default:
+    llvm_unreachable("Unsupported architecture");
+  case Triple::aarch64:
     return isPCRelativeAArch64(Type);
-  if (Arch == Triple::riscv64)
+  case Triple::riscv64:
     return isPCRelativeRISCV(Type);
-  return isPCRelativeX86(Type);
+  case Triple::x86_64:
+    return isPCRelativeX86(Type);
+  }
 }
 
 uint64_t Relocation::getAbs64() {
-  if (Arch == Triple::aarch64)
+  switch (Arch) {
+  default:
+    llvm_unreachable("Unsupported architecture");
+  case Triple::aarch64:
     return ELF::R_AARCH64_ABS64;
-  if (Arch == Triple::riscv64)
+  case Triple::riscv64:
     return ELF::R_RISCV_64;
-  return ELF::R_X86_64_64;
+  case Triple::x86_64:
+    return ELF::R_X86_64_64;
+  }
 }
 
 uint64_t Relocation::getRelative() {
-  if (Arch == Triple::aarch64)
+  switch (Arch) {
+  default:
+    llvm_unreachable("Unsupported architecture");
+  case Triple::aarch64:
     return ELF::R_AARCH64_RELATIVE;
-  return ELF::R_X86_64_RELATIVE;
+  case Triple::riscv64:
+    llvm_unreachable("not implemented");
+  case Triple::x86_64:
+    return ELF::R_X86_64_RELATIVE;
+  }
 }
 
 size_t Relocation::emit(MCStreamer *Streamer) const {
@@ -991,9 +1073,16 @@ void Relocation::print(raw_ostream &OS) const {
   static const char *AArch64RelocNames[] = {
 #include "llvm/BinaryFormat/ELFRelocs/AArch64.def"
   };
-  if (Arch == Triple::aarch64)
+  switch (Arch) {
+  default:
+    OS << "RType:" << Twine::utohexstr(Type);
+    break;
+
+  case Triple::aarch64:
     OS << AArch64RelocNames[Type];
-  else if (Arch == Triple::riscv64) {
+    break;
+
+  case Triple::riscv64:
     // RISC-V relocations are not sequentially numbered so we cannot use an
     // array
     switch (Type) {
@@ -1006,8 +1095,12 @@ void Relocation::print(raw_ostream &OS) const {
     break;
 #include "llvm/BinaryFormat/ELFRelocs/RISCV.def"
     }
-  } else
+    break;
+
+  case Triple::x86_64:
     OS << X86RelocNames[Type];
+    break;
+  }
   OS << ", 0x" << Twine::utohexstr(Offset);
   if (Symbol) {
     OS << ", " << Symbol->getName();

diff --git a/bolt/lib/Passes/CMOVConversion.cpp b/bolt/lib/Passes/CMOVConversion.cpp
@@ -17,7 +17,6 @@
 #include "llvm/ADT/PostOrderIterator.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/ErrorHandling.h"
-#include <numeric>
 
 #define DEBUG_TYPE "cmov"
 

diff --git a/bolt/lib/Passes/FixRISCVCallsPass.cpp b/bolt/lib/Passes/FixRISCVCallsPass.cpp
@@ -1,3 +1,11 @@
+//===- bolt/Passes/FixRISCVCallsPass.cpp ------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
 #include "bolt/Passes/FixRISCVCallsPass.h"
 #include "bolt/Core/ParallelUtilities.h"
 

diff --git a/bolt/lib/Passes/FixRelaxationPass.cpp b/bolt/lib/Passes/FixRelaxationPass.cpp
@@ -1,3 +1,11 @@
+//===- bolt/Passes/FixRelaxationPass.cpp ------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
 #include "bolt/Passes/FixRelaxationPass.h"
 #include "bolt/Core/ParallelUtilities.h"
 

diff --git a/bolt/lib/Passes/FrameOptimizer.cpp b/bolt/lib/Passes/FrameOptimizer.cpp
@@ -20,7 +20,6 @@
 #include "bolt/Utils/CommandLineOpts.h"
 #include "llvm/Support/Timer.h"
 #include <deque>
-#include <unordered_map>
 
 #define DEBUG_TYPE "fop"
 

diff --git a/bolt/lib/Passes/Hugify.cpp b/bolt/lib/Passes/Hugify.cpp
@@ -7,7 +7,6 @@
 //===----------------------------------------------------------------------===//
 
 #include "bolt/Passes/Hugify.h"
-#include "llvm/Support/CommandLine.h"
 
 #define DEBUG_TYPE "bolt-hugify"
 

diff --git a/bolt/lib/Passes/Inliner.cpp b/bolt/lib/Passes/Inliner.cpp
@@ -27,7 +27,6 @@
 #include "bolt/Passes/Inliner.h"
 #include "bolt/Core/MCPlus.h"
 #include "llvm/Support/CommandLine.h"
-#include <map>
 
 #define DEBUG_TYPE "bolt-inliner"
 

diff --git a/bolt/lib/Passes/ShrinkWrapping.cpp b/bolt/lib/Passes/ShrinkWrapping.cpp
@@ -11,7 +11,6 @@
 //===----------------------------------------------------------------------===//
 
 #include "bolt/Passes/ShrinkWrapping.h"
-#include "bolt/Core/MCPlus.h"
 #include "bolt/Passes/DataflowInfoManager.h"
 #include "bolt/Passes/MCF.h"
 #include "bolt/Utils/CommandLineOpts.h"

diff --git a/bolt/lib/Passes/SplitFunctions.cpp b/bolt/lib/Passes/SplitFunctions.cpp
@@ -17,7 +17,6 @@
 #include "bolt/Core/ParallelUtilities.h"
 #include "bolt/Utils/CommandLineOpts.h"
 #include "llvm/ADT/STLExtras.h"
-#include "llvm/ADT/Sequence.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/iterator_range.h"
 #include "llvm/Support/CommandLine.h"

diff --git a/bolt/lib/Passes/TailDuplication.cpp b/bolt/lib/Passes/TailDuplication.cpp
@@ -13,9 +13,9 @@
 #include "bolt/Passes/TailDuplication.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/MC/MCRegisterInfo.h"
-#include <queue>
 
 #include <numeric>
+#include <queue>
 
 #define DEBUG_TYPE "taildup"
 

diff --git a/bolt/lib/Passes/ValidateInternalCalls.cpp b/bolt/lib/Passes/ValidateInternalCalls.cpp
@@ -14,7 +14,6 @@
 #include "bolt/Core/BinaryBasicBlock.h"
 #include "bolt/Passes/DataflowInfoManager.h"
 #include "bolt/Passes/FrameAnalysis.h"
-#include "llvm/ADT/SmallVector.h"
 #include "llvm/MC/MCInstPrinter.h"
 #include <optional>
 #include <queue>

diff --git a/bolt/lib/Profile/BoltAddressTranslation.cpp b/bolt/lib/Profile/BoltAddressTranslation.cpp
@@ -100,7 +100,7 @@ void BoltAddressTranslation::write(const BinaryContext &BC, raw_ostream &OS) {
     LLVM_DEBUG(dbgs() << "Function name: " << Function.getPrintName() << "\n");
     LLVM_DEBUG(dbgs() << " Address reference: 0x"
                       << Twine::utohexstr(Function.getOutputAddress()) << "\n");
-    LLVM_DEBUG(dbgs() << formatv(" Hash: {0:x}\n", getBFHash(OutputAddress)));
+    LLVM_DEBUG(dbgs() << formatv(" Hash: {0:x}\n", getBFHash(InputAddress)));
     LLVM_DEBUG(dbgs() << " Secondary Entry Points: " << NumSecondaryEntryPoints
                       << '\n');
 
@@ -197,8 +197,9 @@ void BoltAddressTranslation::writeMaps(std::map<uint64_t, MapTy> &Maps,
             ? SecondaryEntryPointsMap[Address].size()
             : 0;
     if (Cold) {
-      size_t HotIndex =
-          std::distance(ColdPartSource.begin(), ColdPartSource.find(Address));
+      auto HotEntryIt = Maps.find(ColdPartSource[Address]);
+      assert(HotEntryIt != Maps.end());
+      size_t HotIndex = std::distance(Maps.begin(), HotEntryIt);
       encodeULEB128(HotIndex - PrevIndex, OS);
       PrevIndex = HotIndex;
     } else {
@@ -207,7 +208,7 @@ void BoltAddressTranslation::writeMaps(std::map<uint64_t, MapTy> &Maps,
       LLVM_DEBUG(dbgs() << "Hash: " << formatv("{0:x}\n", BFHash));
       OS.write(reinterpret_cast<char *>(&BFHash), 8);
       // Number of basic blocks
-      size_t NumBasicBlocks = getBBHashMap(HotInputAddress).getNumBasicBlocks();
+      size_t NumBasicBlocks = NumBasicBlocksMap[HotInputAddress];
       LLVM_DEBUG(dbgs() << "Basic blocks: " << NumBasicBlocks << '\n');
       encodeULEB128(NumBasicBlocks, OS);
       // Secondary entry points
@@ -425,8 +426,9 @@ void BoltAddressTranslation::dump(raw_ostream &OS) {
   for (const auto &MapEntry : Maps) {
     const uint64_t Address = MapEntry.first;
     const uint64_t HotAddress = fetchParentAddress(Address);
+    const bool IsHotFunction = HotAddress == 0;
     OS << "Function Address: 0x" << Twine::utohexstr(Address);
-    if (HotAddress == 0)
+    if (IsHotFunction)
       OS << formatv(", hash: {0:x}", getBFHash(Address));
     OS << "\n";
     OS << "BB mappings:\n";
@@ -443,6 +445,8 @@ void BoltAddressTranslation::dump(raw_ostream &OS) {
         OS << formatv(" hash: {0:x}", BBHashMap.getBBHash(Val));
       OS << "\n";
     }
+    if (IsHotFunction)
+      OS << "NumBlocks: " << NumBasicBlocksMap[Address] << '\n';
     if (SecondaryEntryPointsMap.count(Address)) {
       const std::vector<uint32_t> &SecondaryEntryPoints =
           SecondaryEntryPointsMap[Address];
@@ -574,27 +578,52 @@ void BoltAddressTranslation::saveMetadata(BinaryContext &BC) {
     // Set BF/BB metadata
     for (const BinaryBasicBlock &BB : BF)
       BBHashMap.addEntry(BB.getInputOffset(), BB.getIndex(), BB.getHash());
+    NumBasicBlocksMap.emplace(BF.getAddress(), BF.size());
   }
 }
 
-std::unordered_map<uint32_t, std::vector<uint32_t>>
-BoltAddressTranslation::getBFBranches(uint64_t OutputAddress) const {
-  std::unordered_map<uint32_t, std::vector<uint32_t>> Branches;
-  auto FuncIt = Maps.find(OutputAddress);
-  assert(FuncIt != Maps.end());
-  std::vector<uint32_t> InputOffsets;
-  for (const auto &KV : FuncIt->second)
-    InputOffsets.emplace_back(KV.second);
-  // Sort with LSB BRANCHENTRY bit.
-  llvm::sort(InputOffsets);
-  uint32_t BBOffset{0};
-  for (uint32_t InOffset : InputOffsets) {
-    if (InOffset & BRANCHENTRY)
-      Branches[BBOffset].push_back(InOffset >> 1);
-    else
-      BBOffset = InOffset >> 1;
-  }
-  return Branches;
+unsigned
+BoltAddressTranslation::getSecondaryEntryPointId(uint64_t Address,
+                                                 uint32_t Offset) const {
+  auto FunctionIt = SecondaryEntryPointsMap.find(Address);
+  if (FunctionIt == SecondaryEntryPointsMap.end())
+    return 0;
+  const std::vector<uint32_t> &Offsets = FunctionIt->second;
+  auto OffsetIt = std::find(Offsets.begin(), Offsets.end(), Offset);
+  if (OffsetIt == Offsets.end())
+    return 0;
+  // Adding one here because main entry point is not stored in BAT, and
+  // enumeration for secondary entry points starts with 1.
+  return OffsetIt - Offsets.begin() + 1;
+}
+
+std::pair<const BinaryFunction *, unsigned>
+BoltAddressTranslation::translateSymbol(const BinaryContext &BC,
+                                        const MCSymbol &Symbol,
+                                        uint32_t Offset) const {
+  // The symbol could be a secondary entry in a cold fragment.
+  uint64_t SymbolValue = cantFail(errorOrToExpected(BC.getSymbolValue(Symbol)));
+
+  const BinaryFunction *Callee = BC.getFunctionForSymbol(&Symbol);
+  assert(Callee);
+
+  // Containing function, not necessarily the same as symbol value.
+  const uint64_t CalleeAddress = Callee->getAddress();
+  const uint32_t OutputOffset = SymbolValue - CalleeAddress;
+
+  const uint64_t ParentAddress = fetchParentAddress(CalleeAddress);
+  const uint64_t HotAddress = ParentAddress ? ParentAddress : CalleeAddress;
+
+  const BinaryFunction *ParentBF = BC.getBinaryFunctionAtAddress(HotAddress);
+
+  const uint32_t InputOffset =
+      translate(CalleeAddress, OutputOffset, /*IsBranchSrc*/ false) + Offset;
+
+  unsigned SecondaryEntryId{0};
+  if (InputOffset)
+    SecondaryEntryId = getSecondaryEntryPointId(HotAddress, InputOffset);
+
+  return std::pair(ParentBF, SecondaryEntryId);
 }
 
 } // namespace bolt

diff --git a/bolt/lib/Profile/CMakeLists.txt b/bolt/lib/Profile/CMakeLists.txt
@@ -3,7 +3,6 @@ add_llvm_library(LLVMBOLTProfile
   DataAggregator.cpp
   DataReader.cpp
   Heatmap.cpp
-  ProfileReaderBase.cpp
   StaleProfileMatching.cpp
   YAMLProfileReader.cpp
   YAMLProfileWriter.cpp

diff --git a/bolt/lib/Profile/DataAggregator.cpp b/bolt/lib/Profile/DataAggregator.cpp
@@ -604,8 +604,6 @@ Error DataAggregator::readProfile(BinaryContext &BC) {
     // BAT YAML is handled by DataAggregator since normal YAML output requires
     // CFG which is not available in BAT mode.
     if (usesBAT()) {
-      // Postprocess split function profile for BAT
-      fixupBATProfile(BC);
       if (opts::ProfileFormat == opts::ProfileFormatKind::PF_YAML)
         if (std::error_code EC = writeBATYAML(BC, opts::OutputFilename))
           report_error("cannot create output data file", EC);
@@ -664,18 +662,19 @@ DataAggregator::getBinaryFunctionContainingAddress(uint64_t Address) const {
                                                 /*UseMaxSize=*/true);
 }
 
-StringRef DataAggregator::getLocationName(BinaryFunction &Func,
-                                          uint64_t Count) {
+BinaryFunction *
+DataAggregator::getBATParentFunction(const BinaryFunction &Func) const {
+  if (BAT)
+    if (const uint64_t HotAddr = BAT->fetchParentAddress(Func.getAddress()))
+      return getBinaryFunctionContainingAddress(HotAddr);
+  return nullptr;
+}
+
+StringRef DataAggregator::getLocationName(const BinaryFunction &Func) const {
   if (!BAT)
     return Func.getOneName();
 
   const BinaryFunction *OrigFunc = &Func;
-  if (const uint64_t HotAddr = BAT->fetchParentAddress(Func.getAddress())) {
-    NumColdSamples += Count;
-    BinaryFunction *HotFunc = getBinaryFunctionContainingAddress(HotAddr);
-    if (HotFunc)
-      OrigFunc = HotFunc;
-  }
   // If it is a local function, prefer the name containing the file name where
   // the local function was declared
   for (StringRef AlternativeName : OrigFunc->getNames()) {
@@ -690,12 +689,17 @@ StringRef DataAggregator::getLocationName(BinaryFunction &Func,
   return OrigFunc->getOneName();
 }
 
-bool DataAggregator::doSample(BinaryFunction &Func, uint64_t Address,
+bool DataAggregator::doSample(BinaryFunction &OrigFunc, uint64_t Address,
                               uint64_t Count) {
+  BinaryFunction *ParentFunc = getBATParentFunction(OrigFunc);
+  BinaryFunction &Func = ParentFunc ? *ParentFunc : OrigFunc;
+  if (ParentFunc)
+    NumColdSamples += Count;
+
   auto I = NamesToSamples.find(Func.getOneName());
   if (I == NamesToSamples.end()) {
     bool Success;
-    StringRef LocName = getLocationName(Func, Count);
+    StringRef LocName = getLocationName(Func);
     std::tie(I, Success) = NamesToSamples.insert(
         std::make_pair(Func.getOneName(),
                        FuncSampleData(LocName, FuncSampleData::ContainerTy())));
@@ -715,22 +719,12 @@ bool DataAggregator::doIntraBranch(BinaryFunction &Func, uint64_t From,
   FuncBranchData *AggrData = getBranchData(Func);
   if (!AggrData) {
     AggrData = &NamesToBranches[Func.getOneName()];
-    AggrData->Name = getLocationName(Func, Count);
+    AggrData->Name = getLocationName(Func);
     setBranchData(Func, AggrData);
   }
 
-  From -= Func.getAddress();
-  To -= Func.getAddress();
   LLVM_DEBUG(dbgs() << "BOLT-DEBUG: bumpBranchCount: "
                     << formatv("{0} @ {1:x} -> {0} @ {2:x}\n", Func, From, To));
-  if (BAT) {
-    From = BAT->translate(Func.getAddress(), From, /*IsBranchSrc=*/true);
-    To = BAT->translate(Func.getAddress(), To, /*IsBranchSrc=*/false);
-    LLVM_DEBUG(
-        dbgs() << "BOLT-DEBUG: BAT translation on bumpBranchCount: "
-               << formatv("{0} @ {1:x} -> {0} @ {2:x}\n", Func, From, To));
-  }
-
   AggrData->bumpBranchCount(From, To, Count, Mispreds);
   return true;
 }
@@ -744,30 +738,24 @@ bool DataAggregator::doInterBranch(BinaryFunction *FromFunc,
   StringRef SrcFunc;
   StringRef DstFunc;
   if (FromFunc) {
-    SrcFunc = getLocationName(*FromFunc, Count);
+    SrcFunc = getLocationName(*FromFunc);
     FromAggrData = getBranchData(*FromFunc);
     if (!FromAggrData) {
       FromAggrData = &NamesToBranches[FromFunc->getOneName()];
       FromAggrData->Name = SrcFunc;
       setBranchData(*FromFunc, FromAggrData);
     }
-    From -= FromFunc->getAddress();
-    if (BAT)
-      From = BAT->translate(FromFunc->getAddress(), From, /*IsBranchSrc=*/true);
 
     recordExit(*FromFunc, From, Mispreds, Count);
   }
   if (ToFunc) {
-    DstFunc = getLocationName(*ToFunc, 0);
+    DstFunc = getLocationName(*ToFunc);
     ToAggrData = getBranchData(*ToFunc);
     if (!ToAggrData) {
       ToAggrData = &NamesToBranches[ToFunc->getOneName()];
       ToAggrData->Name = DstFunc;
       setBranchData(*ToFunc, ToAggrData);
     }
-    To -= ToFunc->getAddress();
-    if (BAT)
-      To = BAT->translate(ToFunc->getAddress(), To, /*IsBranchSrc=*/false);
 
     recordEntry(*ToFunc, To, Mispreds, Count);
   }
@@ -783,15 +771,32 @@ bool DataAggregator::doInterBranch(BinaryFunction *FromFunc,
 
 bool DataAggregator::doBranch(uint64_t From, uint64_t To, uint64_t Count,
                               uint64_t Mispreds) {
-  BinaryFunction *FromFunc = getBinaryFunctionContainingAddress(From);
-  BinaryFunction *ToFunc = getBinaryFunctionContainingAddress(To);
+  auto handleAddress = [&](uint64_t &Addr, bool IsFrom) -> BinaryFunction * {
+    if (BinaryFunction *Func = getBinaryFunctionContainingAddress(Addr)) {
+      Addr -= Func->getAddress();
+
+      if (BAT)
+        Addr = BAT->translate(Func->getAddress(), Addr, IsFrom);
+
+      if (BinaryFunction *ParentFunc = getBATParentFunction(*Func)) {
+        Func = ParentFunc;
+        if (IsFrom)
+          NumColdSamples += Count;
+      }
+
+      return Func;
+    }
+    return nullptr;
+  };
+
+  BinaryFunction *FromFunc = handleAddress(From, /*IsFrom=*/true);
+  BinaryFunction *ToFunc = handleAddress(To, /*IsFrom=*/false);
   if (!FromFunc && !ToFunc)
     return false;
 
   // Treat recursive control transfers as inter-branches.
-  if (FromFunc == ToFunc && (To != ToFunc->getAddress())) {
-    recordBranch(*FromFunc, From - FromFunc->getAddress(),
-                 To - FromFunc->getAddress(), Count, Mispreds);
+  if (FromFunc == ToFunc && To != 0) {
+    recordBranch(*FromFunc, From, To, Count, Mispreds);
     return doIntraBranch(*FromFunc, From, To, Count, Mispreds);
   }
 
@@ -842,9 +847,14 @@ bool DataAggregator::doTrace(const LBREntry &First, const LBREntry &Second,
                     << FromFunc->getPrintName() << ":"
                     << Twine::utohexstr(First.To) << " to "
                     << Twine::utohexstr(Second.From) << ".\n");
-  for (const std::pair<uint64_t, uint64_t> &Pair : *FTs)
-    doIntraBranch(*FromFunc, Pair.first + FromFunc->getAddress(),
-                  Pair.second + FromFunc->getAddress(), Count, false);
+  BinaryFunction *ParentFunc = getBATParentFunction(*FromFunc);
+  for (auto [From, To] : *FTs) {
+    if (BAT) {
+      From = BAT->translate(FromFunc->getAddress(), From, /*IsBranchSrc=*/true);
+      To = BAT->translate(FromFunc->getAddress(), To, /*IsBranchSrc=*/false);
+    }
+    doIntraBranch(ParentFunc ? *ParentFunc : *FromFunc, From, To, Count, false);
+  }
 
   return true;
 }
@@ -2273,29 +2283,6 @@ DataAggregator::writeAggregatedFile(StringRef OutputFilename) const {
   return std::error_code();
 }
 
-void DataAggregator::fixupBATProfile(BinaryContext &BC) {
-  for (auto &[FuncName, Branches] : NamesToBranches) {
-    BinaryData *BD = BC.getBinaryDataByName(FuncName);
-    assert(BD);
-    uint64_t FuncAddress = BD->getAddress();
-    if (!BAT->isBATFunction(FuncAddress))
-      continue;
-    // Filter out cold fragments
-    if (!BD->getSectionName().equals(BC.getMainCodeSectionName()))
-      continue;
-    // Convert inter-branches between hot and cold fragments into
-    // intra-branches.
-    for (auto &[OffsetFrom, CallToMap] : Branches.InterIndex) {
-      for (auto &[CallToLoc, CallToIdx] : CallToMap) {
-        if (CallToLoc.Name != FuncName)
-          continue;
-        Branches.IntraIndex[OffsetFrom][CallToLoc.Offset] = CallToIdx;
-        Branches.InterIndex[OffsetFrom].erase(CallToLoc);
-      }
-    }
-  }
-}
-
 std::error_code DataAggregator::writeBATYAML(BinaryContext &BC,
                                              StringRef OutputFilename) const {
   std::error_code EC;
@@ -2333,7 +2320,7 @@ std::error_code DataAggregator::writeBATYAML(BinaryContext &BC,
       if (BAT->isBATFunction(Function.getAddress()))
         continue;
       BP.Functions.emplace_back(
-          YAMLProfileWriter::convert(Function, /*UseDFS=*/false));
+          YAMLProfileWriter::convert(Function, /*UseDFS=*/false, BAT));
     }
 
     for (const auto &KV : NamesToBranches) {
@@ -2345,9 +2332,6 @@ std::error_code DataAggregator::writeBATYAML(BinaryContext &BC,
       uint64_t FuncAddress = BD->getAddress();
       if (!BAT->isBATFunction(FuncAddress))
         continue;
-      // Filter out cold fragments
-      if (!BD->getSectionName().equals(BC.getMainCodeSectionName()))
-        continue;
       BinaryFunction *BF = BC.getBinaryFunctionAtAddress(FuncAddress);
       assert(BF);
       YamlBF.Name = FuncName.str();
@@ -2357,87 +2341,68 @@ std::error_code DataAggregator::writeBATYAML(BinaryContext &BC,
       YamlBF.NumBasicBlocks = BAT->getNumBasicBlocks(FuncAddress);
       const BoltAddressTranslation::BBHashMapTy &BlockMap =
           BAT->getBBHashMap(FuncAddress);
+      YamlBF.Blocks.resize(YamlBF.NumBasicBlocks);
 
-      auto addSuccProfile = [&](yaml::bolt::BinaryBasicBlockProfile &YamlBB,
-                                uint64_t SuccOffset, unsigned SuccDataIdx) {
+      for (auto &&[Idx, YamlBB] : llvm::enumerate(YamlBF.Blocks))
+        YamlBB.Index = Idx;
+
+      for (auto BI = BlockMap.begin(), BE = BlockMap.end(); BI != BE; ++BI)
+        YamlBF.Blocks[BI->second.getBBIndex()].Hash = BI->second.getBBHash();
+
+      auto getSuccessorInfo = [&](uint32_t SuccOffset, unsigned SuccDataIdx) {
         const llvm::bolt::BranchInfo &BI = Branches.Data.at(SuccDataIdx);
         yaml::bolt::SuccessorInfo SI;
         SI.Index = BlockMap.getBBIndex(SuccOffset);
         SI.Count = BI.Branches;
         SI.Mispreds = BI.Mispreds;
-        YamlBB.Successors.emplace_back(SI);
+        return SI;
       };
 
-      std::unordered_map<uint32_t, std::vector<uint32_t>> BFBranches =
-          BAT->getBFBranches(FuncAddress);
-
-      auto addCallsProfile = [&](yaml::bolt::BinaryBasicBlockProfile &YamlBB,
-                                 uint64_t Offset) {
-        // Iterate over BRANCHENTRY records in the current block
-        for (uint32_t BranchOffset : BFBranches[Offset]) {
-          if (!Branches.InterIndex.contains(BranchOffset))
-            continue;
-          for (const auto &[CallToLoc, CallToIdx] :
-               Branches.InterIndex.at(BranchOffset)) {
-            const llvm::bolt::BranchInfo &BI = Branches.Data.at(CallToIdx);
-            yaml::bolt::CallSiteInfo YamlCSI;
-            YamlCSI.DestId = 0; // designated for unknown functions
-            YamlCSI.EntryDiscriminator = 0;
-            YamlCSI.Count = BI.Branches;
-            YamlCSI.Mispreds = BI.Mispreds;
-            YamlCSI.Offset = BranchOffset - Offset;
-            BinaryData *CallTargetBD = BC.getBinaryDataByName(CallToLoc.Name);
-            if (!CallTargetBD) {
-              YamlBB.CallSites.emplace_back(YamlCSI);
-              continue;
-            }
-            uint64_t CallTargetAddress = CallTargetBD->getAddress();
-            BinaryFunction *CallTargetBF =
-                BC.getBinaryFunctionAtAddress(CallTargetAddress);
-            if (!CallTargetBF) {
-              YamlBB.CallSites.emplace_back(YamlCSI);
-              continue;
-            }
-            // Calls between hot and cold fragments must be handled in
-            // fixupBATProfile.
-            assert(CallTargetBF != BF && "invalid CallTargetBF");
-            YamlCSI.DestId = CallTargetBF->getFunctionNumber();
-            if (CallToLoc.Offset) {
-              if (BAT->isBATFunction(CallTargetAddress)) {
-                LLVM_DEBUG(dbgs() << "BOLT-DEBUG: Unsupported secondary "
-                                     "entry point in BAT function "
-                                  << CallToLoc.Name << '\n');
-              } else if (const BinaryBasicBlock *CallTargetBB =
-                             CallTargetBF->getBasicBlockAtOffset(
-                                 CallToLoc.Offset)) {
-                // Only record true call information, ignoring returns (normally
-                // won't have a target basic block) and jumps to the landing
-                // pads (not an entry point).
-                if (CallTargetBB->isEntryPoint()) {
-                  YamlCSI.EntryDiscriminator =
-                      CallTargetBF->getEntryIDForSymbol(
-                          CallTargetBB->getLabel());
-                }
-              }
-            }
-            YamlBB.CallSites.emplace_back(YamlCSI);
-          }
-        }
+      auto getCallSiteInfo = [&](Location CallToLoc, unsigned CallToIdx,
+                                 uint32_t Offset) {
+        const llvm::bolt::BranchInfo &BI = Branches.Data.at(CallToIdx);
+        yaml::bolt::CallSiteInfo CSI;
+        CSI.DestId = 0; // designated for unknown functions
+        CSI.EntryDiscriminator = 0;
+        CSI.Count = BI.Branches;
+        CSI.Mispreds = BI.Mispreds;
+        CSI.Offset = Offset;
+        if (BinaryData *BD = BC.getBinaryDataByName(CallToLoc.Name))
+          YAMLProfileWriter::setCSIDestination(BC, CSI, BD->getSymbol(), BAT,
+                                               CallToLoc.Offset);
+        return CSI;
       };
 
       for (const auto &[FromOffset, SuccKV] : Branches.IntraIndex) {
-        yaml::bolt::BinaryBasicBlockProfile YamlBB;
         if (!BlockMap.isInputBlock(FromOffset))
           continue;
-        YamlBB.Index = BlockMap.getBBIndex(FromOffset);
-        YamlBB.Hash = BlockMap.getBBHash(FromOffset);
+        const unsigned Index = BlockMap.getBBIndex(FromOffset);
+        yaml::bolt::BinaryBasicBlockProfile &YamlBB = YamlBF.Blocks[Index];
         for (const auto &[SuccOffset, SuccDataIdx] : SuccKV)
-          addSuccProfile(YamlBB, SuccOffset, SuccDataIdx);
-        addCallsProfile(YamlBB, FromOffset);
-        if (YamlBB.ExecCount || !YamlBB.Successors.empty() ||
-            !YamlBB.CallSites.empty())
-          YamlBF.Blocks.emplace_back(YamlBB);
+          if (BlockMap.isInputBlock(SuccOffset))
+            YamlBB.Successors.emplace_back(
+                getSuccessorInfo(SuccOffset, SuccDataIdx));
+      }
+      for (const auto &[FromOffset, CallTo] : Branches.InterIndex) {
+        auto BlockIt = BlockMap.upper_bound(FromOffset);
+        --BlockIt;
+        const unsigned BlockOffset = BlockIt->first;
+        const unsigned BlockIndex = BlockIt->second.getBBIndex();
+        yaml::bolt::BinaryBasicBlockProfile &YamlBB = YamlBF.Blocks[BlockIndex];
+        const uint32_t Offset = FromOffset - BlockOffset;
+        for (const auto &[CallToLoc, CallToIdx] : CallTo)
+          YamlBB.CallSites.emplace_back(
+              getCallSiteInfo(CallToLoc, CallToIdx, Offset));
+        llvm::sort(YamlBB.CallSites, [](yaml::bolt::CallSiteInfo &A,
+                                        yaml::bolt::CallSiteInfo &B) {
+          return A.Offset < B.Offset;
+        });
       }
+      // Drop blocks without a hash, won't be useful for stale matching.
+      llvm::erase_if(YamlBF.Blocks,
+                     [](const yaml::bolt::BinaryBasicBlockProfile &YamlBB) {
+                       return YamlBB.Hash == (yaml::Hex64)0;
+                     });
       BP.Functions.emplace_back(YamlBF);
     }
   }

diff --git a/bolt/lib/Profile/DataReader.cpp b/bolt/lib/Profile/DataReader.cpp
@@ -18,7 +18,6 @@
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/Errc.h"
-#include <map>
 
 #undef  DEBUG_TYPE
 #define DEBUG_TYPE "bolt-prof"

diff --git a/bolt/lib/Profile/Heatmap.cpp b/bolt/lib/Profile/Heatmap.cpp
@@ -10,7 +10,6 @@
 #include "bolt/Utils/CommandLineOpts.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/Twine.h"
-#include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/Format.h"

diff --git a/bolt/lib/Profile/YAMLProfileWriter.cpp b/bolt/lib/Profile/YAMLProfileWriter.cpp
@@ -9,6 +9,7 @@
 #include "bolt/Profile/YAMLProfileWriter.h"
 #include "bolt/Core/BinaryBasicBlock.h"
 #include "bolt/Core/BinaryFunction.h"
+#include "bolt/Profile/BoltAddressTranslation.h"
 #include "bolt/Profile/ProfileReaderBase.h"
 #include "bolt/Rewrite/RewriteInstance.h"
 #include "llvm/Support/CommandLine.h"
@@ -25,8 +26,30 @@ extern llvm::cl::opt<bool> ProfileUseDFS;
 namespace llvm {
 namespace bolt {
 
+const BinaryFunction *YAMLProfileWriter::setCSIDestination(
+    const BinaryContext &BC, yaml::bolt::CallSiteInfo &CSI,
+    const MCSymbol *Symbol, const BoltAddressTranslation *BAT,
+    uint32_t Offset) {
+  CSI.DestId = 0; // designated for unknown functions
+  CSI.EntryDiscriminator = 0;
+
+  if (Symbol) {
+    uint64_t EntryID = 0;
+    if (const BinaryFunction *Callee =
+            BC.getFunctionForSymbol(Symbol, &EntryID)) {
+      if (BAT && BAT->isBATFunction(Callee->getAddress()))
+        std::tie(Callee, EntryID) = BAT->translateSymbol(BC, *Symbol, Offset);
+      CSI.DestId = Callee->getFunctionNumber();
+      CSI.EntryDiscriminator = EntryID;
+      return Callee;
+    }
+  }
+  return nullptr;
+}
+
 yaml::bolt::BinaryFunctionProfile
-YAMLProfileWriter::convert(const BinaryFunction &BF, bool UseDFS) {
+YAMLProfileWriter::convert(const BinaryFunction &BF, bool UseDFS,
+                           const BoltAddressTranslation *BAT) {
   yaml::bolt::BinaryFunctionProfile YamlBF;
   const BinaryContext &BC = BF.getBinaryContext();
 
@@ -79,31 +102,21 @@ YAMLProfileWriter::convert(const BinaryFunction &BF, bool UseDFS) {
           continue;
         for (const IndirectCallProfile &CSP : ICSP.get()) {
           StringRef TargetName = "";
-          CSI.DestId = 0; // designated for unknown functions
-          CSI.EntryDiscriminator = 0;
-          if (CSP.Symbol) {
-            const BinaryFunction *Callee = BC.getFunctionForSymbol(CSP.Symbol);
-            if (Callee) {
-              CSI.DestId = Callee->getFunctionNumber();
-              TargetName = Callee->getOneName();
-            }
-          }
+          const BinaryFunction *Callee =
+              setCSIDestination(BC, CSI, CSP.Symbol, BAT);
+          if (Callee)
+            TargetName = Callee->getOneName();
           CSI.Count = CSP.Count;
           CSI.Mispreds = CSP.Mispreds;
           CSTargets.emplace_back(TargetName, CSI);
         }
       } else { // direct call or a tail call
-        uint64_t EntryID = 0;
-        CSI.DestId = 0;
         StringRef TargetName = "";
         const MCSymbol *CalleeSymbol = BC.MIB->getTargetSymbol(Instr);
         const BinaryFunction *const Callee =
-            BC.getFunctionForSymbol(CalleeSymbol, &EntryID);
-        if (Callee) {
-          CSI.DestId = Callee->getFunctionNumber();
-          CSI.EntryDiscriminator = EntryID;
+            setCSIDestination(BC, CSI, CalleeSymbol, BAT);
+        if (Callee)
           TargetName = Callee->getOneName();
-        }
 
         auto getAnnotationWithDefault = [&](const MCInst &Inst, StringRef Ann) {
           return BC.MIB->getAnnotationWithDefault(Instr, Ann, 0ull);

diff --git a/bolt/lib/Rewrite/BinaryPassManager.cpp b/bolt/lib/Rewrite/BinaryPassManager.cpp
@@ -377,8 +377,9 @@ Error BinaryFunctionPassManager::runAllPasses(BinaryContext &BC) {
 
   Manager.registerPass(std::make_unique<NormalizeCFG>(PrintNormalized));
 
-  Manager.registerPass(std::make_unique<StripRepRet>(NeverPrint),
-                       opts::StripRepRet);
+  if (BC.isX86())
+    Manager.registerPass(std::make_unique<StripRepRet>(NeverPrint),
+                         opts::StripRepRet);
 
   Manager.registerPass(std::make_unique<IdenticalCodeFolding>(PrintICF),
                        opts::ICF);

diff --git a/bolt/lib/Rewrite/DWARFRewriter.cpp b/bolt/lib/Rewrite/DWARFRewriter.cpp
@@ -14,7 +14,6 @@
 #include "bolt/Core/DynoStats.h"
 #include "bolt/Core/ParallelUtilities.h"
 #include "bolt/Rewrite/RewriteInstance.h"
-#include "bolt/Utils/Utils.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
@@ -1685,7 +1684,7 @@ namespace {
 std::unique_ptr<BinaryContext>
 createDwarfOnlyBC(const object::ObjectFile &File) {
   return cantFail(BinaryContext::createBinaryContext(
-      &File, false,
+      File.makeTriple(), File.getFileName(), nullptr, false,
       DWARFContext::create(File, DWARFContext::ProcessDebugRelocations::Ignore,
                            nullptr, "", WithColor::defaultErrorHandler,
                            WithColor::defaultWarningHandler),

diff --git a/bolt/lib/Rewrite/JITLinkLinker.cpp b/bolt/lib/Rewrite/JITLinkLinker.cpp
@@ -5,9 +5,11 @@
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
+
 #include "bolt/Rewrite/JITLinkLinker.h"
+#include "bolt/Core/BinaryContext.h"
 #include "bolt/Core/BinaryData.h"
-#include "bolt/Rewrite/RewriteInstance.h"
+#include "bolt/Core/BinarySection.h"
 #include "llvm/ExecutionEngine/JITLink/ELF_riscv.h"
 #include "llvm/ExecutionEngine/JITLink/JITLink.h"
 #include "llvm/ExecutionEngine/Orc/Shared/ExecutorAddress.h"

diff --git a/bolt/lib/Rewrite/LinuxKernelRewriter.cpp b/bolt/lib/Rewrite/LinuxKernelRewriter.cpp
@@ -212,6 +212,11 @@ class LinuxKernelRewriter final : public MetadataRewriter {
   /// Size of bug_entry struct.
   static constexpr size_t BUG_TABLE_ENTRY_SIZE = 12;
 
+  /// List of bug entries per function.
+  using FunctionBugListType =
+      DenseMap<BinaryFunction *, SmallVector<uint32_t, 2>>;
+  FunctionBugListType FunctionBugList;
+
   /// .pci_fixup section.
   ErrorOr<BinarySection &> PCIFixupSection = std::errc::bad_address;
   static constexpr size_t PCI_FIXUP_ENTRY_SIZE = 16;
@@ -254,7 +259,9 @@ class LinuxKernelRewriter final : public MetadataRewriter {
   Error readParaInstructions();
   Error rewriteParaInstructions();
 
+  /// __bug_table section handling.
   Error readBugTable();
+  Error rewriteBugTable();
 
   /// Do no process functions containing instruction annotated with
   /// \p Annotation.
@@ -339,6 +346,9 @@ class LinuxKernelRewriter final : public MetadataRewriter {
     if (Error E = rewriteStaticKeysJumpTable())
       return E;
 
+    if (Error E = rewriteBugTable())
+      return E;
+
     return Error::success();
   }
 
@@ -1164,15 +1174,17 @@ Error LinuxKernelRewriter::rewriteParaInstructions() {
 }
 
 /// Process __bug_table section.
-/// This section contains information useful for kernel debugging.
+/// This section contains information useful for kernel debugging, mostly
+/// utilized by WARN()/WARN_ON() macros and deprecated BUG()/BUG_ON().
+///
 /// Each entry in the section is a struct bug_entry that contains a pointer to
 /// the ud2 instruction corresponding to the bug, corresponding file name (both
 /// pointers use PC relative offset addressing), line number, and flags.
 /// The definition of the struct bug_entry can be found in
-/// `include/asm-generic/bug.h`
-///
-/// NB: find_bug() uses linear search to match an address to an entry in the bug
-///     table. Hence there is no need to sort entries when rewriting the table.
+/// `include/asm-generic/bug.h`. The first entry in the struct is an instruction
+/// address encoded as a PC-relative offset. In theory, it could be an absolute
+/// address if CONFIG_GENERIC_BUG_RELATIVE_POINTERS is not set, but in practice
+/// the kernel code relies on it being a relative offset on x86-64.
 Error LinuxKernelRewriter::readBugTable() {
   BugTableSection = BC.getUniqueSectionByName("__bug_table");
   if (!BugTableSection)
@@ -1215,6 +1227,8 @@ Error LinuxKernelRewriter::readBugTable() {
                                  " referenced by bug table entry %d",
                                  InstAddress, EntryID);
       BC.MIB->addAnnotation(*Inst, "BugEntry", EntryID);
+
+      FunctionBugList[BF].push_back(EntryID);
     }
   }
 
@@ -1223,6 +1237,52 @@ Error LinuxKernelRewriter::readBugTable() {
   return Error::success();
 }
 
+/// find_bug() uses linear search to match an address to an entry in the bug
+/// table. Hence, there is no need to sort entries when rewriting the table.
+/// When we need to erase an entry, we set its instruction address to zero.
+Error LinuxKernelRewriter::rewriteBugTable() {
+  if (!BugTableSection)
+    return Error::success();
+
+  for (BinaryFunction &BF : llvm::make_second_range(BC.getBinaryFunctions())) {
+    if (!BC.shouldEmit(BF))
+      continue;
+
+    if (!FunctionBugList.count(&BF))
+      continue;
+
+    // Bugs that will be emitted for this function.
+    DenseSet<uint32_t> EmittedIDs;
+    for (BinaryBasicBlock &BB : BF) {
+      for (MCInst &Inst : BB) {
+        if (!BC.MIB->hasAnnotation(Inst, "BugEntry"))
+          continue;
+        const uint32_t ID = BC.MIB->getAnnotationAs<uint32_t>(Inst, "BugEntry");
+        EmittedIDs.insert(ID);
+
+        // Create a relocation entry for this bug entry.
+        MCSymbol *Label =
+            BC.MIB->getOrCreateInstLabel(Inst, "__BUG_", BC.Ctx.get());
+        const uint64_t EntryOffset = (ID - 1) * BUG_TABLE_ENTRY_SIZE;
+        BugTableSection->addRelocation(EntryOffset, Label, ELF::R_X86_64_PC32,
+                                       /*Addend*/ 0);
+      }
+    }
+
+    // Clear bug entries that were not emitted for this function, e.g. as a
+    // result of DCE, but setting their instruction address to zero.
+    for (const uint32_t ID : FunctionBugList[&BF]) {
+      if (!EmittedIDs.count(ID)) {
+        const uint64_t EntryOffset = (ID - 1) * BUG_TABLE_ENTRY_SIZE;
+        BugTableSection->addRelocation(EntryOffset, nullptr, ELF::R_X86_64_PC32,
+                                       /*Addend*/ 0);
+      }
+    }
+  }
+
+  return Error::success();
+}
+
 /// The kernel can replace certain instruction sequences depending on hardware
 /// it is running on and features specified during boot time. The information
 /// about alternative instruction sequences is stored in .altinstructions

diff --git a/bolt/lib/Rewrite/MachORewriteInstance.cpp b/bolt/lib/Rewrite/MachORewriteInstance.cpp
@@ -18,6 +18,7 @@
 #include "bolt/Rewrite/BinaryPassManager.h"
 #include "bolt/Rewrite/ExecutableFileMemoryManager.h"
 #include "bolt/Rewrite/JITLinkLinker.h"
+#include "bolt/Rewrite/RewriteInstance.h"
 #include "bolt/RuntimeLibs/InstrumentationRuntimeLibrary.h"
 #include "bolt/Utils/Utils.h"
 #include "llvm/MC/MCObjectStreamer.h"
@@ -54,37 +55,6 @@ extern cl::opt<unsigned> Verbosity;
 namespace llvm {
 namespace bolt {
 
-extern MCPlusBuilder *createX86MCPlusBuilder(const MCInstrAnalysis *,
-                                             const MCInstrInfo *,
-                                             const MCRegisterInfo *,
-                                             const MCSubtargetInfo *);
-extern MCPlusBuilder *createAArch64MCPlusBuilder(const MCInstrAnalysis *,
-                                                 const MCInstrInfo *,
-                                                 const MCRegisterInfo *,
-                                                 const MCSubtargetInfo *);
-
-namespace {
-
-MCPlusBuilder *createMCPlusBuilder(const Triple::ArchType Arch,
-                                   const MCInstrAnalysis *Analysis,
-                                   const MCInstrInfo *Info,
-                                   const MCRegisterInfo *RegInfo,
-                                   const MCSubtargetInfo *STI) {
-#ifdef X86_AVAILABLE
-  if (Arch == Triple::x86_64)
-    return createX86MCPlusBuilder(Analysis, Info, RegInfo, STI);
-#endif
-
-#ifdef AARCH64_AVAILABLE
-  if (Arch == Triple::aarch64)
-    return createAArch64MCPlusBuilder(Analysis, Info, RegInfo, STI);
-#endif
-
-  llvm_unreachable("architecture unsupported by MCPlusBuilder");
-}
-
-} // anonymous namespace
-
 #define DEBUG_TYPE "bolt"
 
 Expected<std::unique_ptr<MachORewriteInstance>>
@@ -103,7 +73,8 @@ MachORewriteInstance::MachORewriteInstance(object::MachOObjectFile *InputFile,
     : InputFile(InputFile), ToolPath(ToolPath) {
   ErrorAsOutParameter EAO(&Err);
   auto BCOrErr = BinaryContext::createBinaryContext(
-      InputFile, /* IsPIC */ true, DWARFContext::create(*InputFile),
+      InputFile->makeTriple(), InputFile->getFileName(), nullptr,
+      /* IsPIC */ true, DWARFContext::create(*InputFile),
       {llvm::outs(), llvm::errs()});
   if (Error E = BCOrErr.takeError()) {
     Err = std::move(E);

diff --git a/bolt/lib/Rewrite/RewriteInstance.cpp b/bolt/lib/Rewrite/RewriteInstance.cpp
@@ -84,6 +84,7 @@ extern cl::opt<JumpTableSupportLevel> JumpTables;
 extern cl::opt<bool> KeepNops;
 extern cl::list<std::string> ReorderData;
 extern cl::opt<bolt::ReorderFunctions::ReorderType> ReorderFunctions;
+extern cl::opt<bool> TerminalTrap;
 extern cl::opt<bool> TimeBuild;
 
 cl::opt<bool> AllowStripped("allow-stripped",
@@ -267,6 +268,10 @@ namespace bolt {
 
 extern const char *BoltRevision;
 
+// Weird location for createMCPlusBuilder, but this is here to avoid a
+// cyclic dependency of libCore (its natural place) and libTarget. libRewrite
+// can depend on libTarget, but not libCore. Since libRewrite is the only
+// user of this function, we define it here.
 MCPlusBuilder *createMCPlusBuilder(const Triple::ArchType Arch,
                                    const MCInstrAnalysis *Analysis,
                                    const MCInstrInfo *Info,
@@ -344,8 +349,21 @@ RewriteInstance::RewriteInstance(ELFObjectFileBase *File, const int Argc,
   Stderr.SetUnbuffered();
   LLVM_DEBUG(dbgs().SetUnbuffered());
 
+  // Read RISCV subtarget features from input file
+  std::unique_ptr<SubtargetFeatures> Features;
+  Triple TheTriple = File->makeTriple();
+  if (TheTriple.getArch() == llvm::Triple::riscv64) {
+    Expected<SubtargetFeatures> FeaturesOrErr = File->getFeatures();
+    if (auto E = FeaturesOrErr.takeError()) {
+      Err = std::move(E);
+      return;
+    } else {
+      Features.reset(new SubtargetFeatures(*FeaturesOrErr));
+    }
+  }
+
   auto BCOrErr = BinaryContext::createBinaryContext(
-      File, IsPIC,
+      TheTriple, File->getFileName(), Features.get(), IsPIC,
       DWARFContext::create(*File, DWARFContext::ProcessDebugRelocations::Ignore,
                            nullptr, opts::DWPPathName,
                            WithColor::defaultErrorHandler,
@@ -538,7 +556,7 @@ Error RewriteInstance::discoverStorage() {
     if (Error E = SectionNameOrErr.takeError())
       return E;
     StringRef SectionName = SectionNameOrErr.get();
-    if (SectionName == ".text") {
+    if (SectionName == BC->getMainCodeSectionName()) {
       BC->OldTextSectionAddress = Section.getAddress();
       BC->OldTextSectionSize = Section.getSize();
 
@@ -1846,7 +1864,8 @@ Error RewriteInstance::readSpecialSections() {
                   "Use -update-debug-sections to keep it.\n";
   }
 
-  HasTextRelocations = (bool)BC->getUniqueSectionByName(".rela.text");
+  HasTextRelocations = (bool)BC->getUniqueSectionByName(
+      ".rela" + std::string(BC->getMainCodeSectionName()));
   HasSymbolTable = (bool)BC->getUniqueSectionByName(".symtab");
   EHFrameSection = BC->getUniqueSectionByName(".eh_frame");
   BuildIDSection = BC->getUniqueSectionByName(".note.gnu.build-id");
@@ -2033,8 +2052,14 @@ void RewriteInstance::adjustCommandLineOptions() {
   if (opts::Lite)
     BC->outs() << "BOLT-INFO: enabling lite mode\n";
 
-  if (BC->IsLinuxKernel && !opts::KeepNops.getNumOccurrences())
-    opts::KeepNops = true;
+  if (BC->IsLinuxKernel) {
+    if (!opts::KeepNops.getNumOccurrences())
+      opts::KeepNops = true;
+
+    // Linux kernel may resume execution after a trap instruction in some cases.
+    if (!opts::TerminalTrap.getNumOccurrences())
+      opts::TerminalTrap = false;
+  }
 }
 
 namespace {
@@ -3417,7 +3442,8 @@ void RewriteInstance::emitAndLink() {
   ErrorOr<BinarySection &> TextSection =
       BC->getUniqueSectionByName(BC->getMainCodeSectionName());
   if (BC->HasRelocations && TextSection)
-    BC->renameSection(*TextSection, getOrgSecPrefix() + ".text");
+    BC->renameSection(*TextSection,
+                      getOrgSecPrefix() + BC->getMainCodeSectionName());
 
   //////////////////////////////////////////////////////////////////////////////
   // Assign addresses to new sections.
@@ -4393,6 +4419,7 @@ void RewriteInstance::patchELFSectionHeaderTable(ELFObjectFile<ELFT> *File) {
   raw_fd_ostream &OS = Out->os();
   const ELFFile<ELFT> &Obj = File->getELFFile();
 
+  // Mapping from old section indices to new ones
   std::vector<uint32_t> NewSectionIndex;
   std::vector<ELFShdrTy> OutputSections =
       getOutputSections(File, NewSectionIndex);
@@ -4410,10 +4437,8 @@ void RewriteInstance::patchELFSectionHeaderTable(ELFObjectFile<ELFT> *File) {
   // Write all section header entries while patching section references.
   for (ELFShdrTy &Section : OutputSections) {
     Section.sh_link = NewSectionIndex[Section.sh_link];
-    if (Section.sh_type == ELF::SHT_REL || Section.sh_type == ELF::SHT_RELA) {
-      if (Section.sh_info)
-        Section.sh_info = NewSectionIndex[Section.sh_info];
-    }
+    if (Section.sh_type == ELF::SHT_REL || Section.sh_type == ELF::SHT_RELA)
+      Section.sh_info = NewSectionIndex[Section.sh_info];
     OS.write(reinterpret_cast<const char *>(&Section), sizeof(Section));
   }
 

diff --git a/bolt/lib/RuntimeLibs/HugifyRuntimeLibrary.cpp b/bolt/lib/RuntimeLibs/HugifyRuntimeLibrary.cpp
@@ -11,10 +11,9 @@
 //===----------------------------------------------------------------------===//
 
 #include "bolt/RuntimeLibs/HugifyRuntimeLibrary.h"
-#include "bolt/Core/BinaryFunction.h"
+#include "bolt/Core/BinaryContext.h"
 #include "bolt/Core/Linker.h"
 #include "llvm/MC/MCStreamer.h"
-#include "llvm/Support/Alignment.h"
 #include "llvm/Support/CommandLine.h"
 
 using namespace llvm;

diff --git a/bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp b/bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp
@@ -16,10 +16,7 @@
 #include "llvm/BinaryFormat/ELF.h"
 #include "llvm/MC/MCInst.h"
 #include "llvm/MC/MCSubtargetInfo.h"
-#include "llvm/Support/Debug.h"
 #include "llvm/Support/ErrorHandling.h"
-#include "llvm/Support/Format.h"
-#include "llvm/Support/raw_ostream.h"
 
 #define DEBUG_TYPE "mcplus"
 

diff --git a/bolt/lib/Target/X86/X86MCPlusBuilder.cpp b/bolt/lib/Target/X86/X86MCPlusBuilder.cpp
@@ -211,13 +211,6 @@ class X86MCPlusBuilder : public MCPlusBuilder {
     return false;
   }
 
-  // FIXME: For compatibility with old LLVM only!
-  bool isTerminator(const MCInst &Inst) const override {
-    unsigned Opcode = Inst.getOpcode();
-    return Info->get(Opcode).isTerminator() || X86::isUD1(Opcode) ||
-           X86::isUD2(Opcode);
-  }
-
   bool isIndirectCall(const MCInst &Inst) const override {
     return isCall(Inst) &&
            ((getMemoryOperandNo(Inst) != -1) || Inst.getOperand(0).isReg());

diff --git a/bolt/test/X86/bolt-address-translation-yaml.test b/bolt/test/X86/bolt-address-translation-yaml.test
@@ -36,9 +36,14 @@ YAML-BAT-CHECK-NEXT:   - bid:   0
 YAML-BAT-CHECK-NEXT:     insns: 26
 YAML-BAT-CHECK-NEXT:     hash:  0xA900AE79CFD40000
 YAML-BAT-CHECK-NEXT:     succ:  [ { bid: 3, cnt: 0 }, { bid: 1, cnt: 0 } ]
+# Calls from no-BAT to BAT function
+YAML-BAT-CHECK:        - bid:   28
+YAML-BAT-CHECK-NEXT:     insns: 13
+YAML-BAT-CHECK-NEXT:     hash:  0xB2F04C1F25F00400
+YAML-BAT-CHECK-NEXT:     calls: [ { off: 0x21, fid: [[#SOLVECUBIC:]], cnt: 25 }, { off: 0x2D, fid: [[#]], cnt: 9 } ]
 # Function covered by BAT with calls
 YAML-BAT-CHECK:      - name:    SolveCubic
-YAML-BAT-CHECK-NEXT:   fid:     [[#]]
+YAML-BAT-CHECK-NEXT:   fid:     [[#SOLVECUBIC]]
 YAML-BAT-CHECK-NEXT:   hash:    0x6AF7E61EA3966722
 YAML-BAT-CHECK-NEXT:   exec:    25
 YAML-BAT-CHECK-NEXT:   nblocks: 15

diff --git a/bolt/test/X86/linux-bug-table.s b/bolt/test/X86/linux-bug-table.s
@@ -1,14 +1,21 @@
 # REQUIRES: system-linux
 
-## Check that BOLT correctly parses the Linux kernel __bug_table section.
+## Check that BOLT correctly parses and updates the Linux kernel __bug_table
+## section.
 
 # RUN: llvm-mc -filetype=obj -triple x86_64-unknown-unknown %s -o %t.o
 # RUN: %clang %cflags -nostdlib %t.o -o %t.exe \
 # RUN:   -Wl,--image-base=0xffffffff80000000,--no-dynamic-linker,--no-eh-frame-hdr,--no-pie
 
 ## Verify bug entry bindings to instructions.
 
-# RUN: llvm-bolt %t.exe --print-normalized -o %t.out | FileCheck %s
+# RUN: llvm-bolt %t.exe --print-normalized --print-only=_start -o %t.out \
+# RUN:   --eliminate-unreachable=1 --bolt-info=0 | FileCheck %s
+
+## Verify bug entry bindings again after unreachable code elimination.
+
+# RUN: llvm-bolt %t.out -o %t.out.1 --print-only=_start --print-normalized \
+# RUN:   |& FileCheck --check-prefix=CHECK-REOPT %s
 
 # CHECK:      BOLT-INFO: Linux kernel binary detected
 # CHECK:      BOLT-INFO: parsed 2 bug table entries
@@ -17,18 +24,26 @@
   .globl _start
   .type _start, %function
 _start:
-# CHECK: Binary Function "_start"
-  nop
+  jmp .L1
 .L0:
   ud2
 # CHECK:      ud2
 # CHECK-SAME: BugEntry: 1
-  nop
 .L1:
   ud2
 # CHECK:      ud2
 # CHECK-SAME: BugEntry: 2
+
+## Only the second entry should remain after the first pass.
+
+# CHECK-REOPT: ud2
+# CHECK-REOPT-SAME: BugEntry: 2
+
   ret
+## The return instruction is reachable only via preceding ud2. Test that it is
+## treated as a reachable instruction in the Linux kernel mode.
+
+# CHECK-REOPT-NEXT: ret
   .size _start, .-_start
 
 

diff --git a/bolt/test/X86/patch-entries.test b/bolt/test/X86/patch-entries.test
@@ -7,4 +7,25 @@ REQUIRES: system-linux
 
 RUN: %clang %cflags -no-pie -g %p/Inputs/patch-entries.c -fuse-ld=lld -o %t.exe \
 RUN:   -Wl,-q -I%p/../Inputs
-RUN: llvm-bolt -relocs %t.exe -o %t.out --update-debug-sections --force-patch
+RUN: llvm-bolt -relocs %t.exe -o %t.out --update-debug-sections --force-patch \
+RUN:   --enable-bat
+
+# Check that patched functions can be disassembled (override FDE from the
+# original function)
+# PREAGG: B X:0 #foo.org.0# 1 0
+RUN: link_fdata %s %t.out %t.preagg PREAGG
+RUN: perf2bolt %t.out -p %t.preagg --pa -o %t.yaml --profile-format=yaml \
+RUN:   -print-disasm -print-only=foo.org.0/1 2>&1 | FileCheck %s
+CHECK-NOT: BOLT-WARNING: sizes differ for function foo.org.0/1
+CHECK: Binary Function "foo.org.0/1(*2)" after disassembly {
+
+# Check the expected eh_frame contents
+RUN: llvm-nm --print-size %t.out > %t.foo
+RUN: llvm-objdump %t.out --dwarf=frames >> %t.foo
+RUN: FileCheck %s --input-file %t.foo --check-prefix=CHECK-FOO
+CHECK-FOO: 0000000000[[#%x,FOO:]] [[#%x,OPTSIZE:]] t foo
+CHECK-FOO: 0000000000[[#%x,ORG:]] [[#%x,ORGSIZE:]] t foo.org.0
+# patched FDE comes first
+CHECK-FOO: FDE {{.*}} pc=00[[#%x,ORG]]...00[[#%x,ORG+ORGSIZE]]
+# original FDE comes second
+CHECK-FOO: FDE {{.*}} pc=00[[#%x,ORG]]...00[[#%x,ORG+OPTSIZE]]
diff --git a/bolt/test/X86/yaml-secondary-entry-discriminator.s b/bolt/test/X86/yaml-secondary-entry-discriminator.s
@@ -0,0 +1,145 @@
+## This reproduces a bug with BOLT setting incorrect discriminator for
+## secondary entry points in YAML profile.
+
+# 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: llvm-strip --strip-unneeded %t.o
+# RUN: %clang %cflags %t.o -o %t.exe -Wl,-q -nostdlib
+# RUN: llvm-bolt %t.exe -o %t.out --data %t.fdata -w %t.yaml --print-profile \
+# RUN:   --print-only=main | FileCheck %s --check-prefix=CHECK-CFG
+# RUN: FileCheck %s -input-file %t.yaml
+# CHECK:      - name:    main
+# CHECK-NEXT:   fid:     2
+# CHECK-NEXT:   hash:    {{.*}}
+# CHECK-NEXT:   exec:    0
+# CHECK-NEXT:   nblocks: 4
+# CHECK-NEXT:   blocks:
+# CHECK:          - bid:   1
+# CHECK-NEXT:       insns: 1
+# CHECK-NEXT:       hash:  {{.*}}
+# CHECK-NEXT:       calls: [ { off: 0x0, fid: 1, disc: 1, cnt: 1 } ]
+# CHECK:          - bid:   2
+# CHECK-NEXT:       insns: 5
+# CHECK-NEXT:       hash:  {{.*}}
+# CHECK-NEXT:       calls: [ { off: 0x0, fid: 1, disc: 1, cnt: 1, mis: 1 } ]
+
+## Make sure that the profile is attached correctly
+# RUN: llvm-bolt %t.exe -o %t.out --data %t.yaml --print-profile \
+# RUN:   --print-only=main | FileCheck %s --check-prefix=CHECK-CFG
+
+# CHECK-CFG: Binary Function "main" after attaching profile {
+# CHECK-CFG:      callq secondary_entry # Offset: [[#]] # Count: 1
+# CHECK-CFG:      callq *%rax # Offset: [[#]] # CallProfile: 1 (1 misses) :
+# CHECK-CFG-NEXT:     { secondary_entry: 1 (1 misses) }
+
+## YAML BAT test of calling BAT secondary entry from non-BAT function
+## Now force-split func and skip main (making it call secondary entries)
+# RUN: llvm-bolt %t.exe -o %t.bat --data %t.fdata --funcs=func \
+# RUN:   --split-functions --split-strategy=all --split-all-cold --enable-bat
+
+## Prepare pre-aggregated profile using %t.bat
+# RUN: link_fdata %s %t.bat %t.preagg PREAGG
+## Strip labels used for pre-aggregated profile
+# RUN: llvm-strip -NLcall -NLindcall %t.bat
+
+## Convert pre-aggregated profile using BAT
+# RUN: perf2bolt %t.bat -p %t.preagg --pa -o %t.bat.fdata -w %t.bat.yaml
+
+## Convert BAT fdata into YAML
+# RUN: llvm-bolt %t.exe -data %t.bat.fdata -w %t.bat.fdata-yaml -o /dev/null
+
+## Check fdata YAML - make sure that a direct call has discriminator field
+# RUN: FileCheck %s --input-file %t.bat.fdata-yaml -check-prefix CHECK-BAT-YAML
+
+## Check BAT YAML - make sure that a direct call has discriminator field
+# RUN: FileCheck %s --input-file %t.bat.yaml --check-prefix CHECK-BAT-YAML
+
+## YAML BAT test of calling BAT secondary entry from BAT function
+# RUN: llvm-bolt %t.exe -o %t.bat2 --data %t.fdata --funcs=main,func \
+# RUN:   --split-functions --split-strategy=all --split-all-cold --enable-bat
+
+## Prepare pre-aggregated profile using %t.bat
+# RUN: link_fdata %s %t.bat2 %t.preagg2 PREAGG2
+
+## Strip labels used for pre-aggregated profile
+# RUN: llvm-strip -NLcall -NLindcall %t.bat2
+
+## Convert pre-aggregated profile using BAT
+# RUN: perf2bolt %t.bat2 -p %t.preagg2 --pa -o %t.bat2.fdata -w %t.bat2.yaml
+
+## Convert BAT fdata into YAML
+# RUN: llvm-bolt %t.exe -data %t.bat2.fdata -w %t.bat2.fdata-yaml -o /dev/null
+
+## Check fdata YAML - make sure that a direct call has discriminator field
+# RUN: FileCheck %s --input-file %t.bat2.fdata-yaml -check-prefix CHECK-BAT-YAML
+
+## Check BAT YAML - make sure that a direct call has discriminator field
+# RUN: FileCheck %s --input-file %t.bat2.yaml --check-prefix CHECK-BAT-YAML
+
+# CHECK-BAT-YAML:      - name:    main
+# CHECK-BAT-YAML-NEXT:   fid:     [[#]]
+# CHECK-BAT-YAML-NEXT:   hash:    0xADF270D550151185
+# CHECK-BAT-YAML-NEXT:   exec:    0
+# CHECK-BAT-YAML-NEXT:   nblocks: 4
+# CHECK-BAT-YAML-NEXT:   blocks:
+# CHECK-BAT-YAML:          - bid:   1
+# CHECK-BAT-YAML-NEXT:       insns: [[#]]
+# CHECK-BAT-YAML-NEXT:       hash:  0x36A303CBA4360018
+# CHECK-BAT-YAML-NEXT:       calls: [ { off: 0x0, fid: [[#]], disc: 1, cnt: 1
+
+.globl func
+.type	func, @function
+func:
+# FDATA: 0 [unknown] 0 1 func 0 1 0
+# PREAGG: B X:0 #func# 1 1
+# PREAGG2: B X:0 #func# 1 1
+  .cfi_startproc
+  pushq   %rbp
+  movq    %rsp, %rbp
+## Placeholder code to make splitting profitable
+.rept 5
+  testq   %rax, %rax
+.endr
+.globl secondary_entry
+secondary_entry:
+## Placeholder code to make splitting profitable
+.rept 5
+  testq   %rax, %rax
+.endr
+  popq    %rbp
+  retq
+  nopl    (%rax)
+  .cfi_endproc
+  .size	func, .-func
+
+.globl main
+.type	main, @function
+main:
+  .cfi_startproc
+  pushq   %rbp
+  movq    %rsp, %rbp
+  subq    $16, %rsp
+  movl    $0, -4(%rbp)
+  testq   %rax, %rax
+  jne     Lindcall
+.globl Lcall
+Lcall:
+  call    secondary_entry
+# FDATA: 1 main #Lcall# 1 secondary_entry 0 1 1
+# PREAGG: B #Lcall# #secondary_entry# 1 1
+# PREAGG2: B #main.cold.0# #func.cold.0# 1 1
+.globl Lindcall
+Lindcall:
+  callq   *%rax
+# FDATA: 1 main #Lindcall# 1 secondary_entry 0 1 1
+# PREAGG: B #Lindcall# #secondary_entry# 1 1
+# PREAGG2: B #main.cold.1# #func.cold.0# 1 1
+  xorl    %eax, %eax
+  addq    $16, %rsp
+  popq    %rbp
+  retq
+## For relocations against .text
+  call exit
+  .cfi_endproc
+  .size	main, .-main
diff --git a/bolt/tools/bat-dump/bat-dump.cpp b/bolt/tools/bat-dump/bat-dump.cpp
@@ -1,9 +1,16 @@
+//===- bolt/tools/bat-dump/bat-dump.cpp - BAT dumper utility --------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
 #include "bolt/Profile/BoltAddressTranslation.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/Twine.h"
-#include "llvm/ADT/iterator_range.h"
 #include "llvm/Object/Binary.h"
 #include "llvm/Object/ELFObjectFile.h"
 #include "llvm/Object/Error.h"
@@ -18,7 +25,6 @@
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/Program.h"
 #include "llvm/Support/raw_ostream.h"
-#include <algorithm>
 #include <assert.h>
 #include <cstdint>
 #include <map>
@@ -27,7 +33,6 @@
 #include <system_error>
 #include <type_traits>
 #include <utility>
-#include <vector>
 
 using namespace llvm;
 using namespace bolt;

diff --git a/bolt/tools/heatmap/heatmap.cpp b/bolt/tools/heatmap/heatmap.cpp
@@ -1,12 +1,20 @@
-#include "bolt/Profile/DataAggregator.h"
+//===- bolt/tools/heatmap/heatmap.cpp - Profile heatmap visualization tool ===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
 #include "bolt/Rewrite/RewriteInstance.h"
 #include "bolt/Utils/CommandLineOpts.h"
 #include "llvm/MC/TargetRegistry.h"
 #include "llvm/Object/Binary.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Errc.h"
 #include "llvm/Support/Error.h"
-#include "llvm/Support/Path.h"
+#include "llvm/Support/FileSystem.h"
+#include "llvm/Support/Program.h"
 #include "llvm/Support/TargetSelect.h"
 
 using namespace llvm;

diff --git a/bolt/unittests/Core/BinaryContext.cpp b/bolt/unittests/Core/BinaryContext.cpp
@@ -1,7 +1,14 @@
+//===- bolt/unittest/Core/BinaryContext.cpp -------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
 #include "bolt/Core/BinaryContext.h"
 #include "llvm/BinaryFormat/ELF.h"
 #include "llvm/DebugInfo/DWARF/DWARFContext.h"
-#include "llvm/Object/ELFObjectFile.h"
 #include "llvm/Support/TargetSelect.h"
 #include "gtest/gtest.h"
 
@@ -40,8 +47,8 @@ struct BinaryContextTester : public testing::TestWithParam<Triple::ArchType> {
 
   void initializeBOLT() {
     BC = cantFail(BinaryContext::createBinaryContext(
-        ObjFile.get(), true, DWARFContext::create(*ObjFile.get()),
-        {llvm::outs(), llvm::errs()}));
+        ObjFile->makeTriple(), ObjFile->getFileName(), nullptr, true,
+        DWARFContext::create(*ObjFile.get()), {llvm::outs(), llvm::errs()}));
     ASSERT_FALSE(!BC);
   }
 

diff --git a/bolt/unittests/Core/MCPlusBuilder.cpp b/bolt/unittests/Core/MCPlusBuilder.cpp
@@ -1,3 +1,11 @@
+//===- bolt/unittest/Core/MCPlusBuilder.cpp -------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
 #ifdef AARCH64_AVAILABLE
 #include "AArch64Subtarget.h"
 #endif // AARCH64_AVAILABLE
@@ -11,7 +19,6 @@
 #include "bolt/Rewrite/RewriteInstance.h"
 #include "llvm/BinaryFormat/ELF.h"
 #include "llvm/DebugInfo/DWARF/DWARFContext.h"
-#include "llvm/Object/ELFObjectFile.h"
 #include "llvm/Support/TargetSelect.h"
 #include "gtest/gtest.h"
 
@@ -50,8 +57,8 @@ struct MCPlusBuilderTester : public testing::TestWithParam<Triple::ArchType> {
 
   void initializeBolt() {
     BC = cantFail(BinaryContext::createBinaryContext(
-        ObjFile.get(), true, DWARFContext::create(*ObjFile.get()),
-        {llvm::outs(), llvm::errs()}));
+        ObjFile->makeTriple(), ObjFile->getFileName(), nullptr, true,
+        DWARFContext::create(*ObjFile.get()), {llvm::outs(), llvm::errs()}));
     ASSERT_FALSE(!BC);
     BC->initializeTarget(std::unique_ptr<MCPlusBuilder>(
         createMCPlusBuilder(GetParam(), BC->MIA.get(), BC->MII.get(),

diff --git a/clang-tools-extra/clang-tidy/cppcoreguidelines/MissingStdForwardCheck.cpp b/clang-tools-extra/clang-tidy/cppcoreguidelines/MissingStdForwardCheck.cpp
@@ -9,8 +9,8 @@
 #include "MissingStdForwardCheck.h"
 #include "../utils/Matchers.h"
 #include "clang/AST/ASTContext.h"
-#include "clang/AST/ExprConcepts.h"
 #include "clang/ASTMatchers/ASTMatchFinder.h"
+#include "clang/Basic/IdentifierTable.h"
 
 using namespace clang::ast_matchers;
 
@@ -79,6 +79,11 @@ AST_MATCHER_P(LambdaExpr, hasCaptureDefaultKind, LambdaCaptureDefault, Kind) {
   return Node.getCaptureDefault() == Kind;
 }
 
+AST_MATCHER(VarDecl, hasIdentifier) {
+  const IdentifierInfo *ID = Node.getIdentifier();
+  return ID != NULL && !ID->isPlaceholder();
+}
+
 } // namespace
 
 void MissingStdForwardCheck::registerMatchers(MatchFinder *Finder) {
@@ -125,12 +130,14 @@ void MissingStdForwardCheck::registerMatchers(MatchFinder *Finder) {
                    hasAncestor(expr(hasUnevaluatedContext())))));
 
   Finder->addMatcher(
-      parmVarDecl(parmVarDecl().bind("param"), isTemplateTypeParameter(),
-                  hasAncestor(functionDecl().bind("func")),
-                  hasAncestor(functionDecl(
-                      isDefinition(), equalsBoundNode("func"), ToParam,
-                      unless(anyOf(isDeleted(), hasDescendant(std::move(
-                                                    ForwardCallMatcher))))))),
+      parmVarDecl(
+          parmVarDecl().bind("param"), hasIdentifier(),
+          unless(hasAttr(attr::Kind::Unused)), isTemplateTypeParameter(),
+          hasAncestor(functionDecl().bind("func")),
+          hasAncestor(functionDecl(
+              isDefinition(), equalsBoundNode("func"), ToParam,
+              unless(anyOf(isDeleted(),
+                           hasDescendant(std::move(ForwardCallMatcher))))))),
       this);
 }
 

diff --git a/clang-tools-extra/clang-tidy/hicpp/IgnoredRemoveResultCheck.cpp b/clang-tools-extra/clang-tidy/hicpp/IgnoredRemoveResultCheck.cpp
@@ -14,9 +14,9 @@ IgnoredRemoveResultCheck::IgnoredRemoveResultCheck(llvm::StringRef Name,
                                                    ClangTidyContext *Context)
     : UnusedReturnValueCheck(Name, Context,
                              {
-                                 "::std::remove",
-                                 "::std::remove_if",
-                                 "::std::unique",
+                                 "::std::remove$",
+                                 "::std::remove_if$",
+                                 "::std::unique$",
                              }) {
   // The constructor for ClangTidyCheck needs to have been called
   // before we can access options via Options.get().

diff --git a/clang-tools-extra/clang-tidy/readability/AvoidReturnWithVoidValueCheck.cpp b/clang-tools-extra/clang-tidy/readability/AvoidReturnWithVoidValueCheck.cpp
@@ -7,19 +7,18 @@
 //===----------------------------------------------------------------------===//
 
 #include "AvoidReturnWithVoidValueCheck.h"
-#include "clang/AST/Stmt.h"
-#include "clang/ASTMatchers/ASTMatchFinder.h"
-#include "clang/ASTMatchers/ASTMatchers.h"
+#include "../utils/BracesAroundStatement.h"
+#include "../utils/LexerUtils.h"
 
 using namespace clang::ast_matchers;
 
 namespace clang::tidy::readability {
 
-static constexpr auto IgnoreMacrosName = "IgnoreMacros";
-static constexpr auto IgnoreMacrosDefault = true;
+static constexpr char IgnoreMacrosName[] = "IgnoreMacros";
+static const bool IgnoreMacrosDefault = true;
 
-static constexpr auto StrictModeName = "StrictMode";
-static constexpr auto StrictModeDefault = true;
+static constexpr char StrictModeName[] = "StrictMode";
+static const bool StrictModeDefault = true;
 
 AvoidReturnWithVoidValueCheck::AvoidReturnWithVoidValueCheck(
     StringRef Name, ClangTidyContext *Context)
@@ -32,7 +31,10 @@ void AvoidReturnWithVoidValueCheck::registerMatchers(MatchFinder *Finder) {
   Finder->addMatcher(
       returnStmt(
           hasReturnValue(allOf(hasType(voidType()), unless(initListExpr()))),
-          optionally(hasParent(compoundStmt().bind("compound_parent"))))
+          optionally(hasParent(
+              compoundStmt(
+                  optionally(hasParent(functionDecl().bind("function_parent"))))
+                  .bind("compound_parent"))))
           .bind("void_return"),
       this);
 }
@@ -42,10 +44,30 @@ void AvoidReturnWithVoidValueCheck::check(
   const auto *VoidReturn = Result.Nodes.getNodeAs<ReturnStmt>("void_return");
   if (IgnoreMacros && VoidReturn->getBeginLoc().isMacroID())
     return;
-  if (!StrictMode && !Result.Nodes.getNodeAs<CompoundStmt>("compound_parent"))
+  const auto *SurroundingBlock =
+      Result.Nodes.getNodeAs<CompoundStmt>("compound_parent");
+  if (!StrictMode && !SurroundingBlock)
     return;
-  diag(VoidReturn->getBeginLoc(), "return statement within a void function "
-                                  "should not have a specified return value");
+  DiagnosticBuilder Diag = diag(VoidReturn->getBeginLoc(),
+                                "return statement within a void function "
+                                "should not have a specified return value");
+  const SourceLocation SemicolonPos = utils::lexer::findNextTerminator(
+      VoidReturn->getEndLoc(), *Result.SourceManager, getLangOpts());
+  if (SemicolonPos.isInvalid())
+    return;
+  if (!SurroundingBlock) {
+    const auto BraceInsertionHints = utils::getBraceInsertionsHints(
+        VoidReturn, getLangOpts(), *Result.SourceManager,
+        VoidReturn->getBeginLoc());
+    if (BraceInsertionHints)
+      Diag << BraceInsertionHints.openingBraceFixIt()
+           << BraceInsertionHints.closingBraceFixIt();
+  }
+  Diag << FixItHint::CreateRemoval(VoidReturn->getReturnLoc());
+  if (!Result.Nodes.getNodeAs<FunctionDecl>("function_parent") ||
+      SurroundingBlock->body_back() != VoidReturn)
+    Diag << FixItHint::CreateInsertion(SemicolonPos.getLocWithOffset(1),
+                                       " return;", true);
 }
 
 void AvoidReturnWithVoidValueCheck::storeOptions(

diff --git a/clang-tools-extra/clang-tidy/readability/BracesAroundStatementsCheck.cpp b/clang-tools-extra/clang-tidy/readability/BracesAroundStatementsCheck.cpp
@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "BracesAroundStatementsCheck.h"
+#include "../utils/BracesAroundStatement.h"
 #include "../utils/LexerUtils.h"
 #include "clang/AST/ASTContext.h"
 #include "clang/ASTMatchers/ASTMatchers.h"
@@ -17,12 +18,10 @@ using namespace clang::ast_matchers;
 namespace clang::tidy::readability {
 
 static tok::TokenKind getTokenKind(SourceLocation Loc, const SourceManager &SM,
-                                   const ASTContext *Context) {
+                                   const LangOptions &LangOpts) {
   Token Tok;
-  SourceLocation Beginning =
-      Lexer::GetBeginningOfToken(Loc, SM, Context->getLangOpts());
-  const bool Invalid =
-      Lexer::getRawToken(Beginning, Tok, SM, Context->getLangOpts());
+  SourceLocation Beginning = Lexer::GetBeginningOfToken(Loc, SM, LangOpts);
+  const bool Invalid = Lexer::getRawToken(Beginning, Tok, SM, LangOpts);
   assert(!Invalid && "Expected a valid token.");
 
   if (Invalid)
@@ -33,64 +32,21 @@ static tok::TokenKind getTokenKind(SourceLocation Loc, const SourceManager &SM,
 
 static SourceLocation
 forwardSkipWhitespaceAndComments(SourceLocation Loc, const SourceManager &SM,
-                                 const ASTContext *Context) {
+                                 const LangOptions &LangOpts) {
   assert(Loc.isValid());
   for (;;) {
     while (isWhitespace(*SM.getCharacterData(Loc)))
       Loc = Loc.getLocWithOffset(1);
 
-    tok::TokenKind TokKind = getTokenKind(Loc, SM, Context);
+    tok::TokenKind TokKind = getTokenKind(Loc, SM, LangOpts);
     if (TokKind != tok::comment)
       return Loc;
 
     // Fast-forward current token.
-    Loc = Lexer::getLocForEndOfToken(Loc, 0, SM, Context->getLangOpts());
+    Loc = Lexer::getLocForEndOfToken(Loc, 0, SM, LangOpts);
   }
 }
 
-static SourceLocation findEndLocation(const Stmt &S, const SourceManager &SM,
-                                      const ASTContext *Context) {
-  SourceLocation Loc =
-      utils::lexer::getUnifiedEndLoc(S, SM, Context->getLangOpts());
-  if (!Loc.isValid())
-    return Loc;
-
-  // Start searching right after S.
-  Loc = Loc.getLocWithOffset(1);
-
-  for (;;) {
-    assert(Loc.isValid());
-    while (isHorizontalWhitespace(*SM.getCharacterData(Loc))) {
-      Loc = Loc.getLocWithOffset(1);
-    }
-
-    if (isVerticalWhitespace(*SM.getCharacterData(Loc))) {
-      // EOL, insert brace before.
-      break;
-    }
-    tok::TokenKind TokKind = getTokenKind(Loc, SM, Context);
-    if (TokKind != tok::comment) {
-      // Non-comment token, insert brace before.
-      break;
-    }
-
-    SourceLocation TokEndLoc =
-        Lexer::getLocForEndOfToken(Loc, 0, SM, Context->getLangOpts());
-    SourceRange TokRange(Loc, TokEndLoc);
-    StringRef Comment = Lexer::getSourceText(
-        CharSourceRange::getTokenRange(TokRange), SM, Context->getLangOpts());
-    if (Comment.starts_with("/*") && Comment.contains('\n')) {
-      // Multi-line block comment, insert brace before.
-      break;
-    }
-    // else: Trailing comment, insert brace after the newline.
-
-    // Fast-forward current token.
-    Loc = TokEndLoc;
-  }
-  return Loc;
-}
-
 BracesAroundStatementsCheck::BracesAroundStatementsCheck(
     StringRef Name, ClangTidyContext *Context)
     : ClangTidyCheck(Name, Context),
@@ -124,7 +80,7 @@ void BracesAroundStatementsCheck::check(
   } else if (const auto *S = Result.Nodes.getNodeAs<DoStmt>("do")) {
     checkStmt(Result, S->getBody(), S->getDoLoc(), S->getWhileLoc());
   } else if (const auto *S = Result.Nodes.getNodeAs<WhileStmt>("while")) {
-    SourceLocation StartLoc = findRParenLoc(S, SM, Context);
+    SourceLocation StartLoc = findRParenLoc(S, SM, Context->getLangOpts());
     if (StartLoc.isInvalid())
       return;
     checkStmt(Result, S->getBody(), StartLoc);
@@ -133,7 +89,7 @@ void BracesAroundStatementsCheck::check(
     if (S->isConsteval())
       return;
 
-    SourceLocation StartLoc = findRParenLoc(S, SM, Context);
+    SourceLocation StartLoc = findRParenLoc(S, SM, Context->getLangOpts());
     if (StartLoc.isInvalid())
       return;
     if (ForceBracesStmts.erase(S))
@@ -156,7 +112,7 @@ template <typename IfOrWhileStmt>
 SourceLocation
 BracesAroundStatementsCheck::findRParenLoc(const IfOrWhileStmt *S,
                                            const SourceManager &SM,
-                                           const ASTContext *Context) {
+                                           const LangOptions &LangOpts) {
   // Skip macros.
   if (S->getBeginLoc().isMacroID())
     return {};
@@ -170,14 +126,14 @@ BracesAroundStatementsCheck::findRParenLoc(const IfOrWhileStmt *S,
   }
 
   SourceLocation PastCondEndLoc =
-      Lexer::getLocForEndOfToken(CondEndLoc, 0, SM, Context->getLangOpts());
+      Lexer::getLocForEndOfToken(CondEndLoc, 0, SM, LangOpts);
   if (PastCondEndLoc.isInvalid())
     return {};
   SourceLocation RParenLoc =
-      forwardSkipWhitespaceAndComments(PastCondEndLoc, SM, Context);
+      forwardSkipWhitespaceAndComments(PastCondEndLoc, SM, LangOpts);
   if (RParenLoc.isInvalid())
     return {};
-  tok::TokenKind TokKind = getTokenKind(RParenLoc, SM, Context);
+  tok::TokenKind TokKind = getTokenKind(RParenLoc, SM, LangOpts);
   if (TokKind != tok::r_paren)
     return {};
   return RParenLoc;
@@ -188,86 +144,23 @@ BracesAroundStatementsCheck::findRParenLoc(const IfOrWhileStmt *S,
 bool BracesAroundStatementsCheck::checkStmt(
     const MatchFinder::MatchResult &Result, const Stmt *S,
     SourceLocation StartLoc, SourceLocation EndLocHint) {
-
   while (const auto *AS = dyn_cast<AttributedStmt>(S))
     S = AS->getSubStmt();
 
-  const SourceManager &SM = *Result.SourceManager;
-  const ASTContext *Context = Result.Context;
-
-  // 1) If there's a corresponding "else" or "while", the check inserts "} "
-  // right before that token.
-  // 2) If there's a multi-line block comment starting on the same line after
-  // the location we're inserting the closing brace at, or there's a non-comment
-  // token, the check inserts "\n}" right before that token.
-  // 3) Otherwise the check finds the end of line (possibly after some block or
-  // line comments) and inserts "\n}" right before that EOL.
-  if (!S || isa<CompoundStmt>(S)) {
-    // Already inside braces.
-    return false;
-  }
-
-  // When TreeTransform, Stmt in constexpr IfStmt will be transform to NullStmt.
-  // This NullStmt can be detected according to beginning token.
-  const SourceLocation StmtBeginLoc = S->getBeginLoc();
-  if (isa<NullStmt>(S) && StmtBeginLoc.isValid() &&
-      getTokenKind(StmtBeginLoc, SM, Context) == tok::l_brace)
-    return false;
-
-  if (StartLoc.isInvalid())
-    return false;
-
-  // Convert StartLoc to file location, if it's on the same macro expansion
-  // level as the start of the statement. We also need file locations for
-  // Lexer::getLocForEndOfToken working properly.
-  StartLoc = Lexer::makeFileCharRange(
-                 CharSourceRange::getCharRange(StartLoc, S->getBeginLoc()), SM,
-                 Context->getLangOpts())
-                 .getBegin();
-  if (StartLoc.isInvalid())
-    return false;
-  StartLoc =
-      Lexer::getLocForEndOfToken(StartLoc, 0, SM, Context->getLangOpts());
-
-  // StartLoc points at the location of the opening brace to be inserted.
-  SourceLocation EndLoc;
-  std::string ClosingInsertion;
-  if (EndLocHint.isValid()) {
-    EndLoc = EndLocHint;
-    ClosingInsertion = "} ";
-  } else {
-    EndLoc = findEndLocation(*S, SM, Context);
-    ClosingInsertion = "\n}";
-  }
-
-  assert(StartLoc.isValid());
-
-  // Don't require braces for statements spanning less than certain number of
-  // lines.
-  if (ShortStatementLines && !ForceBracesStmts.erase(S)) {
-    unsigned StartLine = SM.getSpellingLineNumber(StartLoc);
-    unsigned EndLine = SM.getSpellingLineNumber(EndLoc);
-    if (EndLine - StartLine < ShortStatementLines)
+  const auto BraceInsertionHints = utils::getBraceInsertionsHints(
+      S, Result.Context->getLangOpts(), *Result.SourceManager, StartLoc,
+      EndLocHint);
+  if (BraceInsertionHints) {
+    if (ShortStatementLines && !ForceBracesStmts.erase(S) &&
+        BraceInsertionHints.resultingCompoundLineExtent(*Result.SourceManager) <
+            ShortStatementLines)
       return false;
+    auto Diag = diag(BraceInsertionHints.DiagnosticPos,
+                     "statement should be inside braces");
+    if (BraceInsertionHints.offersFixIts())
+      Diag << BraceInsertionHints.openingBraceFixIt()
+           << BraceInsertionHints.closingBraceFixIt();
   }
-
-  auto Diag = diag(StartLoc, "statement should be inside braces");
-
-  // Change only if StartLoc and EndLoc are on the same macro expansion level.
-  // This will also catch invalid EndLoc.
-  // Example: LLVM_DEBUG( for(...) do_something() );
-  // In this case fix-it cannot be provided as the semicolon which is not
-  // visible here is part of the macro. Adding braces here would require adding
-  // another semicolon.
-  if (Lexer::makeFileCharRange(
-          CharSourceRange::getTokenRange(SourceRange(
-              SM.getSpellingLoc(StartLoc), SM.getSpellingLoc(EndLoc))),
-          SM, Context->getLangOpts())
-          .isInvalid())
-    return false;
-
-  Diag << FixItHint::CreateInsertion(StartLoc, " {")
-       << FixItHint::CreateInsertion(EndLoc, ClosingInsertion);
   return true;
 }
 

diff --git a/clang-tools-extra/clang-tidy/readability/BracesAroundStatementsCheck.h b/clang-tools-extra/clang-tidy/readability/BracesAroundStatementsCheck.h
@@ -52,7 +52,7 @@ class BracesAroundStatementsCheck : public ClangTidyCheck {
                  SourceLocation EndLocHint = SourceLocation());
   template <typename IfOrWhileStmt>
   SourceLocation findRParenLoc(const IfOrWhileStmt *S, const SourceManager &SM,
-                               const ASTContext *Context);
+                               const LangOptions &LangOpts);
   std::optional<TraversalKind> getCheckTraversalKind() const override {
     return TK_IgnoreUnlessSpelledInSource;
   }

diff --git a/clang-tools-extra/clang-tidy/readability/CMakeLists.txt b/clang-tools-extra/clang-tidy/readability/CMakeLists.txt
@@ -17,6 +17,7 @@ add_clang_library(clangTidyReadabilityModule
   DeleteNullPointerCheck.cpp
   DuplicateIncludeCheck.cpp
   ElseAfterReturnCheck.cpp
+  EnumInitialValueCheck.cpp
   FunctionCognitiveComplexityCheck.cpp
   FunctionSizeCheck.cpp
   IdentifierLengthCheck.cpp

diff --git a/clang-tools-extra/clang-tidy/readability/DuplicateIncludeCheck.cpp b/clang-tools-extra/clang-tidy/readability/DuplicateIncludeCheck.cpp
@@ -79,6 +79,10 @@ void DuplicateIncludeCallbacks::InclusionDirective(
     bool IsAngled, CharSourceRange FilenameRange, OptionalFileEntryRef File,
     StringRef SearchPath, StringRef RelativePath, const Module *SuggestedModule,
     bool ModuleImported, SrcMgr::CharacteristicKind FileType) {
+  // Skip includes behind macros
+  if (FilenameRange.getBegin().isMacroID() ||
+      FilenameRange.getEnd().isMacroID())
+    return;
   if (llvm::is_contained(Files.back(), FileName)) {
     // We want to delete the entire line, so make sure that [Start,End] covers
     // everything.