Fix the block stack used by the bitstream cursor.

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 <climits>
 #include <vector>
 
 namespace llvm {
 
 class Deserializer;
+class NaClBitstreamCursor;
 
 namespace naclbitc {
 
 /// Returns the Bit as a Byte:BitInByte string.
 std::string getBitAddress(uint64_t Bit);
 
 /// Severity levels for reporting errors.
 enum ErrorLevel {
   Warning,
   Error,
   Fatal
 };
 
 // Basic printing routine to generate the beginning of an error
 // message. BitPosition is the bit position the error was found.
 // Level is the severity of the error.
 raw_ostream &ErrorAt(raw_ostream &Out, ErrorLevel Level,
                      uint64_t BitPosition);
 
 } // End namespace naclbitc.
 
 /// This class is used to read from a NaCl bitcode wire format stream,
 /// maintaining information that is global to decoding the entire file.
 /// While a file is being read, multiple cursors can be independently
 /// advanced or skipped around within the file.  These are represented by
 /// the NaClBitstreamCursor class.
 class NaClBitstreamReader {
 public:
-  /// This contains information emitted to BLOCKINFO_BLOCK blocks. These
-  /// describe abbreviations that all blocks of the specified ID inherit.
-  struct BlockInfo {
+  // Models a raw list of abbreviations.
+  static const size_t DefaultAbbrevListSize = 12;
+  using AbbrevListVector = SmallVector<NaClBitCodeAbbrev *,
+                                       DefaultAbbrevListSize>;
+
+  // Models and maintains a list of abbreviations. In particular, it maintains
+  // updating reference counts of abbreviation operators within the abbreviation
+  // list.
+  class AbbrevList {
+  public:
+    AbbrevList() = default;
+    explicit AbbrevList(const AbbrevList &NewAbbrevs) {
+      appendList(NewAbbrevs);
+    }
+    AbbrevList &operator=(const AbbrevList &Rhs) {
+      clear();
+      appendList(Rhs);
+      return *this;
+    }
+    // appends Abbrev assuming that Abbv is not shared (i.e. already has the
+    // refernce count incremented). Takes ownership of Abbv.
+    void appendUnique(NaClBitCodeAbbrev *Abbv) {

[Derek Schuff, 2016/03/17 16:20:17]

Does appendUnique mean that nothing else will ever get a reference to it? If so
then this naming is fine. If this is just the first reference (and we just have
to avoid incrementing the refcount because it starts at 1) then maybe we could
call it appendFirst, or something?

[Karl, 2016/03/17 18:37:08]

Simplified the class by adding a "appendCreate" method that creates a new
abbreviation and appends it to the list. Then, we only need append() and
appendList(), and the user is no longer exposed to reference counts.

+      Abbrevs.push_back(Abbv);
+    }
+    // appends Abbrev assuming that that Abbv is shared elsewhere (i.e. has not
+    // yet had the reference count incremented for sharing).
+    void appendShare(NaClBitCodeAbbrev *Abbrv) {
+      Abbrv->addRef();
+      Abbrevs.push_back(Abbrv);
+    }
+    // Returns last abbreviation on list, and returns an unique reference (i.e.
+    // the reference count has been incremented before returning).
+    NaClBitCodeAbbrev *lastUnique() {
+      NaClBitCodeAbbrev *Abbrv = Abbrevs.back();
+      Abbrv->addRef();
+      return Abbrv;
+    }
+    // Removes the last element of the list.
+    void popLast() {
+      Abbrevs.back()->dropRef();
+      Abbrevs.pop_back();
+    }
+    // Copies contents of NewAbbrevs and appends to this.
+    void appendList(const AbbrevList &NewAbbrevs) {
+      for (NaClBitCodeAbbrev *Abbrv : NewAbbrevs.Abbrevs)
+        appendShare(Abbrv);
+    }
+    // Empties abbreviation list.
+    void clear() {
+      while(!Abbrevs.empty())
+        popLast();
+    }
+    // Allow read access to vector defining list.
+    const AbbrevListVector &getVector() const { return Abbrevs; }
+    ~AbbrevList() { clear(); }
+  private:
+    AbbrevListVector Abbrevs;
+  };
+
+  /// This contains information about abbreviations in blocks defined in the
+  /// BLOCKINFO_BLOCK block.  These describe global abbreviations that apply to
+  /// all succeeding blocks of the specified ID.
+  class BlockInfo {
+    BlockInfo &operator=(const BlockInfo&) = delete;
+  public:
+    BlockInfo() = default;
+    explicit BlockInfo(unsigned BlockID)
+        : BlockID(BlockID), Abbrevs() {}
+    BlockInfo(const BlockInfo&) = default;
+    unsigned getBlockID() const { return BlockID; }
+    void setBlockID(unsigned ID) { BlockID = ID; }
+    AbbrevList &getAbbrevs() { return Abbrevs; }
+    // Allow read access to vector defining abbreviation list.
+    const AbbrevListVector &getVector() const {
+      return Abbrevs.getVector();
+    }
+    // appends Abbrev assuming that Abbv is not shared (i.e. already has the
+    // refernce count incremented). Takes ownership of Abbv.
+    void appendUnique(NaClBitCodeAbbrev *Abbv) {
+      Abbrevs.appendUnique(Abbv);
+    }
+    ~BlockInfo() {}
+  private:
     unsigned BlockID;
-    std::vector<NaClBitCodeAbbrev*> Abbrevs;
+    AbbrevList Abbrevs;
   };
+
 private:
   friend class NaClBitstreamCursor;
 
   std::unique_ptr<MemoryObject> BitcodeBytes;
 
   std::vector<BlockInfo> BlockInfoRecords;
 
+  // True if the BlockInfo block has been read.
+  bool HasReadBlockInfoBlock = false;
+
   /// \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) {
@@ skipping 18 matching lines @@
 
   /// 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) {
   }
 
   // Returns the memory object that is being read.
   MemoryObject &getBitcodeBytes() { return *BitcodeBytes; }
 
-  ~NaClBitstreamReader() {
-    // Free the BlockInfoRecords.
-    while (!BlockInfoRecords.empty()) {
-      BlockInfo &Info = BlockInfoRecords.back();
-      // Free blockinfo abbrev info.
-      for (unsigned i = 0, e = static_cast<unsigned>(Info.Abbrevs.size());
-           i != e; ++i)
-        Info.Abbrevs[i]->dropRef();
-      BlockInfoRecords.pop_back();
-    }
-  }
+  ~NaClBitstreamReader() {}
 
   /// \brief Returns the initial address (after the header) of the input stream.
   size_t getInitialAddress() const {
     return InitialAddress;
   }
 
   //===--------------------------------------------------------------------===//
   // Block Manipulation
   //===--------------------------------------------------------------------===//
 
-  /// Return true if we've already read and processed the block info block for
-  /// this Bitstream. We only process it for the first cursor that walks over
-  /// it.
-  bool hasBlockInfoRecords() const { return !BlockInfoRecords.empty(); }
-
   /// 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().BlockID == 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].BlockID == BlockID)
+      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());
-    BlockInfoRecords.back().BlockID = BlockID;
+    BlockInfoRecords.push_back(BlockInfo(BlockID));
     return BlockInfoRecords.back();
   }
 };
 
 /// 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
