Switch out use of std::string and std::vector for large allocations for a buffer class that doesn...

courgette/assembly_program.cc

 // 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.
 
 #include "courgette/assembly_program.h"
 
 #include <memory.h>
 #include <algorithm>
 #include <map>
 #include <set>
@@ skipping 68 matching lines @@
     if (label == NULL) NOTREACHED();
   }
   Label* label() const { return label_; }
  private:
   Label* label_;
 };
 
 }  // namespace
 
 AssemblyProgram::AssemblyProgram()
-  : byte_instruction_cache_(NULL),
-    image_base_(0) {
+  : image_base_(0) {
 }
 
 static void DeleteContainedLabels(const RVAToLabel& labels) {
   for (RVAToLabel::const_iterator p = labels.begin();  p != labels.end();  ++p)
     delete p->second;
 }
 
 AssemblyProgram::~AssemblyProgram() {
   for (size_t i = 0;  i < instructions_.size();  ++i) {
     Instruction* instruction = instructions_[i];
     if (instruction->op() != DEFBYTE)  // Will be in byte_instruction_cache_.
       delete instruction;
   }
-  if (byte_instruction_cache_) {
+  if (byte_instruction_cache_.get()) {
     for (size_t i = 0;  i < 256;  ++i)
       delete byte_instruction_cache_[i];
-    delete[] byte_instruction_cache_;
   }
   DeleteContainedLabels(rel32_labels_);
   DeleteContainedLabels(abs32_labels_);
 }
 
-void AssemblyProgram::EmitMakeRelocsInstruction() {
-  Emit(new MakeRelocsInstruction());
+CheckBool AssemblyProgram::EmitMakeRelocsInstruction() {
+  return Emit(new(std::nothrow) MakeRelocsInstruction());
 }
 
-void AssemblyProgram::EmitOriginInstruction(RVA rva) {
-  Emit(new OriginInstruction(rva));
+CheckBool AssemblyProgram::EmitOriginInstruction(RVA rva) {
+  return Emit(new(std::nothrow) OriginInstruction(rva));
 }
 
-void AssemblyProgram::EmitByteInstruction(uint8 byte) {
-  Emit(GetByteInstruction(byte));
+CheckBool AssemblyProgram::EmitByteInstruction(uint8 byte) {
+  return Emit(GetByteInstruction(byte));
 }
 
-void AssemblyProgram::EmitRel32(Label* label) {
-  Emit(new InstructionWithLabel(REL32, label));
+CheckBool AssemblyProgram::EmitRel32(Label* label) {
+  return Emit(new(std::nothrow) InstructionWithLabel(REL32, label));
 }
 
-void AssemblyProgram::EmitAbs32(Label* label) {
-  Emit(new InstructionWithLabel(ABS32, label));
+CheckBool AssemblyProgram::EmitAbs32(Label* label) {
+  return Emit(new(std::nothrow) InstructionWithLabel(ABS32, label));
 }
 
 Label* AssemblyProgram::FindOrMakeAbs32Label(RVA rva) {
   return FindLabel(rva, &abs32_labels_);
 }
 
 Label* AssemblyProgram::FindOrMakeRel32Label(RVA rva) {
   return FindLabel(rva, &rel32_labels_);
 }
 
@@ skipping 19 matching lines @@
   return NULL;
 }
 
 Label* AssemblyProgram::InstructionRel32Label(
     const Instruction* instruction) const {
   if (instruction->op() == REL32)
     return static_cast<const InstructionWithLabel*>(instruction)->label();
   return NULL;
 }
 
+CheckBool AssemblyProgram::Emit(Instruction* instruction) {
+  if (!instruction)
+    return false;
+  bool ok = instructions_.push_back(instruction);
+  if (!ok)
+    delete instruction;
+  return ok;
+}
+
 Label* AssemblyProgram::FindLabel(RVA rva, RVAToLabel* labels) {
   Label*& slot = (*labels)[rva];
-  if (slot == 0) {
-    slot = new Label(rva);
+  if (slot == NULL) {
+    slot = new(std::nothrow) Label(rva);
   }
   return slot;
 }
 
 void AssemblyProgram::UnassignIndexes(RVAToLabel* labels) {
   for (RVAToLabel::iterator p = labels->begin();  p != labels->end();  ++p) {
     Label* current = p->second;
     current->index_ = Label::kNoIndex;
   }
 }
@@ skipping 116 matching lines @@
   for (RVAToLabel::const_iterator p = labels.begin();
        ok && p != labels.end();
        ++p) {
     Label* label = p->second;
     ok = (encoded_format->*define_label)(label->index_, label->rva_);
   }
   return ok;
 }
 
 EncodedProgram* AssemblyProgram::Encode() const {
-  scoped_ptr<EncodedProgram> encoded(new EncodedProgram());
+  scoped_ptr<EncodedProgram> encoded(new(std::nothrow) EncodedProgram());
+  if (!encoded.get())
+    return NULL;
+
   encoded->set_image_base(image_base_);
 
   if (!DefineLabels(abs32_labels_, encoded.get(),
                     &EncodedProgram::DefineAbs32Label) ||
       !DefineLabels(rel32_labels_, encoded.get(),
                     &EncodedProgram::DefineRel32Label)) {
     return NULL;
   }
 
   encoded->EndLabels();
@@ skipping 34 matching lines @@
       default: {
         NOTREACHED() << "Unknown Insn OP kind";
       }
     }
   }
 
   return encoded.release();
 }
 
 Instruction* AssemblyProgram::GetByteInstruction(uint8 byte) {
-  if (!byte_instruction_cache_) {
-    byte_instruction_cache_ = new Instruction*[256];
+  if (!byte_instruction_cache_.get()) {
+    byte_instruction_cache_.reset(new(std::nothrow) Instruction*[256]);
+    if (!byte_instruction_cache_.get())
+      return NULL;
+
     for (int i = 0; i < 256; ++i) {
-      byte_instruction_cache_[i] = new ByteInstruction(static_cast<uint8>(i));
+      byte_instruction_cache_[i] =
+          new(std::nothrow) ByteInstruction(static_cast<uint8>(i));
+      if (!byte_instruction_cache_[i]) {
+        for (int j = 0; j < i; ++j)
+          delete byte_instruction_cache_[j];
+        byte_instruction_cache_.reset();
+        return NULL;
+      }
     }
   }
 
   return byte_instruction_cache_[byte];
 }
 
 ////////////////////////////////////////////////////////////////////////////////
 
 Status Encode(AssemblyProgram* program, EncodedProgram** output) {
   *output = NULL;
   EncodedProgram *encoded = program->Encode();
   if (encoded) {
     *output = encoded;
     return C_OK;
   } else {
     return C_GENERAL_ERROR;
   }
 }
 
 }  // namespace courgette