[Courgette] Refactor: Store Label Annotation in AssemblyProgram for p…

…atch generation.

For Courgette-gen, label adjustment needs lists of abs32 & rel32 Label*
sorted by file offset. Let these lists be "Label Annotations".
Previously these were extracted during label adjustment from list of
instructions in AssemblyProgram, but now we wish to remove these stored
instructions.

This CL make AssemblyProgram store Label annotations. These are computed
only when needed (Courgette-gen / -gen1a). Details:

- Add ParseDetecteExecutableWithAnnotation() alongside
  ParseDetecteExecutable(), to avoid affecting flows that don't require
  Label Annotations.
- AssemblyProgram: Add |*_label_annotations_| as storage vectors. These
  are optionally populated in InstructionStoreReceptor when
  Disassembler::Disassemble() gets called.
  - InstructionCountReceptor now put into use.
- Simplify Label adjustment: AssemblyProgram::HandleInstructionLabels()
  is replaced with visiting AssemblyProgram's Label Annotations.
- Subtle: GraphAdjuster now includes abs64. The class doesn't get used,
  but the change is logical anyway.

BUG=660980

Review-Url: https://codereview.chromium.org/2793153003
Cr-Commit-Position: refs/heads/master@{#464536}

courgette/adjustment_method.cc

@@ -14,7 +14,6 @@
 #include <string>
 #include <vector>
 
-#include "base/bind.h"
 #include "base/logging.h"
 #include "base/macros.h"
 #include "base/strings/string_number_conversions.h"
@@ -595,16 +594,10 @@ class GraphAdjuster : public AdjustmentMethod {
  private:
   void CollectTraces(const AssemblyProgram* program, Trace* abs32, Trace* rel32,
                      bool is_model) {
-    AssemblyProgram::LabelHandler abs32_handler =
-        base::Bind(&GraphAdjuster::ReferenceLabel, base::Unretained(this),
-                   abs32, is_model);
-    AssemblyProgram::LabelHandler rel32_handler =
-        base::Bind(&GraphAdjuster::ReferenceLabel, base::Unretained(this),
-                   rel32, is_model);
-
-    program->HandleInstructionLabels({{ABS32, abs32_handler},
-                                      {REL32, rel32_handler},
-                                      {REL32ARM, rel32_handler}});
+    for (Label* label : program->abs32_label_annotations())
+      ReferenceLabel(abs32, is_model, label);
+    for (Label* label : program->rel32_label_annotations())
+      ReferenceLabel(rel32, is_model, label);
 
     // TODO(sra): we could simply append all the labels in index order to
     // incorporate some costing for entropy (bigger deltas) that will be

courgette/adjustment_method_2.cc

@@ -15,7 +15,6 @@
 #include <string>
 #include <vector>
 
-#include "base/bind.h"
 #include "base/format_macros.h"
 #include "base/logging.h"
 #include "base/macros.h"
@@ -1254,15 +1253,10 @@ class Adjuster : public AdjustmentMethod {
                      bool is_model) {
     label_info_maker_.ResetDebugLabel();
 
-    AssemblyProgram::LabelHandler abs32_handler = base::Bind(
-        &Adjuster::ReferenceLabel, base::Unretained(this), abs32, is_model);
-    AssemblyProgram::LabelHandler rel32_handler = base::Bind(
-        &Adjuster::ReferenceLabel, base::Unretained(this), rel32, is_model);
-
-    program->HandleInstructionLabels({{ABS32, abs32_handler},
-                                      {ABS64, abs32_handler},
-                                      {REL32, rel32_handler},
-                                      {REL32ARM, rel32_handler}});
+    for (Label* label : program->abs32_label_annotations())
+      ReferenceLabel(abs32, is_model, label);
+    for (Label* label : program->rel32_label_annotations())
+      ReferenceLabel(rel32, is_model, label);
 
     // TODO(sra): we could simply append all the labels in index order to
     // incorporate some costing for entropy (bigger deltas) that will be

courgette/adjustment_method_unittest.cc

@@ -65,7 +65,9 @@ class AdjustmentMethodTest : public testing::Test {
         },
         labelA, labelB);
 
-    EXPECT_TRUE(prog->GenerateInstructions(gen));
+    EXPECT_TRUE(prog->GenerateInstructions(gen, true));
+    EXPECT_EQ(6U, prog->abs32_label_annotations().size());
+    EXPECT_EQ(0U, prog->rel32_label_annotations().size());
 
     if (kind == 0) {
       labelA->index_ = 0;

courgette/assembly_program.cc

@@ -7,6 +7,7 @@
 #include "base/callback.h"
 #include "base/logging.h"
 #include "courgette/courgette.h"
+#include "courgette/disassembler.h"
 #include "courgette/encoded_program.h"
 
 namespace courgette {
@@ -104,10 +105,10 @@ class InstructionCountReceptor : public InstructionReceptor {
  public:
   InstructionCountReceptor() = default;
 
-  size_t size() const { return size_; }
+  size_t abs_count() const { return abs_count_; }
+  size_t rel_count() const { return rel_count_; }
 
   // InstructionReceptor:
-  // TODO(huangs): 2016/11: Populate these with size_ += ...
   CheckBool EmitPeRelocs() override { return true; }
   CheckBool EmitElfRelocation() override { return true; }
   CheckBool EmitElfARMRelocation() override { return true; }
@@ -116,18 +117,29 @@ class InstructionCountReceptor : public InstructionReceptor {
   CheckBool EmitMultipleBytes(const uint8_t* bytes, size_t len) override {
     return true;
   }
-  CheckBool EmitRel32(Label* label) override { return true; }
+  CheckBool EmitRel32(Label* label) override {
+    ++rel_count_;
+    return true;
+  }
   CheckBool EmitRel32ARM(uint16_t op,
                          Label* label,
                          const uint8_t* arm_op,
                          uint16_t op_size) override {
+    ++rel_count_;
+    return true;
+  }
+  CheckBool EmitAbs32(Label* label) override {
+    ++abs_count_;
+    return true;
+  }
+  CheckBool EmitAbs64(Label* label) override {
+    ++abs_count_;
     return true;
   }
-  CheckBool EmitAbs32(Label* label) override { return true; }
-  CheckBool EmitAbs64(Label* label) override { return true; }
 
  private:
-  size_t size_ = 0;
+  size_t abs_count_ = 0;
+  size_t rel_count_ = 0;
 
   DISALLOW_COPY_AND_ASSIGN(InstructionCountReceptor);
 };
@@ -137,15 +149,15 @@ class InstructionCountReceptor : public InstructionReceptor {
 // An InstructionReceptor that stores emitted instructions.
 class InstructionStoreReceptor : public InstructionReceptor {
  public:
-  explicit InstructionStoreReceptor(AssemblyProgram* program)
-      : program_(program) {
+  InstructionStoreReceptor(AssemblyProgram* program, bool annotate_labels)
+      : program_(program), annotate_labels_(annotate_labels) {
     CHECK(program_);
   }
 
-  // TODO(huangs): 2016/11: Add Reserve().
+  // TODO(huangs): 2017/04: Add Reserve().
 
   // InstructionReceptor:
-  // TODO(huangs): 2016/11: Replace stub with implementation.
+  // TODO(huangs): 2017/04: Move implementations here.
   CheckBool EmitPeRelocs() override { return program_->EmitPeRelocs(); }
   CheckBool EmitElfRelocation() override {
     return program_->EmitElfRelocation();
@@ -161,23 +173,32 @@ class InstructionStoreReceptor : public InstructionReceptor {
     return program_->EmitMultipleBytes(bytes, len);
   }
   CheckBool EmitRel32(Label* label) override {
+    if (annotate_labels_)
+      program_->mutable_rel32_label_annotations()->push_back(label);
     return program_->EmitRel32(label);
   }
   CheckBool EmitRel32ARM(uint16_t op,
                          Label* label,
                          const uint8_t* arm_op,
                          uint16_t op_size) override {
+    if (annotate_labels_)
+      program_->mutable_rel32_label_annotations()->push_back(label);
     return program_->EmitRel32ARM(op, label, arm_op, op_size);
   }
   CheckBool EmitAbs32(Label* label) override {
+    if (annotate_labels_)
+      program_->mutable_abs32_label_annotations()->push_back(label);
     return program_->EmitAbs32(label);
   }
   CheckBool EmitAbs64(Label* label) override {
+    if (annotate_labels_)
+      program_->mutable_abs32_label_annotations()->push_back(label);
     return program_->EmitAbs64(label);
   }
 
  private:
   AssemblyProgram* program_;
+  const bool annotate_labels_;
 
   DISALLOW_COPY_AND_ASSIGN(InstructionStoreReceptor);
 };
@@ -294,28 +315,24 @@ Label* AssemblyProgram::FindRel32Label(RVA rva) {
   return rel32_label_manager_.Find(rva);
 }
 
-void AssemblyProgram::HandleInstructionLabels(
-    const AssemblyProgram::LabelHandlerMap& handler_map) const {
-  for (const Instruction* instruction : instructions_) {
-    LabelHandlerMap::const_iterator it = handler_map.find(instruction->op());
-    if (it != handler_map.end()) {
-      it->second.Run(
-          static_cast<const InstructionWithLabel*>(instruction)->label());
-    }
-  }
-}
-
-CheckBool AssemblyProgram::GenerateInstructions(
-    const InstructionGenerator& gen) {
-  // Pass 1: Count the space needed to store instructions.
+CheckBool AssemblyProgram::GenerateInstructions(const InstructionGenerator& gen,
+                                                bool annotate_labels) {
+  // Pass 1: Count storage space required and reserve in advance.
   InstructionCountReceptor count_receptor;
   if (!gen.Run(&count_receptor))
     return false;
 
-  // Pass 2: Emit all instructions to preallocated buffer (uses Phase 1 count).
-  InstructionStoreReceptor store_receptor(this);
-  // TODO(huangs): 2017/03: Pass |count_receptor->size()| to |store_receptor_|
-  // to reserve space for raw data.
+  if (annotate_labels) {
+    DCHECK(abs32_label_annotations_.empty());
+    abs32_label_annotations_.reserve(count_receptor.abs_count());
+    DCHECK(rel32_label_annotations_.empty());
+    rel32_label_annotations_.reserve(count_receptor.rel_count());
+  }
+
+  // Pass 2: Emit all instructions to reserved buffer (uses Phase 1 count).
+  // Populates |abs32_label_annotations_| and |re32_label_annotations_| if
+  // |annotate_labels| is true.
+  InstructionStoreReceptor store_receptor(this, annotate_labels);
   return gen.Run(&store_receptor);
 }
 

courgette/assembly_program.h

@@ -8,8 +8,8 @@
 #include <stddef.h>
 #include <stdint.h>
 
-#include <map>
 #include <memory>
+#include <vector>
 
 #include "base/callback_forward.h"
 #include "base/macros.h"
@@ -78,12 +78,23 @@ class Instruction {
 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);
@@ -102,18 +113,15 @@ class AssemblyProgram {
   // 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.
@@ -181,6 +189,11 @@ class AssemblyProgram {
   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);
 };
 

courgette/courgette_tool.cc

@@ -124,8 +124,7 @@ bool Supported(const base::FilePath& input_file) {
   // If the detection fails, we just fall back on UNKNOWN
   std::string format = "Unsupported";
 
-  switch (type)
-  {
+  switch (type) {
     case courgette::EXE_UNKNOWN:
       break;
 
@@ -214,17 +213,18 @@ void DisassembleAdjustDiff(const base::FilePath& model_file,
   BufferedFileReader model_buffer(model_file, "old");
   BufferedFileReader program_buffer(program_file, "new");
 
+  auto parser = adjust ? courgette::ParseDetectedExecutableWithAnnotation
+                       : courgette::ParseDetectedExecutable;
+
   std::unique_ptr<courgette::AssemblyProgram> model;
   const courgette::Status parse_model_status =
-      courgette::ParseDetectedExecutable(model_buffer.data(),
-                                         model_buffer.length(), &model);
+      parser(model_buffer.data(), model_buffer.length(), &model);
   if (parse_model_status != courgette::C_OK)
     Problem("Can't parse model input (code = %d).", parse_model_status);
 
   std::unique_ptr<courgette::AssemblyProgram> program;
   const courgette::Status parse_program_status =
-      courgette::ParseDetectedExecutable(program_buffer.data(),
-                                         program_buffer.length(), &program);
+      parser(program_buffer.data(), program_buffer.length(), &program);
   if (parse_program_status != courgette::C_OK)
     Problem("Can't parse program input (code = %d).", parse_program_status);
 

courgette/disassembler.cc

@@ -54,7 +54,8 @@ const uint8_t* Disassembler::RVAToPointer(RVA rva) const {
   return FileOffsetToPointer(file_offset);
 }
 
-std::unique_ptr<AssemblyProgram> Disassembler::Disassemble() {
+std::unique_ptr<AssemblyProgram> Disassembler::Disassemble(
+    bool annotate_labels) {
   if (!ok() || !ExtractAbs32Locations() || !ExtractRel32Locations())
     return nullptr;
 
@@ -64,8 +65,10 @@ std::unique_ptr<AssemblyProgram> Disassembler::Disassemble() {
   PrecomputeLabels(program.get());
   RemoveUnusedRel32Locations(program.get());
 
-  if (!program->GenerateInstructions(GetInstructionGenerator(program.get())))
+  if (!program->GenerateInstructions(GetInstructionGenerator(program.get()),
+                                     annotate_labels)) {
     return nullptr;
+  }
 
   program->DefaultAssignIndexes();
   return program;

courgette/disassembler.h

@@ -91,9 +91,9 @@ class Disassembler : public AddressTranslator {
   virtual bool ParseHeader() = 0;
 
   // Extracts and stores references from the main image. Returns a new
-  // AssemblyProgram initialized using data parsed from the main image, or null
-  // on failure.
-  std::unique_ptr<AssemblyProgram> Disassemble();
+  // AssemblyProgram initialized using data parsed from the main image and
+  // |annotate_labels|, or null on failure.
+  std::unique_ptr<AssemblyProgram> Disassemble(bool annotate_labels);
 
   // ok() may always be called but returns true only after ParseHeader()
   // succeeds.

courgette/disassembler_elf_32_x86_unittest.cc

@@ -91,7 +91,7 @@ void DisassemblerElf32X86Test::TestExe(const char* file_name,
   EXPECT_EQ('L', offset_p[2]);
   EXPECT_EQ('F', offset_p[3]);
 
-  std::unique_ptr<AssemblyProgram> program = disassembler->Disassemble();
+  std::unique_ptr<AssemblyProgram> program = disassembler->Disassemble(false);
   EXPECT_TRUE(nullptr != program.get());
 
   const std::vector<RVA>& abs32_list = disassembler->Abs32Locations();