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

courgette/encoded_program.cc

Index: courgette/encoded_program.cc
===================================================================
--- courgette/encoded_program.cc	(revision 78901)
+++ courgette/encoded_program.cc	(working copy)
@@ -40,14 +40,15 @@
 
 // Serializes a vector of integral values using Varint32 coding.
 template<typename T, typename A>
-void WriteVector(const std::vector<T, A>& items, SinkStream* buffer) {
+CheckBool WriteVector(const std::vector<T, A>& items, SinkStream* buffer) {
   size_t count = items.size();
-  buffer->WriteSizeVarint32(count);
-  for (size_t i = 0;  i < count;  ++i) {
+  bool ok = buffer->WriteSizeVarint32(count);
+  for (size_t i = 0; ok && i < count;  ++i) {
     COMPILE_ASSERT(sizeof(T) <= sizeof(uint32),  // NOLINT
                    T_must_fit_in_uint32);
-    buffer->WriteSizeVarint32(items[i]);
+    ok = buffer->WriteSizeVarint32(items[i]);
   }
+  return ok;
 }
 
 template<typename T, typename A>
@@ -62,6 +63,7 @@
     uint32 item;
     if (!buffer->ReadVarint32(&item))
       return false;
+    // TODO(tommi): Handle errors.
     items->push_back(static_cast<T>(item));
   }
 
@@ -70,26 +72,29 @@
 
 // Serializes a vector, using delta coding followed by Varint32 coding.
 template<typename A>
-void WriteU32Delta(const std::vector<uint32, A>& set, SinkStream* buffer) {
+CheckBool WriteU32Delta(const std::vector<uint32, A>& set, SinkStream* buffer) {
   size_t count = set.size();
-  buffer->WriteSizeVarint32(count);
+  bool ok = buffer->WriteSizeVarint32(count);
   uint32 prev = 0;
-  for (size_t i = 0;  i < count;  ++i) {
+  for (size_t i = 0;  ok && i < count;  ++i) {
     uint32 current = set[i];
     uint32 delta = current - prev;
-    buffer->WriteVarint32(delta);
+    ok = buffer->WriteVarint32(delta);
     prev = current;
   }
+  return ok;
 }
 
 template <typename A>
-static bool ReadU32Delta(std::vector<uint32, A>* set, SourceStream* buffer) {
+static CheckBool ReadU32Delta(std::vector<uint32, A>* set,
+                              SourceStream* buffer) {
   uint32 count;
 
   if (!buffer->ReadVarint32(&count))
     return false;
 
   set->clear();
+  // TODO(tommi): Handle errors.
   set->reserve(count);
   uint32 prev = 0;
 
@@ -98,28 +103,31 @@
     if (!buffer->ReadVarint32(&delta))
       return false;
     uint32 current = prev + delta;
+    // TODO(tommi): handle errors
     set->push_back(current);
     prev = current;
   }
 
+  // TODO(tommi): Handle errors.
   return true;
 }
 
 // Write a vector as the byte representation of the contents.
 //
 // (This only really makes sense for a type T that has sizeof(T)==1, otherwise
-// serilized representation is not endian-agnositic.  But it is useful to keep
+// serialized representation is not endian-agnositic.  But it is useful to keep
 // the possibility of a greater size for experiments comparing Varint32 encoding
 // of a vector of larger integrals vs a plain form.)
 //
 template<typename T, typename A>
-void WriteVectorU8(const std::vector<T, A>& items, SinkStream* buffer) {
+CheckBool WriteVectorU8(const std::vector<T, A>& items, SinkStream* buffer) {
   size_t count = items.size();
-  buffer->WriteSizeVarint32(count);
-  if (count != 0) {
+  bool ok = buffer->WriteSizeVarint32(count);
+  if (count != 0 && ok) {
     size_t byte_count = count * sizeof(T);
-    buffer->Write(static_cast<const void*>(&items[0]), byte_count);
+    ok = buffer->Write(static_cast<const void*>(&items[0]), byte_count);
   }
+  return ok;
 }
 
 template<typename T, typename A>
@@ -129,6 +137,7 @@
     return false;
 
   items->clear();
+  // TODO(tommi): check error
   items->resize(count);
   if (count != 0) {
     size_t byte_count = count * sizeof(T);
@@ -139,26 +148,29 @@
 
 ////////////////////////////////////////////////////////////////////////////////
 
-void EncodedProgram::DefineRel32Label(int index, RVA value) {
-  DefineLabelCommon(&rel32_rva_, index, value);
+CheckBool EncodedProgram::DefineRel32Label(int index, RVA value) {
+  return DefineLabelCommon(&rel32_rva_, index, value);
 }
 
-void EncodedProgram::DefineAbs32Label(int index, RVA value) {
-  DefineLabelCommon(&abs32_rva_, index, value);
+CheckBool EncodedProgram::DefineAbs32Label(int index, RVA value) {
+  return DefineLabelCommon(&abs32_rva_, index, value);
 }
 
 static const RVA kUnassignedRVA = static_cast<RVA>(-1);
 
-void EncodedProgram::DefineLabelCommon(RvaVector* rvas,
-                                       int index,
-                                       RVA rva) {
+CheckBool EncodedProgram::DefineLabelCommon(RvaVector* rvas,
+                                            int index,
+                                            RVA rva) {
   if (static_cast<int>(rvas->size()) <= index) {
+    // TODO(tommi): handle error
     rvas->resize(index + 1, kUnassignedRVA);
   }
   if ((*rvas)[index] != kUnassignedRVA) {
     NOTREACHED() << "DefineLabel double assigned " << index;
   }
   (*rvas)[index] = rva;
+  // TODO(tommi): Handle errors
+  return true;
 }
 
 void EncodedProgram::EndLabels() {
@@ -181,12 +193,15 @@
   }
 }
 
-void EncodedProgram::AddOrigin(RVA origin) {
+CheckBool EncodedProgram::AddOrigin(RVA origin) {
+  //TODO(tommi): Handle errors
   ops_.push_back(ORIGIN);
   origins_.push_back(origin);
+  return true;
 }
 
-void EncodedProgram::AddCopy(uint32 count, const void* bytes) {
+CheckBool EncodedProgram::AddCopy(uint32 count, const void* bytes) {
+  //TODO(tommi): Handle errors
   const uint8* source = static_cast<const uint8*>(bytes);
 
   // Fold adjacent COPY instructions into one.  This nearly halves the size of
@@ -205,7 +220,7 @@
       for (uint32 i = 0;  i < count;  ++i) {
         copy_bytes_.push_back(source[i]);
       }
-      return;
+      return true;
     }
   }
 
@@ -219,20 +234,28 @@
       copy_bytes_.push_back(source[i]);
     }
   }
+
+  return true;
 }
 
-void EncodedProgram::AddAbs32(int label_index) {
+CheckBool EncodedProgram::AddAbs32(int label_index) {
+  //TODO(tommi): Handle errors
   ops_.push_back(ABS32);
   abs32_ix_.push_back(label_index);
+  return true;
 }
 
-void EncodedProgram::AddRel32(int label_index) {
+CheckBool EncodedProgram::AddRel32(int label_index) {
+  //TODO(tommi): Handle errors
   ops_.push_back(REL32);
   rel32_ix_.push_back(label_index);
+  return true;
 }
 
-void EncodedProgram::AddMakeRelocs() {
+CheckBool EncodedProgram::AddMakeRelocs() {
+  //TODO(tommi): Handle errors
   ops_.push_back(MAKE_BASE_RELOCATION_TABLE);
+  return true;
 }
 
 void EncodedProgram::DebuggingSummary() {
@@ -279,7 +302,7 @@
   return  static_cast<FieldSelect>(~0);
 }
 
-void EncodedProgram::WriteTo(SinkStreamSet* streams) {
+CheckBool EncodedProgram::WriteTo(SinkStreamSet* streams) {
   FieldSelect select = GetFieldSelect();
 
   // The order of fields must be consistent in WriteTo and ReadFrom, regardless
@@ -293,28 +316,46 @@
 
   if (select & INCLUDE_MISC) {
     // TODO(sra): write 64 bits.
-    streams->stream(kStreamMisc)->WriteVarint32(
-      static_cast<uint32>(image_base_));
+    if (!streams->stream(kStreamMisc)->WriteVarint32(
+            static_cast<uint32>(image_base_))) {
+      return false;
+    }
   }
 
-  if (select & INCLUDE_ABS32_ADDRESSES)
-    WriteU32Delta(abs32_rva_, streams->stream(kStreamAbs32Addresses));
-  if (select & INCLUDE_REL32_ADDRESSES)
-    WriteU32Delta(rel32_rva_, streams->stream(kStreamRel32Addresses));
+  bool success = true;
+
+  if (select & INCLUDE_ABS32_ADDRESSES) {
+    success &= WriteU32Delta(abs32_rva_,
+                             streams->stream(kStreamAbs32Addresses));
+  }
+
+  if (select & INCLUDE_REL32_ADDRESSES) {
+    success &= WriteU32Delta(rel32_rva_,
+                             streams->stream(kStreamRel32Addresses));
+  }
+
   if (select & INCLUDE_MISC)
-    WriteVector(origins_, streams->stream(kStreamOriginAddresses));
+    success &= WriteVector(origins_, streams->stream(kStreamOriginAddresses));
+
   if (select & INCLUDE_OPS) {
-    streams->stream(kStreamOps)->Reserve(ops_.size() + 5);  // 5 for length.
-    WriteVector(ops_, streams->stream(kStreamOps));
+    // 5 for length.
+    success &= streams->stream(kStreamOps)->Reserve(ops_.size() + 5);
+    success &= WriteVector(ops_, streams->stream(kStreamOps));
   }
+
   if (select & INCLUDE_COPY_COUNTS)
-    WriteVector(copy_counts_, streams->stream(kStreamCopyCounts));
+    success &= WriteVector(copy_counts_, streams->stream(kStreamCopyCounts));
+
   if (select & INCLUDE_BYTES)
-    WriteVectorU8(copy_bytes_, streams->stream(kStreamBytes));
+    success &= WriteVectorU8(copy_bytes_, streams->stream(kStreamBytes));
+
   if (select & INCLUDE_ABS32_INDEXES)
-    WriteVector(abs32_ix_, streams->stream(kStreamAbs32Indexes));
+    success &= WriteVector(abs32_ix_, streams->stream(kStreamAbs32Indexes));
+
   if (select & INCLUDE_REL32_INDEXES)
-    WriteVector(rel32_ix_, streams->stream(kStreamRel32Indexes));
+    success &= WriteVector(rel32_ix_, streams->stream(kStreamRel32Indexes));
+
+  return success;
 }
 
 bool EncodedProgram::ReadFrom(SourceStreamSet* streams) {
@@ -360,7 +401,7 @@
   return true;
 }
 
-bool EncodedProgram::AssembleTo(SinkStream* final_buffer) {
+CheckBool EncodedProgram::AssembleTo(SinkStream* final_buffer) {
   // For the most part, the assembly process walks the various tables.
   // ix_mumble is the index into the mumble table.
   size_t ix_origins = 0;
@@ -402,7 +443,8 @@
           if (!VectorAt(copy_bytes_, ix_copy_bytes, &b))
             return false;
           ++ix_copy_bytes;
-          output->Write(&b, 1);
+          if (!output->Write(&b, 1))
+            return false;
         }
         current_rva += count;
         break;
@@ -413,7 +455,8 @@
         if (!VectorAt(copy_bytes_, ix_copy_bytes, &b))
           return false;
         ++ix_copy_bytes;
-        output->Write(&b, 1);
+        if (!output->Write(&b, 1))
+          return false;
         current_rva += 1;
         break;
       }
@@ -427,7 +470,8 @@
         if (!VectorAt(rel32_rva_, index, &rva))
           return false;
         uint32 offset = (rva - (current_rva + 4));
-        output->Write(&offset, 4);
+        if (!output->Write(&offset, 4))
+          return false;
         current_rva += 4;
         break;
       }
@@ -442,7 +486,8 @@
           return false;
         uint32 abs32 = static_cast<uint32>(rva + image_base_);
         abs32_relocs_.push_back(current_rva);
-        output->Write(&abs32, 4);
+        if (!output->Write(&abs32, 4))
+          return false;
         current_rva += 4;
         break;
       }
@@ -471,8 +516,9 @@
   }
 
   if (pending_base_relocation_table) {
-    GenerateBaseRelocations(final_buffer);
-    final_buffer->Append(&bytes_following_base_relocation_table);
+    if (!GenerateBaseRelocations(final_buffer) ||
+        !final_buffer->Append(&bytes_following_base_relocation_table))
+      return false;
   }
 
   // Final verification check: did we consume all lists?
@@ -488,7 +534,6 @@
   return true;
 }
 
-
 // RelocBlock has the layout of a block of relocations in the base relocation
 // table file format.
 //
@@ -512,39 +557,45 @@
     pod.block_size += 2;
   }
 
-  void Flush(SinkStream* buffer) {
+  CheckBool Flush(SinkStream* buffer) {
+    bool ok = true;
     if (pod.block_size != 8) {
       if (pod.block_size % 4 != 0) {  // Pad to make size multiple of 4 bytes.
         Add(0);
       }
-      buffer->Write(&pod, pod.block_size);
+      ok = buffer->Write(&pod, pod.block_size);
       pod.block_size = 8;
     }
+    return ok;
   }
   RelocBlockPOD pod;
 };
 
-void EncodedProgram::GenerateBaseRelocations(SinkStream* buffer) {
+CheckBool EncodedProgram::GenerateBaseRelocations(SinkStream* buffer) {
   std::sort(abs32_relocs_.begin(), abs32_relocs_.end());
 
   RelocBlock block;
 
-  for (size_t i = 0;  i < abs32_relocs_.size();  ++i) {
+  bool ok = true;
+  for (size_t i = 0;  ok && i < abs32_relocs_.size();  ++i) {
     uint32 rva = abs32_relocs_[i];
     uint32 page_rva = rva & ~0xFFF;
     if (page_rva != block.pod.page_rva) {
-      block.Flush(buffer);
+      ok &= block.Flush(buffer);
       block.pod.page_rva = page_rva;
     }
-    block.Add(0x3000 | (rva & 0xFFF));
+    if (ok)
+      block.Add(0x3000 | (rva & 0xFFF));
   }
-  block.Flush(buffer);
+  ok &= block.Flush(buffer);
+  return ok;
 }
 
 ////////////////////////////////////////////////////////////////////////////////
 
 Status WriteEncodedProgram(EncodedProgram* encoded, SinkStreamSet* sink) {
-  encoded->WriteTo(sink);
+  if (!encoded->WriteTo(sink))
+    return C_STREAM_ERROR;
   return C_OK;
 }