Modify NaCl bitstream reader to allow parallel parses.

include/llvm/Bitcode/NaCl/NaClBitstreamReader.h

 //===- NaClBitstreamReader.h -----------------------------------*- C++ -*-===//
 //     Low-level bitstream reader interface
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 //
 // This header defines the BitstreamReader class.  This class can be used to
 // read an arbitrary bitstream, regardless of its contents.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_BITCODE_NACL_NACLBITSTREAMREADER_H
 #define LLVM_BITCODE_NACL_NACLBITSTREAMREADER_H
 
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Bitcode/NaCl/NaClBitcodeHeader.h"
 #include "llvm/Bitcode/NaCl/NaClLLVMBitCodes.h"
 #include "llvm/Support/Endian.h"
 #include "llvm/Support/StreamingMemoryObject.h"
+#include <atomic>
 #include <climits>
+#include <mutex>
+#include <unordered_map>
 #include <vector>
 
 namespace llvm {
 
 class Deserializer;
 class NaClBitstreamCursor;
 
 namespace naclbitc {
 
 /// Returns the Bit as a Byte:BitInByte string.
@@ skipping 88 matching lines @@
     BlockInfo(const BlockInfo&) = default;
     unsigned getBlockID() const { return BlockID; }
     void setBlockID(unsigned ID) { BlockID = ID; }
     AbbrevList &getAbbrevs() { return Abbrevs; }
     ~BlockInfo() {}
   private:
     unsigned BlockID;
     AbbrevList Abbrevs;
   };
 
+  // Holds the global abbreviations in the BlockInfo block of the bitcode file.
+  // Sharing is used to allow parallel parses. Share by using std::share_ptr's
+  // and std::shared_from_this().
+  //
+  // Note: The BlockInfo block must be parsed before sharing of the
+  // BlockInfoRecordsMap.  Therefore, before changing to a parallel parse, the
+  // BlockInfoRecordsMap must be frozen.  Failure to do so, can lead to
+  // unexpected behaviour.
+  //
+  // In practice, this means that only function blocks can be parsed in
+  // parallel.
+  class BlockInfoRecordsMap :
+      public std::enable_shared_from_this<BlockInfoRecordsMap> {
+    friend class NaClBitstreamReader;
+    BlockInfoRecordsMap(const BlockInfoRecordsMap&) = delete;
+    BlockInfoRecordsMap &operator=(const BlockInfoRecordsMap&) = delete;
+  public:
+    using InfosMap = std::unordered_map<unsigned, BlockInfo>;
+
+    static std::shared_ptr<BlockInfoRecordsMap> create() {
+      return std::shared_ptr<BlockInfoRecordsMap>(new BlockInfoRecordsMap());
+    }
+    ~BlockInfoRecordsMap() = default;
+
+    bool isFrozen() const {
+      return IsFrozen.load();
+    }
+
+    // Returns true if already frozen.
+    bool freeze() {
+      return IsFrozen.exchange(true);
+    }
+
+    BlockInfo &getBlockInfo(unsigned BlockID) {
+      auto Pos = Infos.find(BlockID);
+      if (Pos != Infos.end())
+        return Pos->second;
+      report_fatal_error("Invalid block ID: " + std::to_string(BlockID));
+    }
+
+    // Locks the BlockInfoRecordsMap for the lifetime of the UpdateLock.  Used
+    // to allow the parsing of a BlockInfo block, and install global
+    // abbreviations.
+    //
+    // Verifies that the BlockInfoRecordsMap didn't get frozen during the
+    // instance's lifetime as a safety precaution. That is, it checks that no
+    // bitstream reader was created to share the global abbreviations before the
+    // global abbreviations are defined.
+    class UpdateLock {
+      UpdateLock() = delete;
+      UpdateLock &operator=(const UpdateLock&) = delete;
+    public:
+      explicit UpdateLock(BlockInfoRecordsMap &BlockInfoRecords);
+      ~UpdateLock();
+    private:
+      // The BlockInfoRecordsMap to update.
+      BlockInfoRecordsMap &BlockInfoRecords;
+      // The locked mutex from BlockInfoRecordsMap;
+      std::unique_lock<std::mutex> Lock;
+    };
+
+  private:
+    // The set of known BlockInfo's.
+    InfosMap Infos;
+    // True if the known BlockInfo blocks are frozen (i.e. the bitstream reader
+    // will ignore the BlockInfo block).
+    std::atomic_bool IsFrozen;
+    // Lock to use to update this data structure.
+    std::mutex Lock;
+
+    BlockInfoRecordsMap();
+  };
+
 private:
   friend class NaClBitstreamCursor;
 
   std::unique_ptr<MemoryObject> BitcodeBytes;
 
-  std::vector<BlockInfo> BlockInfoRecords;
-
-  // True if the BlockInfo block has been read.
-  bool HasReadBlockInfoBlock = false;
+  std::shared_ptr<BlockInfoRecordsMap> BlockInfoRecords;
 
   /// \brief Holds the offset of the first byte after the header.
   size_t InitialAddress;
 
   // True if filler should be added to byte align records.
   bool AlignBitcodeRecords = false;
   NaClBitstreamReader(const NaClBitstreamReader&) = delete;
   void operator=(const NaClBitstreamReader&) = delete;
 
 
   void initFromHeader(NaClBitcodeHeader &Header) {
     InitialAddress = Header.getHeaderSize();
     AlignBitcodeRecords = Header.getAlignBitcodeRecords();
   }
 
 public:
   /// Read stream from sequence of bytes [Start .. End) after parsing
   /// the given bitcode header.
   NaClBitstreamReader(const unsigned char *Start, const unsigned char *End,
                       NaClBitcodeHeader &Header)
-      : BitcodeBytes(getNonStreamedMemoryObject(Start, End)) {
+      : BitcodeBytes(getNonStreamedMemoryObject(Start, End)),
+        BlockInfoRecords(BlockInfoRecordsMap::create()) {
     initFromHeader(Header);
   }
 
   /// Read stream from Bytes, after parsing the given bitcode header.
   NaClBitstreamReader(MemoryObject *Bytes, NaClBitcodeHeader &Header)
-      : BitcodeBytes(Bytes) {
+      : BitcodeBytes(Bytes), BlockInfoRecords(BlockInfoRecordsMap::create()) {
     initFromHeader(Header);
   }
 
   /// Read stream from bytes, starting at the given initial address.
   /// Provides simple API for unit testing.
   NaClBitstreamReader(MemoryObject *Bytes, size_t InitialAddress)
-      : BitcodeBytes(Bytes), InitialAddress(InitialAddress) {
+      : BitcodeBytes(Bytes), BlockInfoRecords(BlockInfoRecordsMap::create()),
+        InitialAddress(InitialAddress) {
   }
 
+  /// Read stream from sequence of bytes [Start .. End), using the global
+  /// abbreviations of the given bitstream reader. Assumes that [Start .. End)
+  /// is copied from Reader's memory object.
+  NaClBitstreamReader(const unsigned char *Start,
+                      const unsigned char *End, NaClBitstreamReader *Reader)
+      : BitcodeBytes(getNonStreamedMemoryObject(Start, End)),
+        BlockInfoRecords(Reader->BlockInfoRecords), InitialAddress(0)
+  { BlockInfoRecords->freeze(); }
+
   // Returns the memory object that is being read.
   MemoryObject &getBitcodeBytes() { return *BitcodeBytes; }
 
   ~NaClBitstreamReader() {}
 
   /// \brief Returns the initial address (after the header) of the input stream.
   size_t getInitialAddress() const {
     return InitialAddress;
   }
 
   //===--------------------------------------------------------------------===//
   // Block Manipulation
   //===--------------------------------------------------------------------===//
 
-  /// If there is block info for the specified ID, return it, otherwise return
-  /// null.
-  const BlockInfo *getBlockInfo(unsigned BlockID) const {
-    // Common case, the most recent entry matches BlockID.
-    if (!BlockInfoRecords.empty() &&
-        BlockInfoRecords.back().getBlockID() == BlockID)
-      return &BlockInfoRecords.back();
-
-    for (unsigned i = 0, e = static_cast<unsigned>(BlockInfoRecords.size());
-         i != e; ++i)
-      if (BlockInfoRecords[i].getBlockID() == BlockID)
-        return &BlockInfoRecords[i];
-    return nullptr;
-  }
-
-  BlockInfo &getOrCreateBlockInfo(unsigned BlockID) {
-    if (const BlockInfo *BI = getBlockInfo(BlockID))
-      return *const_cast<BlockInfo*>(BI);
-
-    // Otherwise, add a new record.
-    BlockInfoRecords.push_back(BlockInfo(BlockID));
-    return BlockInfoRecords.back();
+  BlockInfo &getBlockInfo(unsigned BlockID) {
+    return BlockInfoRecords->getBlockInfo(BlockID);
   }
 };
 
 /// When advancing through a bitstream cursor, each advance can discover a few
 /// different kinds of entries:
 struct NaClBitstreamEntry {
   enum {
     Error,    // Malformed bitcode was found.
     EndBlock, // We've reached the end of the current block, (or the end of the
               // file, which is treated like a series of EndBlock records.
@@ skipping 182 matching lines @@
       : ErrHandler(new ErrorHandler(*this)) { init(&R); }
 
   void init(NaClBitstreamReader *R) {
     freeState();
     BitStream = R;
     NextChar = (BitStream == nullptr) ? 0 : BitStream->getInitialAddress();
     Size = 0;
     BitsInCurWord = 0;
     if (BitStream) {
       BlockScope.push_back(
-          Block(&BitStream->getOrCreateBlockInfo(naclbitc::TOP_LEVEL_BLOCKID)));
+          Block(&BitStream->getBlockInfo(naclbitc::TOP_LEVEL_BLOCKID)));
     }
   }
 
   ~NaClBitstreamCursor() {
     freeState();
   }
 
   void freeState() {
     while (!BlockScope.empty())
       BlockScope.pop_back();
@@ skipping 89 matching lines @@
       NaClBitstreamEntry Entry = advance(Flags, 0);
       if (Entry.Kind != NaClBitstreamEntry::SubBlock)
         return Entry;
 
       // If we found a sub-block, just skip over it and check the next entry.
       if (SkipBlock())
         return NaClBitstreamEntry::getError();
     }
   }
 
+  /// Returns the starting byte of the word containing BitNo.
+  uintptr_t getStartWordByteForBit(uint64_t BitNo) const {
+    return uintptr_t(BitNo/8) & ~(sizeof(word_t)-1);
+  }
+
+  /// Returns the index of BitNo within the word it appears in.
+  unsigned getWordBitNo(uint64_t BitNo) const {
+    return unsigned(BitNo & (sizeof(word_t)*8-1));
+  }
+
+  /// Returns the ending byte of the word containing BitNo.
+  uintptr_t getEndWordByteForBit(uint64_t BitNo) const {
+    return getStartWordByteForBit(BitNo) +
+        (getWordBitNo(BitNo)
+         ? sizeof(word_t)
+         : 0);
+  }
+
+  /// Fills Buffer[Size] using bytes at Address (in the memory object being
+  /// read). Returns number of bytes filled (less than Size if at end of memory
+  /// object).
+  uint64_t fillBuffer(uint8_t *Buffer, size_t Size, size_t Address) const {
+    return BitStream->getBitcodeBytes().readBytes(Buffer, Size, Address);
+  }
+
   /// Reset the stream to the specified bit number.
   void JumpToBit(uint64_t BitNo) {
-    uintptr_t ByteNo = uintptr_t(BitNo/8) & ~(sizeof(word_t)-1);
-    unsigned WordBitNo = unsigned(BitNo & (sizeof(word_t)*8-1));
+    uintptr_t ByteNo = getStartWordByteForBit(BitNo);
+    unsigned WordBitNo = getWordBitNo(BitNo);
     if (!canSkipToPos(ByteNo))
       reportInvalidJumpToBit(BitNo);
 
     // Move the cursor to the right word.
     NextChar = ByteNo;
     BitsInCurWord = 0;
 
     // Skip over any bits that are already consumed.
     if (WordBitNo)
       Read(WordBitNo);
   }
 
   void fillCurWord() {
     assert(Size == 0 || NextChar < (unsigned)Size);
 
     // Read the next word from the stream.
     uint8_t Array[sizeof(word_t)] = {0};
 
-    uint64_t BytesRead =
-        BitStream->getBitcodeBytes().readBytes(Array, sizeof(Array), NextChar);
+    uint64_t BytesRead = fillBuffer(Array, sizeof(Array), NextChar);
 
     // If we run out of data, stop at the end of the stream.
     if (BytesRead == 0) {
       Size = NextChar;
       return;
     }
 
     CurWord =
         support::endian::read<word_t, support::little, support::unaligned>(
             Array);
@@ skipping 225 matching lines @@
   // Skips over an abbreviation record. Duplicates code of ReadAbbrevRecord,
   // except that no abbreviation is built.
   void SkipAbbrevRecord();
   
   bool ReadBlockInfoBlock(NaClAbbrevListener *Listener);
 };
 
 } // End llvm namespace
 
 #endif