Fix the block stack used by the bitstream cursor.

include/llvm/Bitcode/NaCl/NaClBitstreamReader.h

Index: include/llvm/Bitcode/NaCl/NaClBitstreamReader.h
diff --git a/include/llvm/Bitcode/NaCl/NaClBitstreamReader.h b/include/llvm/Bitcode/NaCl/NaClBitstreamReader.h
index 05cd91eb2aa92af2eb788d8c7289919604754b2e..29b476e4c1722a84594e4ac7d735c36c4d45e2da 100644
--- a/include/llvm/Bitcode/NaCl/NaClBitstreamReader.h
+++ b/include/llvm/Bitcode/NaCl/NaClBitstreamReader.h
@@ -27,6 +27,7 @@
 namespace llvm {
 
 class Deserializer;
+class NaClBitstreamCursor;
 
 namespace naclbitc {
 
@@ -55,12 +56,93 @@ raw_ostream &ErrorAt(raw_ostream &Out, ErrorLevel Level,
 /// 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;
 
@@ -68,6 +150,9 @@ private:
 
   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;
 
@@ -106,17 +191,7 @@ public:
   // 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 {
@@ -127,21 +202,17 @@ public:
   // 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;
   }
@@ -151,8 +222,7 @@ public:
       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();
   }
 };
@@ -205,7 +275,7 @@ public:
   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
@@ -276,22 +346,67 @@ private:
   /// 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;
@@ -311,13 +426,20 @@ public:
     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
@@ -342,7 +464,9 @@ public:
   }
 
   /// 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 {
@@ -559,9 +683,11 @@ private:
 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:
@@ -602,24 +728,12 @@ public:
     //    [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
   //===--------------------------------------------------------------------===//
@@ -657,9 +771,16 @@ 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.