From ec1ba207be2320e398741e4f220bf5323ffc40a8 Mon Sep 17 00:00:00 2001
From: Maksim Panchenko <maks@fb.com>
Date: Sun, 12 Nov 2023 19:34:42 -0800
Subject: [PATCH] [BOLT] Enhance LowerAnnotations pass. NFCI. (#71847)

After #70147, all primary annotation types are stored directly in the
instruction and hence there's no need for the temporary storage we've
used previously for repopulating preserved annotations.
---
 bolt/lib/Passes/BinaryPasses.cpp | 58 +++++++++++---------------------
 1 file changed, 20 insertions(+), 38 deletions(-)

diff --git a/bolt/lib/Passes/BinaryPasses.cpp b/bolt/lib/Passes/BinaryPasses.cpp
index 5366a62d73281..83c7138e5fe5c 100644
--- a/bolt/lib/Passes/BinaryPasses.cpp
+++ b/bolt/lib/Passes/BinaryPasses.cpp
@@ -582,65 +582,47 @@ bool CheckLargeFunctions::shouldOptimize(const BinaryFunction &BF) const {
 }
 
 void LowerAnnotations::runOnFunctions(BinaryContext &BC) {
-  std::vector<std::pair<MCInst *, uint32_t>> PreservedOffsetAnnotations;
-  std::vector<std::pair<MCInst *, MCSymbol *>> PreservedLabelAnnotations;
-
-  for (auto &It : BC.getBinaryFunctions()) {
-    BinaryFunction &BF = It.second;
-
-    for (FunctionFragment &FF : BF.getLayout().fragments()) {
+  for (BinaryFunction *BF : BC.getAllBinaryFunctions()) {
+    for (FunctionFragment &FF : BF->getLayout().fragments()) {
+      // Reset at the start of the new fragment.
       int64_t CurrentGnuArgsSize = 0;
 
       for (BinaryBasicBlock *const BB : FF) {
-        // First convert GnuArgsSize annotations into CFIs. This may change
-        // instr pointers, so do it before recording ptrs for preserved
-        // annotations
-        if (BF.usesGnuArgsSize()) {
-          for (auto II = BB->begin(); II != BB->end(); ++II) {
-            if (!BC.MIB->isInvoke(*II))
-              continue;
+        for (auto II = BB->begin(); II != BB->end(); ++II) {
+
+          // Convert GnuArgsSize annotations into CFIs.
+          if (BF->usesGnuArgsSize() && BC.MIB->isInvoke(*II)) {
             const int64_t NewGnuArgsSize = BC.MIB->getGnuArgsSize(*II);
             assert(NewGnuArgsSize >= 0 &&
-                   "expected non-negative GNU_args_size");
+                   "Expected non-negative GNU_args_size.");
             if (NewGnuArgsSize != CurrentGnuArgsSize) {
-              auto InsertII = BF.addCFIInstruction(
+              auto InsertII = BF->addCFIInstruction(
                   BB, II,
                   MCCFIInstruction::createGnuArgsSize(nullptr, NewGnuArgsSize));
               CurrentGnuArgsSize = NewGnuArgsSize;
               II = std::next(InsertII);
             }
           }
-        }
 
-        // Now record preserved annotations separately and then strip
-        // annotations.
-        for (auto II = BB->begin(); II != BB->end(); ++II) {
-          if (BF.requiresAddressTranslation() && BC.MIB->getOffset(*II))
-            PreservedOffsetAnnotations.emplace_back(&(*II),
-                                                    *BC.MIB->getOffset(*II));
-          if (MCSymbol *Label = BC.MIB->getLabel(*II))
-            PreservedLabelAnnotations.emplace_back(&*II, Label);
+          // Preserve selected annotations and strip the rest.
+          std::optional<uint32_t> Offset = BF->requiresAddressTranslation()
+                                               ? BC.MIB->getOffset(*II)
+                                               : std::nullopt;
+          MCSymbol *Label = BC.MIB->getLabel(*II);
+
           BC.MIB->stripAnnotations(*II);
+
+          if (Offset)
+            BC.MIB->setOffset(*II, *Offset);
+          if (Label)
+            BC.MIB->setLabel(*II, Label);
         }
       }
     }
   }
-  for (BinaryFunction *BF : BC.getInjectedBinaryFunctions())
-    for (BinaryBasicBlock &BB : *BF)
-      for (MCInst &Instruction : BB) {
-        if (MCSymbol *Label = BC.MIB->getLabel(Instruction))
-          PreservedLabelAnnotations.emplace_back(&Instruction, Label);
-        BC.MIB->stripAnnotations(Instruction);
-      }
 
   // Release all memory taken by annotations
   BC.MIB->freeAnnotations();
-
-  // Reinsert preserved annotations we need during code emission.
-  for (const std::pair<MCInst *, uint32_t> &Item : PreservedOffsetAnnotations)
-    BC.MIB->setOffset(*Item.first, Item.second);
-  for (auto [Instr, Label] : PreservedLabelAnnotations)
-    BC.MIB->setLabel(*Instr, Label);
 }
 
 // Check for dirty state in MCSymbol objects that might be a consequence