@@ skipping 32 matching lines @@
 /// to allow the readers to update internals of the bit stream reader
 /// appropriately.
 class NaClAbbrevListener {
   NaClAbbrevListener(const NaClAbbrevListener&) = delete;
   void operator=(const NaClAbbrevListener&) = delete;
 public:
   NaClAbbrevListener() {}
   virtual ~NaClAbbrevListener() {}
 
   /// Called to process the read abbreviation.
-  virtual void ProcessAbbreviation(NaClBitCodeAbbrev *Abbrev,
+  virtual void ProcessAbbreviation(NaClBitCodeAbbrev *Abbrv,
                                    bool IsLocal) = 0;
 
   /// Called after entering block. NumWords is the number of words
   /// in the block.
   virtual void BeginBlockInfoBlock(unsigned NumWords) = 0;
 
   /// Called if a naclbitc::BLOCKINFO_CODE_SETBID record is found in
   /// NaClBitstreamCursor::ReadBlockInfoBlock.
   virtual void SetBID() = 0;
 
@@ skipping 50 matching lines @@
   /// follow the word size of the host machine for efficiency. We use word_t in
   /// places that are aware of this to make it perfectly explicit what is going
   /// on.
   typedef size_t word_t;
   word_t CurWord;
 
   /// This is the number of bits in CurWord that are valid. This
   /// is always from [0...bits_of(word_t)-1] inclusive.
   unsigned BitsInCurWord;
 
-  /// This is the declared size of code values used for the current
-  /// block, in bits.
-  NaClBitcodeSelectorAbbrev CurCodeSize;
-
-  /// Abbrevs installed in this block.
-  std::vector<NaClBitCodeAbbrev*> CurAbbrevs;
-
-  struct Block {
-    NaClBitcodeSelectorAbbrev PrevCodeSize;
-    std::vector<NaClBitCodeAbbrev*> PrevAbbrevs;
-    Block() : PrevCodeSize() {}
-    explicit Block(const NaClBitcodeSelectorAbbrev& PCS)
-        : PrevCodeSize(PCS) {}
+  // Data specific to a block being scanned.
+  class Block {
+  public:
+    Block() = delete;
+    Block &operator=(const Block &Rhs) {
+      GlobalAbbrevs = Rhs.GlobalAbbrevs;
+      NumGlobalAbbrevs = Rhs.NumGlobalAbbrevs;
+      LocalAbbrevs = Rhs.LocalAbbrevs;
+      CodeAbbrev = Rhs.CodeAbbrev;
+      return *this;
+    }
+    Block(NaClBitstreamReader::BlockInfo *GlobalAbbrevs,
+          NaClBitcodeSelectorAbbrev& CodeAbbrev)
+        : GlobalAbbrevs(GlobalAbbrevs),
+          NumGlobalAbbrevs(GlobalAbbrevs->getVector().size()),
+          LocalAbbrevs(), CodeAbbrev(CodeAbbrev) {}
+    Block(NaClBitstreamReader::BlockInfo *GlobalAbbrevs)
+        : GlobalAbbrevs(GlobalAbbrevs),
+          NumGlobalAbbrevs(GlobalAbbrevs->getVector().size()),
+          LocalAbbrevs(), CodeAbbrev() {}
+    ~Block() = default;
+    const NaClBitstreamReader::AbbrevList &getGlobalAbbrevs() const {
+      return GlobalAbbrevs->getAbbrevs();
+    }
+    unsigned getNumGlobalAbbrevs() const { return NumGlobalAbbrevs; }
+    const NaClBitstreamReader::AbbrevList &getLocalAbbrevs() const {
+      return LocalAbbrevs;
+    }
+    const NaClBitcodeSelectorAbbrev &getCodeAbbrev() const {
+      return CodeAbbrev;
+    }
+    void setCodeAbbrev(NaClBitcodeSelectorAbbrev &Abbrev) {
+      CodeAbbrev = Abbrev;
+    }
+    // Appends Abbrev to local abbreviations assuming that Abbv is not shared
+    // (i.e. already has the refernce count incremented). Takes ownship of Abbv.
+    void appendUniqueLocalAbbrev(NaClBitCodeAbbrev *Abbv) {
+      LocalAbbrevs.appendUnique(Abbv);
+    }
+    /// removes the last added local abbreviation and moves it to the given
+    /// abbreviation list.
+    void moveLocalAbbrevToAbbrevList(NaClBitstreamReader::AbbrevList *List) {
+      NaClBitCodeAbbrev *Abbv = LocalAbbrevs.lastUnique();
+      LocalAbbrevs.popLast();
+      List->appendUnique(Abbv);
+    }
+  private:
+    friend class NaClBitstreamCursor;
+    // The global abbreviations associated with this scope.
+    NaClBitstreamReader::BlockInfo *GlobalAbbrevs;
+    // Number of abbreviations when block was entered. Used to limit scope of
+    // CurBlockInfo, since any abbreviation added inside a BlockInfo block
+    // (within this block) must not effect global abbreviations.
+    unsigned NumGlobalAbbrevs;
+    NaClBitstreamReader::AbbrevList LocalAbbrevs;
+    // This is the declared size of code values used for the current block, in
+    // bits.
+    NaClBitcodeSelectorAbbrev CodeAbbrev;
   };
 
-  /// This tracks the codesize of parent blocks.
+  /// This tracks the Block-specific information for each nested block.
   SmallVector<Block, 8> BlockScope;
 
   NaClBitstreamCursor(const NaClBitstreamCursor &) = delete;
   NaClBitstreamCursor &operator=(const NaClBitstreamCursor &) = delete;
 
 public:
   NaClBitstreamCursor() : ErrHandler(new ErrorHandler(*this)) {
     init(nullptr);
   }
 
   explicit NaClBitstreamCursor(NaClBitstreamReader &R)
       : 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)));
