Add abilities to generate bitcode buffers from munged bitcode.

include/llvm/Bitcode/NaCl/NaClBitcodeMungeUtils.h

 //===- NaClBitcodeMungeUtils.h - Munge bitcode records --------*- C++ -*-===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 //
 // This file defines utility class NaClMungedBitcode to edit a base
@@ skipping 63 matching lines @@
 #ifndef LLVM_BITCODE_NACL_NACLBITCODEMUNGEUTILS_H
 #define LLVM_BITCODE_NACL_NACLBITCODEMUNGEUTILS_H
 
 #include "llvm/Bitcode/NaCl/NaClBitcodeParser.h"
 
 #include <list>
 #include <map>
 
 namespace llvm {
 
+class MemoryBuffer;            // Buffer to read bitcode from
 class NaClBitcodeAbbrevRecord; // bitcode record.
 class NaClMungedBitcodeIter;   // iterator over edited bitcode records.
 
 /// \brief Defines a list of bitcode records.
 typedef std::vector<std::unique_ptr<NaClBitcodeAbbrevRecord>>
     NaClBitcodeRecordList;
 
 /// \brief Extracts out the records in Records, and puts them into RecordList.
 ///
 /// \brief RecordList[in/out] Record list to read into.
 /// \brief Records Array containing data defining records.
 /// \brief RecordsSize The size of array Records.
 /// \brief RecordTerminator The value used to terminate records.
 void readNaClBitcodeRecordList(NaClBitcodeRecordList &RecordList,
                                const uint64_t Records[],
                                size_t RecordsSize,
                                uint64_t RecordTerminator);
 
+/// \brief Read in the list of records from bitcode in a memory buffer.
+void readNaClBitcodeRecordList(NaClBitcodeRecordList &RecordList,
+                               std::unique_ptr<MemoryBuffer> InputBuffer);
+
 /// \brief An edited (i.e. munged) list of bitcode records. Edits are
 /// always relative to the initial list of records.
 class NaClMungedBitcode {
   friend class NaClMungedBitcodeIter;
   NaClMungedBitcode() = delete;
   NaClMungedBitcode(const NaClMungedBitcode &) = delete;
   NaClMungedBitcode &operator=(const NaClMungedBitcode &) = delete;
 
 public:
   /// \brief Iterator over edited records.
   typedef NaClMungedBitcodeIter iterator;
 
+  /// \brief Read in initial list of records from bitcode in a memory buffer.
+  explicit NaClMungedBitcode(std::unique_ptr<MemoryBuffer> InputBuffer);
+
   /// \brief Initialize the list of records to be edited.
   explicit NaClMungedBitcode(std::unique_ptr<NaClBitcodeRecordList> BaseRecords)
       : BaseRecords(std::move(BaseRecords)) {}
 
   /// \brief Initialize the list of records to be edited using
   /// array specification.
   ///
   /// \brief Records Array containing data defining records.
   /// \brief RecordsSize The size of array Records.
   /// \brief RecordTerminator The value used to terminate records.
@@ skipping 29 matching lines @@
 
   /// \brief Replace the record at RecordIndex with Record. Note that
   /// because the record index is based on the base records, and those
   /// records are not modified, this editing action effectively undoes
   /// all previous remove/replace editing actions for this index.
   void replace(size_t RecordIndex, NaClBitcodeAbbrevRecord &Record);
 
   /// \brief Print out the resulting edited list of records.
   void print(raw_ostream &Out) const;
 
+  /// Defines set of possible write flags.
+  struct WriteFlags {
+    /// True if error recovery should be applied.
+    bool getTryToRecover() const { return TryToRecover; }
+
+    /// Define that error recovery should be applied when writing.
+    void setTryToRecover(bool NewValue) {
+      TryToRecover = NewValue;
+      assert(!(TryToRecover && WriteBadAbbrevIndex));
+    }
+
+    /// True if a bad abbreviation index should be written (rather than
+    /// trying error recovery) so that bitcode readers can be tested for
+    /// this condition.
+    bool getWriteBadAbbrevIndex() const { return WriteBadAbbrevIndex; }
+
+    /// Define that the first bad abbreviation index should be written,
+    /// and corresponding minimal context added so that the bitcode can
+    /// be used to test reading the erroneous written bitcode.
+    void setWriteBadAbbrevIndex(bool NewValue) {
+      WriteBadAbbrevIndex = NewValue;
+      assert(!(TryToRecover && WriteBadAbbrevIndex));
+    }
+
+    /// Get the stream to print errors while writing bitcode.
+    raw_ostream &getErrStream() const {
+      return ErrStream ? *ErrStream : errs();
+    }
+
+    /// Set the stream to print errors to.
+    void setErrStream(raw_ostream &NewValue) {
+      ErrStream = &NewValue;
+    }
+
+    void reset() {
+      TryToRecover = false;
+      WriteBadAbbrevIndex = false;
+      ErrStream = nullptr;
+    }
+
+  private:
+    bool TryToRecover = false;
+    bool WriteBadAbbrevIndex = false;
+    raw_ostream *ErrStream = nullptr;
+  };
+
+  /// Defines the results associated with writing bitcode.
+  struct WriteResults {
+    /// Number of errors generated.
+    size_t NumErrors = 0;
+    /// Number of repairs (via error recovery) that were applied.
+    size_t NumRepairs = 0;
+    /// True if a bad abbreviation index were written.
+    bool WroteBadAbbrevIndex = false;
+  };
+
+  /// \brief Write out the edited list of bitcode records using
+  /// the given buffer.
+  ///
+  /// \param Buffer The buffer to write into.
+  /// \param AddHeader Add header block when true.
+  /// \param Flags Write flags to use.
+  ///
+  /// \return Returns the results of the write.
+  WriteResults writeMaybeRepair(
+      SmallVectorImpl<char> &Buffer, bool AddHeader,
+      const WriteFlags &Flags) const;
+
+  bool write(SmallVectorImpl<char> &Buffer, bool AddHeader,
+             const WriteFlags &Flags) const {
+    return writeMaybeRepair(Buffer, AddHeader, Flags).NumErrors == 0;
+  }
+
+  bool write(SmallVectorImpl<char> &Buffer, bool AddHeader) const {
+    WriteFlags Flags;
+    return write(Buffer, AddHeader, Flags);
+  }
+
   /// \brief The types of editing actions that can be applied.
   enum EditAction {
     AddBefore, // Insert new record before base record at index.
     AddAfter,  // Insert new record after base record at index.
     Remove,    // Remove record at index.
     Replace    // Replace base record at index with new record.
   };
 
   /// \brief Apply a set of edits defined in the given array.
   ///
@@ skipping 10 matching lines @@
   ///
   /// \param Munges The array containing the edits to apply.
   /// \param MungesSize The size of Munges.
   /// \param Terminator The value used to terminate records in editing actions.
   void munge(const uint64_t Munges[], size_t MungesSize, uint64_t Terminator);
 
   /// \brief Removes all editing actions and resets back to the original
   /// set of base records.
   void removeEdits();
 
+  /// Returns the unedited list of bitcode records.
+  const NaClBitcodeRecordList &getBaseRecords() const {
+    return *BaseRecords;
+  }
+
 private:
   typedef std::list<NaClBitcodeAbbrevRecord *> RecordListType;
   typedef std::map<size_t, RecordListType *> InsertionsMapType;
   typedef std::map<size_t, NaClBitcodeAbbrevRecord *> ReplaceMapType;
 
   /// \brief The list of base records that will be edited.
   std::unique_ptr<NaClBitcodeRecordList> BaseRecords;
   // Holds map from record index to list of records added before
   // the corresponding record in the list of base records.
   InsertionsMapType BeforeInsertionsMap;
@@ skipping 31 matching lines @@
 public:
   /// \brief The abbreviation associated with the bitcode record.
   unsigned Abbrev;
 
   /// \brief Creates a bitcode record.
   ///
   /// \param Abbrev Abbreviation index associated with record.
   /// \param Code The selector code of the record.
   /// \param Values The values associated with the selector code.
   NaClBitcodeAbbrevRecord(unsigned Abbrev, unsigned Code,
-                          NaClRecordVector &Values)
+                          const NaClRecordVector &Values)
       : NaClBitcodeRecordData(Code, Values), Abbrev(Abbrev) {}
 
   /// \brief Creates a copy of the given abbreviated bitcode record.
   explicit NaClBitcodeAbbrevRecord(const NaClBitcodeAbbrevRecord &Rcd)
       : NaClBitcodeRecordData(Rcd), Abbrev(Rcd.Abbrev) {}
 
   /// \brief Creates a default abbreviated bitcode record.
   NaClBitcodeAbbrevRecord() : Abbrev(naclbitc::UNABBREV_RECORD) {}
 
   ~NaClBitcodeAbbrevRecord() {}
@@ skipping 106 matching lines @@
   }
 
   // \brief Moves the iterator to the position of the next edited
   // record.
   void updatePosition();
 };
 
 } // end namespace llvm.
 
 #endif // LLVM_BITCODE_NACL_NACLBITCODEMUNGE_H