[Courgette] Using LabelManager to reduce Courgette-apply peak RAM by 25%.

courgette/assembly_program.cc

Index: courgette/assembly_program.cc
diff --git a/courgette/assembly_program.cc b/courgette/assembly_program.cc
index 5344cff95f84868f83c670b143b88f9fef736564..f20ed67ea3fa3e6a462f7d99e49153b7eec3d22c 100644
--- a/courgette/assembly_program.cc
+++ b/courgette/assembly_program.cc
@@ -111,11 +111,6 @@ AssemblyProgram::AssemblyProgram(ExecutableType kind)
   : kind_(kind), image_base_(0) {
 }
 
-static void DeleteContainedLabels(const RVAToLabel& labels) {
-  for (RVAToLabel::const_iterator p = labels.begin();  p != labels.end();  ++p)
-    UncheckedDelete(p->second);
-}
-
 AssemblyProgram::~AssemblyProgram() {
   for (size_t i = 0;  i < instructions_.size();  ++i) {
     Instruction* instruction = instructions_[i];
@@ -126,8 +121,6 @@ AssemblyProgram::~AssemblyProgram() {
     for (size_t i = 0;  i < 256;  ++i)
       UncheckedDelete(byte_instruction_cache_[i]);
   }
-  DeleteContainedLabels(rel32_labels_);
-  DeleteContainedLabels(abs32_labels_);
 }
 
 CheckBool AssemblyProgram::EmitPeRelocsInstruction() {
@@ -178,27 +171,50 @@ CheckBool AssemblyProgram::EmitAbs64(Label* label) {
       ScopedInstruction(UncheckedNew<InstructionWithLabel>(ABS64, label)));
 }
 
-Label* AssemblyProgram::FindOrMakeAbs32Label(RVA rva) {
-  return FindLabel(rva, &abs32_labels_);
+void AssemblyProgram::PrecomputeLabels(RvaVisitor* abs32_visitor,
+                                       RvaVisitor* rel32_visitor) {
+  abs32_label_manager_.Read(abs32_visitor);
+  rel32_label_manager_.Read(rel32_visitor);
+  TrimLabels();
 }
 
-Label* AssemblyProgram::FindOrMakeRel32Label(RVA rva) {
-  return FindLabel(rva, &rel32_labels_);
-}
+// Chosen empirically to give the best reduction in payload size for
+// an update from daisy_3701.98.0 to daisy_4206.0.0.
+const int AssemblyProgram::kLabelLowerLimit = 5;
 
-void AssemblyProgram::DefaultAssignIndexes() {
-  DefaultAssignIndexes(&abs32_labels_);
-  DefaultAssignIndexes(&rel32_labels_);
+void AssemblyProgram::TrimLabels() {
+  // For now only trim for ARM binaries.
+  if (kind() != EXE_ELF_32_ARM)
+    return;
+
+  int lower_limit = kLabelLowerLimit;
+
+  VLOG(1) << "TrimLabels: threshold " << lower_limit;
+
+  rel32_label_manager_.RemoveUnderusedLabels(lower_limit);
 }
 
 void AssemblyProgram::UnassignIndexes() {
-  UnassignIndexes(&abs32_labels_);
-  UnassignIndexes(&rel32_labels_);
+  abs32_label_manager_.UnassignIndexes();
+  rel32_label_manager_.UnassignIndexes();
+}
+
+void AssemblyProgram::DefaultAssignIndexes() {
+  abs32_label_manager_.DefaultAssignIndexes();
+  rel32_label_manager_.DefaultAssignIndexes();
 }
 
 void AssemblyProgram::AssignRemainingIndexes() {
-  AssignRemainingIndexes(&abs32_labels_);
-  AssignRemainingIndexes(&rel32_labels_);
+  abs32_label_manager_.AssignRemainingIndexes();
+  rel32_label_manager_.AssignRemainingIndexes();
+}
+
+Label* AssemblyProgram::FindAbs32Label(RVA rva) {
+  return abs32_label_manager_.Find(rva);
+}
+
+Label* AssemblyProgram::FindRel32Label(RVA rva) {
+  return rel32_label_manager_.Find(rva);
 }
 
 Label* AssemblyProgram::InstructionAbs32Label(
@@ -238,17 +254,6 @@ CheckBool AssemblyProgram::EmitShared(Instruction* instruction) {
   return instruction && instructions_.push_back(instruction);
 }
 
-Label* AssemblyProgram::FindLabel(RVA rva, RVAToLabel* labels) {
-  Label*& slot = (*labels)[rva];
-  if (slot == NULL) {
-    slot = UncheckedNew<Label>(rva);
-    if (slot == NULL)
-      return NULL;
-  }
-  slot->count_++;
-  return slot;
-}
-
 void AssemblyProgram::UnassignIndexes(RVAToLabel* labels) {
   for (RVAToLabel::iterator p = labels->begin();  p != labels->end();  ++p) {
     Label* current = p->second;
@@ -367,7 +372,7 @@ std::unique_ptr<EncodedProgram> AssemblyProgram::Encode() const {
 
   encoded->set_image_base(image_base_);
 
-  if (!encoded->DefineLabels(abs32_labels_, rel32_labels_))
+  if (!encoded->ImportLabels(abs32_label_manager_, rel32_label_manager_))
     return nullptr;
 
   for (size_t i = 0;  i < instructions_.size();  ++i) {
@@ -470,61 +475,6 @@ Instruction* AssemblyProgram::GetByteInstruction(uint8_t byte) {
   return byte_instruction_cache_[byte];
 }
 
-// Chosen empirically to give the best reduction in payload size for
-// an update from daisy_3701.98.0 to daisy_4206.0.0.
-const int AssemblyProgram::kLabelLowerLimit = 5;
-
-CheckBool AssemblyProgram::TrimLabels() {
-  // For now only trim for ARM binaries.
-  if (kind() != EXE_ELF_32_ARM)
-    return true;
-
-  int lower_limit = kLabelLowerLimit;
-
-  VLOG(1) << "TrimLabels: threshold " << lower_limit;
-
-  // Walk through the list of instructions, replacing trimmed labels
-  // with the original machine instruction.
-  for (size_t i = 0; i < instructions_.size(); ++i) {
-    Instruction* instruction = instructions_[i];
-    switch (instruction->op()) {
-      case REL32ARM: {
-        Label* label =
-            static_cast<InstructionWithLabelARM*>(instruction)->label();
-        if (label->count_ <= lower_limit) {
-          const uint8_t* arm_op =
-              static_cast<InstructionWithLabelARM*>(instruction)->arm_op();
-          uint16_t op_size =
-              static_cast<InstructionWithLabelARM*>(instruction)->op_size();
-
-          if (op_size < 1)
-            return false;
-          UncheckedDelete(instruction);
-          instructions_[i] = UncheckedNew<BytesInstruction>(arm_op, op_size);
-          if (!instructions_[i])
-            return false;
-        }
-        break;
-      }
-      default:
-        break;
-    }
-  }
-
-  // Remove and deallocate underused Labels.
-  RVAToLabel::iterator it = rel32_labels_.begin();
-  while (it != rel32_labels_.end()) {
-    if (it->second->count_ <= lower_limit) {
-      UncheckedDelete(it->second);
-      rel32_labels_.erase(it++);
-    } else {
-      ++it;
-    }
-  }
-
-  return true;
-}
-
 ////////////////////////////////////////////////////////////////////////////////
 
 Status Encode(const AssemblyProgram& program,