+    }
   }
 
   ~NaClBitstreamCursor() {
     freeState();
   }
 
-  void freeState();
+  void freeState() {
+    while (!BlockScope.empty())
+      BlockScope.pop_back();
+  }
 
   // Replaces the current bitstream error handler with the new
   // handler. Takes ownership of the new handler and deletes it when
   // it is no longer needed.
   void setErrorHandler(std::unique_ptr<ErrorHandler> &NewHandler) {
     ErrHandler = std::move(NewHandler);
   }
 
   bool canSkipToPos(size_t pos) const {
     // pos can be skipped to if it is a valid address or one byte past the end.
     return pos == 0 || BitStream->getBitcodeBytes().isValidAddress(
         static_cast<uint64_t>(pos - 1));
   }
 
   bool AtEndOfStream() {
     if (BitsInCurWord != 0)
       return false;
     if (Size != 0)
       return Size == NextChar;
     fillCurWord();
     return BitsInCurWord == 0;
   }
 
   /// Return the number of bits used to encode an abbrev #.
-  unsigned getAbbrevIDWidth() const { return CurCodeSize.NumBits; }
+  unsigned getAbbrevIDWidth() const {
+    return BlockScope.back().getCodeAbbrev().NumBits;
+  }
 
   /// Return the bit # of the bit we are reading.
   uint64_t GetCurrentBitNo() const {
     return NextChar*CHAR_BIT - BitsInCurWord;
   }
 
   NaClBitstreamReader *getBitStreamReader() {
     return BitStream;
   }
   const NaClBitstreamReader *getBitStreamReader() const {
@@ skipping 196 matching lines @@
       CurWord >>= BitsInCurWord-32;
       BitsInCurWord = 32;
       return;
     }
 
     BitsInCurWord = 0;
   }
 public:
 
   unsigned ReadCode() {
-    return CurCodeSize.IsFixed
-        ? Read(CurCodeSize.NumBits)
-        : ReadVBR(CurCodeSize.NumBits);
+    const NaClBitcodeSelectorAbbrev &CodeAbbrev =
+        BlockScope.back().getCodeAbbrev();
+    return CodeAbbrev.IsFixed
+        ? Read(CodeAbbrev.NumBits)
+        : ReadVBR(CodeAbbrev.NumBits);
   }
 
   // Block header:
   //    [ENTER_SUBBLOCK, blockid, newcodelen, <align4bytes>, blocklen]
 
   /// Having read the ENTER_SUBBLOCK code, read the BlockID for the block.
   unsigned ReadSubBlockID() {
     return ReadVBR(naclbitc::BlockIDWidth);
   }
 
