Add abilities to generate bitcode buffers from munged bitcode.

lib/Bitcode/NaCl/TestUtils/NaClBitcodeMungeWriter.cpp

Index: lib/Bitcode/NaCl/TestUtils/NaClBitcodeMungeWriter.cpp
diff --git a/lib/Bitcode/NaCl/TestUtils/NaClBitcodeMunge.cpp b/lib/Bitcode/NaCl/TestUtils/NaClBitcodeMungeWriter.cpp
similarity index 51%
copy from lib/Bitcode/NaCl/TestUtils/NaClBitcodeMunge.cpp
copy to lib/Bitcode/NaCl/TestUtils/NaClBitcodeMungeWriter.cpp
index 2d88585ab5c81813db593a15b3473cc1d74627f9..76c83859b12e7b46f17d1c705c3e8ceba6755537 100644
--- a/lib/Bitcode/NaCl/TestUtils/NaClBitcodeMunge.cpp
+++ b/lib/Bitcode/NaCl/TestUtils/NaClBitcodeMungeWriter.cpp
@@ -1,4 +1,4 @@
-//===--- Bitcode/NaCl/TestUtils/NaClBitcodeMunge.cpp - Bitcode Munger -----===//
+//===- NaClBitcodeMungeWriter.cpp - Write munged bitcode --------*- C++ -*-===//
 //
 //                     The LLVM Compiler Infrastructure
 //
@@ -7,121 +7,66 @@
 //
 //===----------------------------------------------------------------------===//
 //
-// Bitcode writer/munger implementation for testing.
-//
-//===----------------------------------------------------------------------===//
+// Implements method NaClMungedBitcode.write(), which writes out a munged
+// list of bitcode records using a bitstream writer.
 
-#include "llvm/Bitcode/NaCl/NaClBitcodeMunge.h"
+#include "llvm/Bitcode/NaCl/NaClBitcodeMungeUtils.h"
 
-#include "llvm/Bitcode/NaCl/NaClBitcodeHeader.h"
-#include "llvm/Bitcode/NaCl/NaClBitcodeParser.h"
 #include "llvm/Bitcode/NaCl/NaClBitstreamWriter.h"
-#include "llvm/Bitcode/NaCl/NaClCompress.h"
 #include "llvm/Bitcode/NaCl/NaClReaderWriter.h"
-#include "llvm/IR/LLVMContext.h"
-#include "llvm/IR/Module.h"
-#include "llvm/Support/ErrorHandling.h"
-#include "llvm/Support/MemoryBuffer.h"
-
-#include <memory>
 
 using namespace llvm;
 
