Identifying call sites that need to handle out of memory situations in Courgette.

courgette/assembly_program.cc

 // Copyright (c) 2010 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>
 #include <sstream>
 #include <vector>
 
 #include "base/logging.h"
+#include "base/scoped_ptr.h"
 
 #include "courgette/courgette.h"
 #include "courgette/encoded_program.h"
 
 namespace courgette {
 
 // Opcodes of simple assembly language
 enum OP {
   ORIGIN,         // ORIGIN <rva> - set current address for assembly.
   MAKERELOCS,     // Generates a base relocation table.
@@ skipping 255 matching lines @@
       ++index;
       ++fill_infill_count;
     }
   }
 
   VLOG(1) << "  fill forward " << fill_forward_count
           << "  backward " << fill_backward_count
           << "  infill " << fill_infill_count;
 }
 
-typedef void (EncodedProgram::*DefineLabelMethod)(int index, RVA value);
+typedef CheckBool (EncodedProgram::*DefineLabelMethod)(int index, RVA value);
 
 #if defined(OS_WIN)
 __declspec(noinline)
 #endif
-static void DefineLabels(const RVAToLabel& labels,
-                         EncodedProgram* encoded_format,
-                         DefineLabelMethod define_label) {
-  for (RVAToLabel::const_iterator p = labels.begin(); p != labels.end(); ++p) {
+static CheckBool DefineLabels(const RVAToLabel& labels,
+                              EncodedProgram* encoded_format,
+                              DefineLabelMethod define_label) {
+  bool ok = true;
+  for (RVAToLabel::const_iterator p = labels.begin();
+       ok && p != labels.end();
+       ++p) {
     Label* label = p->second;
-    (encoded_format->*define_label)(label->index_, label->rva_);
+    ok = (encoded_format->*define_label)(label->index_, label->rva_);
   }
+  return ok;
 }
 
 EncodedProgram* AssemblyProgram::Encode() const {
-  EncodedProgram* encoded = new EncodedProgram();
+  scoped_ptr<EncodedProgram> encoded(new EncodedProgram());
+  encoded->set_image_base(image_base_);
 
-  encoded->set_image_base(image_base_);
-  DefineLabels(abs32_labels_, encoded, &EncodedProgram::DefineAbs32Label);
-  DefineLabels(rel32_labels_, encoded, &EncodedProgram::DefineRel32Label);
+  if (!DefineLabels(abs32_labels_, encoded.get(),
+                    &EncodedProgram::DefineAbs32Label) ||
+      !DefineLabels(rel32_labels_, encoded.get(),
+                    &EncodedProgram::DefineRel32Label)) {
+    return NULL;
+  }
+
   encoded->EndLabels();
 
   for (size_t i = 0;  i < instructions_.size();  ++i) {
     Instruction* instruction = instructions_[i];
 
     switch (instruction->op()) {
       case ORIGIN: {
         OriginInstruction* org = static_cast<OriginInstruction*>(instruction);
-        encoded->AddOrigin(org->origin_rva());
+        if (!encoded->AddOrigin(org->origin_rva()))
+          return NULL;
         break;
       }
       case DEFBYTE: {
         uint8 b = static_cast<ByteInstruction*>(instruction)->byte_value();
-        encoded->AddCopy(1, &b);
+        if (!encoded->AddCopy(1, &b))
+          return NULL;
         break;
       }
       case REL32: {
         Label* label = static_cast<InstructionWithLabel*>(instruction)->label();
-        encoded->AddRel32(label->index_);
+        if (!encoded->AddRel32(label->index_))
+          return NULL;
         break;
       }
       case ABS32: {
         Label* label = static_cast<InstructionWithLabel*>(instruction)->label();
-        encoded->AddAbs32(label->index_);
+        if (!encoded->AddAbs32(label->index_))
+          return NULL;
         break;
       }
       case MAKERELOCS: {
-        encoded->AddMakeRelocs();
+        if (!encoded->AddMakeRelocs())
+          return NULL;
         break;
       }
       default: {
         NOTREACHED() << "Unknown Insn OP kind";
       }
     }
   }
 
-  return encoded;
+  return encoded.release();
 }
 
 Instruction* AssemblyProgram::GetByteInstruction(uint8 byte) {
   if (!byte_instruction_cache_) {
     byte_instruction_cache_ = new Instruction*[256];
     for (int i = 0; i < 256; ++i) {
       byte_instruction_cache_[i] = new ByteInstruction(static_cast<uint8>(i));
     }
   }
 
   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