@@ skipping 20 matching lines @@
   /// if the block has an error.
   bool EnterSubBlock(unsigned BlockID, unsigned *NumWordsP = nullptr);
 
   bool ReadBlockEnd() {
     if (BlockScope.empty()) return true;
 
     // Block tail:
     //    [END_BLOCK, <align4bytes>]
     SkipToFourByteBoundary();
 
-    popBlockScope();
+    BlockScope.pop_back();
     return false;
   }
 
 private:
 
-  void popBlockScope() {
-    CurCodeSize = BlockScope.back().PrevCodeSize;
-
-    // Delete abbrevs from popped scope.
-    for (unsigned i = 0, e = static_cast<unsigned>(CurAbbrevs.size());
-         i != e; ++i)
-      CurAbbrevs[i]->dropRef();
-
-    BlockScope.back().PrevAbbrevs.swap(CurAbbrevs);
-    BlockScope.pop_back();
-  }
-
   //===--------------------------------------------------------------------===//
   // Record Processing
   //===--------------------------------------------------------------------===//
 
 private:
   // Returns abbreviation encoding associated with Value.
   NaClBitCodeAbbrevOp::Encoding getEncoding(uint64_t Value);
 
   void skipAbbreviatedField(const NaClBitCodeAbbrevOp &Op);
 
@@ skipping 17 matching lines @@
   void reportInvalidAbbrevNumber(unsigned Index) const;
 
   // Reports that jumping to Bit is not valid.
   void reportInvalidJumpToBit(uint64_t Bit) const;
 
 public:
 
   /// Return the abbreviation for the specified AbbrevId.
   const NaClBitCodeAbbrev *getAbbrev(unsigned AbbrevID) const {
     unsigned AbbrevNo = AbbrevID-naclbitc::FIRST_APPLICATION_ABBREV;
-    if (AbbrevNo >= CurAbbrevs.size())
+    const Block &CurBlock = BlockScope.back();
+    const unsigned NumGlobalAbbrevs = CurBlock.getNumGlobalAbbrevs();
+    if (AbbrevNo < NumGlobalAbbrevs)
+      return CurBlock.getGlobalAbbrevs().getVector()[AbbrevNo];
+    unsigned LocalAbbrevNo = AbbrevNo - NumGlobalAbbrevs;
+    NaClBitstreamReader::AbbrevListVector
+        LocalAbbrevs = CurBlock.getLocalAbbrevs().getVector();
+    if (LocalAbbrevNo >= LocalAbbrevs.size())
       reportInvalidAbbrevNumber(AbbrevID);
-    return CurAbbrevs[AbbrevNo];
+    return LocalAbbrevs[LocalAbbrevNo];
   }
 
   /// Read the current record and discard it.
   void skipRecord(unsigned AbbrevID);
   
   unsigned readRecord(unsigned AbbrevID, SmallVectorImpl<uint64_t> &Vals);
 
   //===--------------------------------------------------------------------===//
   // Abbrev Processing
   //===--------------------------------------------------------------------===//
   // IsLocal indicates where the abbreviation occurs. If it is in the
   // BlockInfo block, IsLocal is false. In all other cases, IsLocal is
   // true.
   void ReadAbbrevRecord(bool IsLocal,
                         NaClAbbrevListener *Listener);
 
   // 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