[Courgette] Refactor: Store Label Annotation in AssemblyProgram for patch generation.

courgette/assembly_program.h

 // Copyright (c) 2011 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 #ifndef COURGETTE_ASSEMBLY_PROGRAM_H_
 #define COURGETTE_ASSEMBLY_PROGRAM_H_
 
 #include <stddef.h>
 #include <stdint.h>
 
-#include <map>
 #include <memory>
+#include <vector>
 
 #include "base/callback_forward.h"
 #include "base/macros.h"
 #include "base/memory/free_deleter.h"
 #include "courgette/courgette.h"
 #include "courgette/image_utils.h"
 #include "courgette/instruction_utils.h"
 #include "courgette/label_manager.h"
 #include "courgette/memory_allocator.h"
 
@@ skipping 48 matching lines @@
 // The optional step is to modify the AssemblyProgram.  One form of modification
 // is to assign indexes in such a way as to make the EncodedProgram for this
 // AssemblyProgram look more like the EncodedProgram for some other
 // AssemblyProgram.  The modification process should call UnassignIndexes, do
 // its own assignment, and then call AssignRemainingIndexes to ensure all
 // indexes are assigned.
 
 class AssemblyProgram {
  public:
   using LabelHandler = base::Callback<void(Label*)>;
-  using LabelHandlerMap = std::map<OP, LabelHandler>;
 
   AssemblyProgram(ExecutableType kind, uint64_t image_base);
   ~AssemblyProgram();
 
   ExecutableType kind() const { return kind_; }
+  const std::vector<Label*>& abs32_label_annotations() const {
+    return abs32_label_annotations_;
+  }
+  const std::vector<Label*>& rel32_label_annotations() const {
+    return rel32_label_annotations_;
+  }
+  std::vector<Label*>* mutable_abs32_label_annotations() {
+    return &abs32_label_annotations_;
+  }
+  std::vector<Label*>* mutable_rel32_label_annotations() {
+    return &rel32_label_annotations_;
+  }
 
   // Traverses RVAs in |abs32_visitor| and |rel32_visitor| to precompute Labels.
   void PrecomputeLabels(RvaVisitor* abs32_visitor, RvaVisitor* rel32_visitor);
 
   // Removes underused Labels. Thresholds used (0 = no trimming) is
   // architecture-dependent.
   void TrimLabels();
 
   void UnassignIndexes();
   void DefaultAssignIndexes();
   void AssignRemainingIndexes();
 
   // Looks up abs32 label. Returns null if none found.
   Label* FindAbs32Label(RVA rva);
 
   // Looks up rel32 label. Returns null if none found.
   Label* FindRel32Label(RVA rva);
 
-  std::unique_ptr<EncodedProgram> Encode() const;
-
-  // For each |instruction| in |instructions_|, looks up its opcode from
-  // |handler_map| for a handler. If a handler exists, invoke it by passing the
-  // |instruction|'s label. We assume that |handler_map| has correct keys, i.e.,
-  // opcodes for an instruction that have label.
-  void HandleInstructionLabels(const LabelHandlerMap& handler_map) const;
-
   // Calls |gen| in 2 passes to emit instructions. In pass 1 we provide a
   // receptor to count space requirement. In pass 2 we provide a receptor to
-  // store instructions.
-  CheckBool GenerateInstructions(const InstructionGenerator& gen);
+  // store instructions. If |annotate_labels| is true, then extracts Label
+  // annotations into |*_label_annotations_|.
+  CheckBool GenerateInstructions(const InstructionGenerator& gen,
+                                 bool annotate_labels);
+
+  // Returns an EncodeProgram that converts program to encoded form.
+  std::unique_ptr<EncodedProgram> Encode() const;
 
   // TODO(huangs): Implement these in InstructionStoreReceptor.
   // Instructions will be assembled in the order they are emitted.
 
   // Generates an entire base relocation table.
   CheckBool EmitPeRelocs() WARN_UNUSED_RESULT;
 
   // Generates an ELF style relocation table for X86.
   CheckBool EmitElfRelocation() WARN_UNUSED_RESULT;
 
@@ skipping 47 matching lines @@
 
   std::unique_ptr<Instruction* [], base::FreeDeleter> byte_instruction_cache_;
 
   InstructionVector instructions_;  // All the instructions in program.
 
   // Storage and lookup of Labels associated with target addresses. We use
   // separate abs32 and rel32 labels.
   LabelManager abs32_label_manager_;
   LabelManager rel32_label_manager_;
 
+  // Label pointers for each abs32 and rel32 location, sorted by file offset.
+  // These are used by Label adjustment during patch generation.
+  std::vector<Label*> abs32_label_annotations_;
+  std::vector<Label*> rel32_label_annotations_;
+
   DISALLOW_COPY_AND_ASSIGN(AssemblyProgram);
 };
 
 // Converts |program| into encoded form, returning it as |*output|.
 // Returns C_OK if succeeded, otherwise returns an error status and sets
 // |*output| to null.
 Status Encode(const AssemblyProgram& program,
               std::unique_ptr<EncodedProgram>* output);
 
 }  // namespace courgette
 
 #endif  // COURGETTE_ASSEMBLY_PROGRAM_H_