+namespace {
+
 // For debugging. When true, shows each PNaCl record that is
 // emitted to the bitcode file.
 static bool DebugEmitRecord = false;
 
-void NaClBitcodeMunger::setupTest(
-    const char *TestName, const uint64_t Munges[], size_t MungesSize,
-    bool AddHeader) {
-  assert(DumpStream == nullptr && "Test run with DumpStream already defined");
-  assert(MungedInput.get() == nullptr
-         && "Test run with MungedInput already defined");
-  FoundErrors = false;
-  DumpResults.clear(); // Throw away any previous results.
-  std::string DumpBuffer;
-  DumpStream = new raw_string_ostream(DumpResults);
-  MungedInputBuffer.clear();
-  NaClBitstreamWriter OutStream(MungedInputBuffer);
-  Writer = &OutStream;
-
-  if (DebugEmitRecord) {
-    errs() << "*** Run test: " << TestName << "\n";
-  }
-
-  WriteBlockID = -1;
-  SetBID = -1;
-
-  MungedBitcode.munge(Munges, MungesSize, RecordTerminator);
-  writeMungedBitcode(MungedBitcode, AddHeader);
+// State of bitcode writing.
+struct WriteState {
+  // The block ID associated with the block being written.
+  int WriteBlockID = -1;
+  // The SetBID for the blockinfo block.
+  int SetBID = -1;
+  // The stack of maximum abbreviation indices allowed by block enter record.
+  SmallVector<uint64_t, 3> AbbrevIndexLimitStack;
+  // The buffer to hold the generated fatal message.
+  std::string FatalBuffer;
+  // The stream to write the fatal message to.
+  raw_string_ostream FatalStream;
 
-  // Add null terminator, so that we meet the requirements of the
-  // MemoryBuffer API.
-  MungedInputBuffer.push_back('\0');
+  WriteState() :  FatalBuffer(), FatalStream(FatalBuffer) {}
 
-  MungedInput = MemoryBuffer::getMemBuffer(
-      StringRef(MungedInputBuffer.data(), MungedInputBuffer.size()-1),
-      TestName);
-}
-
-void NaClBitcodeMunger::cleanupTest() {
-  MungedBitcode.removeEdits();
-  MungedInput.reset();
-  assert(DumpStream && "Dump stream removed before cleanup!");
-  DumpStream->flush();
-  delete DumpStream;
-  DumpStream = nullptr;
-  Writer = nullptr;
-}
-
-// Return the next line of input (including eoln), starting from
-// Pos. Then increment Pos past the end of that line.
-static std::string getLine(const std::string &Input, size_t &Pos) {
-  std::string Line;
-  if (Pos >= Input.size()) {
-    Pos = std::string::npos;
-    return Line;
-  }
-  size_t Eoln = Input.find_first_of("\n", Pos);
-  if (Eoln != std::string::npos) {
-    for (size_t i = Pos; i <= Eoln; ++i)
-      Line.push_back(Input[i]);
-    Pos = Eoln + 1;
-    return Line;
+  // Returns stream to print fatal error message to.
+  // Note: Once the fatal error message has been dumped to the stream,
+  // one must call ReportFatalError to display the error and terminate
+  // execution.
+  raw_ostream &Fatal() {
+    return FatalStream << "Fatal: ";
   }
-  Pos = std::string::npos;
-  return Input.substr(Pos);
-}
 
-std::string NaClBitcodeMunger::
-getLinesWithTextMatch(const std::string &Substring, bool MustBePrefix) const {
-  std::string Messages;
-  size_t LastLinePos = 0;
-  while (1) {
-    std::string Line = getLine(DumpResults, LastLinePos);
-    if (LastLinePos == std::string::npos) break;
-    size_t Pos = Line.find(Substring);
-    if (Pos != std::string::npos && (!MustBePrefix || Pos == 0)) {
-      Messages.append(Line);
-    }
+  // Displays the fatal error message and exits the program.
+  void ReportFatalError() {
+    report_fatal_error(FatalStream.str());
   }
-  return Messages;
-}
 
-void NaClBitcodeMunger::writeMungedBitcode(const NaClMungedBitcode &Bitcode,
-                                           bool AddHeader) {
-  if (AddHeader) {
-    NaClWriteHeader(*Writer, true);
-  }
-  for (const NaClBitcodeAbbrevRecord &Record : Bitcode) {
-    emitRecord(Record);
-  }
-}
+  // Converts the abbreviation record to the corresponding abbreviation.
+  NaClBitCodeAbbrev *buildAbbrev(const NaClBitcodeAbbrevRecord &Record);
 
-void NaClBitcodeMunger::emitRecord(const NaClBitcodeAbbrevRecord &Record) {
+  // Emits the given record to the bitcode file.
+  void emitRecord(NaClBitstreamWriter &Writer,
+                  const NaClBitcodeAbbrevRecord &Record,
+                  bool RecoverOnBadAbbrevId);
+};
+
+void WriteState::emitRecord(NaClBitstreamWriter &Writer,
+                            const NaClBitcodeAbbrevRecord &Record,
+                            bool RecoverOnBadAbbrevId) {
+  size_t NumValues = Record.Values.size();
   if (DebugEmitRecord) {
     errs() << "Emit " << Record.Abbrev << ": <" << Record.Code;
-    for (size_t i = 0, e = Record.Values.size(); i < e; ++i) {
+    for (size_t i = 0; i < NumValues; ++i) {
       errs() << ", " << Record.Values[i];
     }
     errs() << ">\n";
@@ -136,7 +81,7 @@ void NaClBitcodeMunger::emitRecord(const NaClBitcodeAbbrevRecord &Record) {
               << " uses illegal abbreviation index " << Record.Abbrev << "\n";
       ReportFatalError();
     }
-    if (Record.Values.size() == 2) {
+    if (NumValues == 2) {
       WriteBlockID = Record.Values[0];
       NumBits = Record.Values[1];
       if (NumBits > 32) {
@@ -151,22 +96,24 @@ void NaClBitcodeMunger::emitRecord(const NaClBitcodeAbbrevRecord &Record) {
       }
     } else {
       Fatal() << "Error: Values for enter record should be of size 2. Found: "
-              << Record.Values.size();
+              << NumValues;
       ReportFatalError();
     }
     uint64_t MaxAbbrev = (static_cast<uint64_t>(1) << NumBits) - 1;
     AbbrevIndexLimitStack.push_back(MaxAbbrev);
     if (WriteBlockID == naclbitc::BLOCKINFO_BLOCK_ID) {
-      if (NumBits != naclbitc::DEFAULT_MAX_ABBREV) {
+      unsigned DefaultMaxBits =

[jvoung (off chromium), 2015/05/05 00:27:23]

Hmm bug fix / got lucky before?

[Karl, 2015/05/05 22:38:04]

Yep. Because the test examples didn't really use abbreviations, this problem
didn't get caught. The new test case uses abbreviations and caused the discovery
of the problem.

+          NaClBitsNeededForValue(naclbitc::DEFAULT_MAX_ABBREV);
+      if (NumBits != DefaultMaxBits) {
         Fatal()
-            << "Error: Numbits entry for abbreviations record not 2. Found: "
-            << NumBits << "\n";
+            << "Error: Numbits entry for abbreviations record not "
+            << DefaultMaxBits << ". Found: " << NumBits << "\n";
         ReportFatalError();
       }
-      Writer->EnterBlockInfoBlock();
+      Writer.EnterBlockInfoBlock();
     } else {
       NaClBitcodeSelectorAbbrev CurCodeLen(MaxAbbrev);
-      Writer->EnterSubblock(WriteBlockID, CurCodeLen);
+      Writer.EnterSubblock(WriteBlockID, CurCodeLen);
     }
     return;
   }
@@ -176,14 +123,14 @@ void NaClBitcodeMunger::emitRecord(const NaClBitcodeAbbrevRecord &Record) {
               << " uses illegal abbreviation index " << Record.Abbrev << "\n";
       ReportFatalError();
     }
-    if (!Record.Values.empty()) {
+    if (NumValues != 0) {
       Fatal() << "Error: Exit block should not have values. Found: "
               << Record.Values.size() << "\n";
       ReportFatalError();
     }
     if (!AbbrevIndexLimitStack.empty())
       AbbrevIndexLimitStack.pop_back();
-    Writer->ExitBlock();
+    Writer.ExitBlock();
     return;
   case naclbitc::BLK_CODE_DEFINE_ABBREV: {
     if (Record.Abbrev != naclbitc::DEFINE_ABBREV) {
@@ -194,27 +141,33 @@ void NaClBitcodeMunger::emitRecord(const NaClBitcodeAbbrevRecord &Record) {
     NaClBitCodeAbbrev *Abbrev = buildAbbrev(Record);
     if (Abbrev == NULL) return;
     if (WriteBlockID == naclbitc::BLOCKINFO_BLOCK_ID) {
-      Writer->EmitBlockInfoAbbrev(SetBID, Abbrev);
+      Writer.EmitBlockInfoAbbrev(SetBID, Abbrev);
     } else {
-      Writer->EmitAbbrev(Abbrev);
+      Writer.EmitAbbrev(Abbrev);
     }
     return;
   }
   case naclbitc::BLK_CODE_HEADER:
     // Note: There is no abbreviation index here. Ignore.
     for (uint64_t Value : Record.Values)
-      Writer->Emit(Value, 8);
+      Writer.Emit(Value, 8);
     return;
   default:
     if ((Record.Abbrev != naclbitc::UNABBREV_RECORD
-         && !Writer->isUserRecordAbbreviation(Record.Abbrev))) {
+         && !Writer.isUserRecordAbbreviation(Record.Abbrev))) {
       uint64_t BlockAbbrevIndexLimit = 0;
       if (!AbbrevIndexLimitStack.empty())
         BlockAbbrevIndexLimit = AbbrevIndexLimitStack.back();
-      if (Record.Abbrev > BlockAbbrevIndexLimit) {
-        Fatal() << "Error: Record code " << Record.Code
-                << " uses illegal abbreviation index " << Record.Abbrev
-                << ". Must not exceed " << BlockAbbrevIndexLimit << "\n";
+      if (!RecoverOnBadAbbrevId || Record.Abbrev > BlockAbbrevIndexLimit) {
+        // Can't generate abbreviation ID because not enough bits
+        // specified in block ID.
+        raw_ostream &ErrStrm = Fatal();
+        ErrStrm << "Error: Block ID " << WriteBlockID
+                << " Record code " << Record.Code
+                << " uses illegal abbreviation index " << Record.Abbrev;
+        if (Record.Abbrev > BlockAbbrevIndexLimit)
+          ErrStrm << ". Must not exceed " << BlockAbbrevIndexLimit;
+        ErrStrm << "\n";
         ReportFatalError();
       }
       // Note: If this point is reached, the abbreviation is
@@ -225,35 +178,48 @@ void NaClBitcodeMunger::emitRecord(const NaClBitcodeAbbrevRecord &Record) {
       // abbreviation pattern.
       errs() << "Warning: Record code " << Record.Code
              << " uses illegal abbreviation index " << Record.Abbrev << "\n";
-      Writer->EmitCode(Record.Abbrev);
-      Writer->EmitVBR(Record.Code, 6);
-      uint32_t NumValues = static_cast<uint32_t>(Record.Values.size());
-      Writer->EmitVBR(NumValues, 6);
+      Writer.EmitCode(Record.Abbrev);

[jvoung (off chromium), 2015/05/05 00:27:23]

Just revisiting this since it's been a while... this says "we output the record
data using the default abbreviation pattern"

Should this be Writer.EmitCode(naclbitc::UNABBREV_RECORD) instead or the idea is
that you go ahead and emit the bad abbrev, but the reader knows to recover by
following this format?

[Karl, 2015/05/05 22:38:04]

Actually, your close, but not quite right. The point of this code is to provide
a way to write out bitcode with a bad abbreviation value, so that we can test
the bitcode reader to see if it catches the problem. It probably is enough to
just generate the abbreviation code above, but I went ahead and added output
that looks close to what the abbreviation would look like.

I was letting the writer continue so that the proper number of exit blocks get
added (otherwise Writer fails on asserts). To make this clearer, I have modified
the code to stop writing as soon as something went wrong. This means that this
function now returns a boolean.

For this case, it generates the bad abbreviation index, followed by the proper
number of close blocks, and then returns false.

I have also modified buildAbbrev to return nullptr if an error occurs.

I have also modified Fatal() to record if an error occurred, rather than calling
ReportFatalError(). This allows the caller to NaClMungedBitcode::write to decide
what to do if there was a fatal error.

+      Writer.EmitVBR(Record.Code, 6);
+      Writer.EmitVBR(NumValues, 6);
       for (uint32_t i = 0; i < NumValues; ++i) {
-        Writer->EmitVBR64(Record.Values[i], 6);
+        Writer.EmitVBR64(Record.Values[i], 6);
+      }
+      return;
+    }
+    if (WriteBlockID == naclbitc::BLOCKINFO_BLOCK_ID
+        && Record.Code == naclbitc::BLOCKINFO_CODE_SETBID) {
+      // Note: SetBID records are handled by Writer->EmitBlockInfoAbbrev,
+      // based on the SetBID value. Don't bother to generate SetBID record here.
+      // Rathter just set SetBID and let call to Writer->EmitBlockInfoAbbrev

[jvoung (off chromium), 2015/05/05 00:27:23]

Rathter -> Rather

[Karl, 2015/05/05 22:38:04]

Done.

+      // generate the SetBID record.
+      if (NumValues != 1) {
+        Fatal() << "SetBID record expects 1 value, found: "
+                << NumValues << "\n";
+        ReportFatalError();
       }
+      if (NumValues > 0)
+        SetBID = Record.Values[0];
       return;
     }
     if (Record.Abbrev == naclbitc::UNABBREV_RECORD)
-      Writer->EmitRecord(Record.Code, Record.Values);
+      Writer.EmitRecord(Record.Code, Record.Values);
     else
-      Writer->EmitRecord(Record.Code, Record.Values, Record.Abbrev);
+      Writer.EmitRecord(Record.Code, Record.Values, Record.Abbrev);
     return;
   }
-  Fatal() << "emitRecord on unimplemented code" << "\n";
-  ReportFatalError();
 }
 
-NaClBitCodeAbbrev *NaClBitcodeMunger::buildAbbrev(
+// BUilds an abbreviation from the given bitcode record.

[jvoung (off chromium), 2015/05/05 00:27:23]

BUilds -> Builds

[Karl, 2015/05/05 22:38:04]

Removed comment since the declaration already has a comment.

+NaClBitCodeAbbrev *WriteState::buildAbbrev(
     const NaClBitcodeAbbrevRecord &Record) {
   NaClBitCodeAbbrev *Abbrev = new NaClBitCodeAbbrev();
   size_t Index = 0;
-  if (Record.Values.empty()) {
+  size_t NumValues = Record.Values.size();
+  if (NumValues == 0) {
     Fatal() << "Empty abbreviation record not allowed\n";
     ReportFatalError();
   }
   size_t NumAbbreviations = Record.Values[Index++];
-  size_t NumValues = Record.Values.size();
   if (NumAbbreviations == 0) {
     Fatal() << "Abbreviation must contain at least one operator\n";
     ReportFatalError();
@@ -321,49 +287,17 @@ NaClBitCodeAbbrev *NaClBitcodeMunger::buildAbbrev(
   return Abbrev;
 }
 
-bool NaClObjDumpMunger::runTestWithFlags(
-    const char *Name, const uint64_t Munges[], size_t MungesSize,
-    bool AddHeader, bool NoRecords, bool NoAssembly) {
-  setupTest(Name, Munges, MungesSize, AddHeader);
-
-  /// If running in death mode, redirect output directly to the
-  /// error stream (rather than buffering in DumpStream), so that
-  /// output can be seen in gtest death test.
-  raw_ostream &Output = RunAsDeathTest ? errs() : *DumpStream;
-  if (NaClObjDump(MungedInput.get()->getMemBufferRef(),
-                  Output, NoRecords, NoAssembly))
-    FoundErrors = true;
-  cleanupTest();
-  return !FoundErrors;
-}
+} // end of anonymous namespace.
 
-bool NaClParseBitcodeMunger::runTest(
-    const char *Name, const uint64_t Munges[], size_t MungesSize,
-    bool VerboseErrors) {
-  bool AddHeader = true;
-  setupTest(Name, Munges, MungesSize, AddHeader);
-  LLVMContext &Context = getGlobalContext();
-  raw_ostream *VerboseStrm = VerboseErrors ? DumpStream : nullptr;
-  ErrorOr<Module *> ModuleOrError =
-      NaClParseBitcodeFile(MungedInput->getMemBufferRef(), Context,
-                           VerboseStrm);
-  if (ModuleOrError) {
-    if (VerboseErrors)
-      *DumpStream << "Successful parse!\n";
-    delete ModuleOrError.get();
-  } else {
-    Error() << ModuleOrError.getError().message() << "\n";
+void NaClMungedBitcode::write(SmallVectorImpl<char> &Buffer,
+                              bool AddHeader,
+                              bool RecoverOnBadAbbrevId) const {
+  NaClBitstreamWriter Writer(Buffer);
+  WriteState State;
+  if (AddHeader) {
+    NaClWriteHeader(Writer, true);
+  }
+  for (const NaClBitcodeAbbrevRecord &Record : *this) {
+    State.emitRecord(Writer, Record, RecoverOnBadAbbrevId);

[jvoung (off chromium), 2015/05/05 00:27:23]

nit: Might have been able to make RecoverOnBadAbbrevId a field of State vs a
parameter here, but either way is fine.

[Karl, 2015/05/05 22:38:04]

Good point. The state is probably the better place to put it. Changing.

   }
-  cleanupTest();
-  return !FoundErrors;
-}
-
-bool NaClCompressMunger::runTest(const char* Name, const uint64_t Munges[],
-                                 size_t MungesSize) {
-  bool AddHeader = true;
-  setupTest(Name, Munges, MungesSize, AddHeader);
-  NaClBitcodeCompressor Compressor;
-  bool Result = Compressor.compress(MungedInput.get(), *DumpStream);
-  cleanupTest();
-  return Result;
 }