Fix abbreviation handling in munged bitcode writer.

lib/Bitcode/NaCl/TestUtils/NaClBitcodeMungeWriter.cpp

 //===- NaClBitcodeMungeWriter.cpp - Write munged bitcode --------*- C++ -*-===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 //
 // Implements method NaClMungedBitcode.write(), which writes out a munged
 // list of bitcode records using a bitstream writer.
 
 #include "llvm/Bitcode/NaCl/NaClBitcodeMungeUtils.h"
 
 #include "llvm/Bitcode/NaCl/NaClBitstreamWriter.h"
 #include "llvm/Bitcode/NaCl/NaClReaderWriter.h"
 #include "llvm/Support/raw_ostream.h"
 
+#include <set>
+
 using namespace llvm;
 
 namespace {
 
 // For debugging. When true, shows trace information of emitting
 // bitcode.
 static bool DebugEmit = false;
 
 // Description of current block scope
 struct BlockScope {
   BlockScope(unsigned CurBlockID, size_t AbbrevIndexLimit)
-      : CurBlockID(CurBlockID), AbbrevIndexLimit(AbbrevIndexLimit) {}
+      : CurBlockID(CurBlockID), AbbrevIndexLimit(AbbrevIndexLimit),
+        OmittedAbbreviations(false) {}
   unsigned CurBlockID;
   // Defines the maximum value for abbreviation indices in block.
   size_t AbbrevIndexLimit;
+  // Defines if an abbreviation definition was omitted (i.e. not
+  // written) from this block. Used to turn off writing further
+  // abbreviation definitions for this block.
+  bool OmittedAbbreviations;
   void print(raw_ostream &Out) const {
     Out << "BlockScope(ID=" << CurBlockID << ", AbbrevIndexLimit="
-        << AbbrevIndexLimit << ")";
+        << AbbrevIndexLimit << "OmittedAbbreviations="
+        << OmittedAbbreviations << ")";
   }
 };
 
 inline raw_ostream &operator<<(raw_ostream &Out, const BlockScope &Scope) {
   Scope.print(Out);
   return Out;
 }
 
 // State of bitcode writing.
 struct WriteState {
   // The block ID associated with records not in any block.
   static const unsigned UnknownWriteBlockID = UINT_MAX;
   // The SetBID for the blockinfo block.
   unsigned SetBID = UnknownWriteBlockID;
   // The stack of scopes the writer is in.
-  SmallVector<BlockScope, 3> ScopeStack;
+  SmallVector<BlockScope, 4> ScopeStack;
   // The set of write flags to use.
   const NaClMungedBitcode::WriteFlags &Flags;
   // The results of the attempted write.
   NaClMungedBitcode::WriteResults Results;
   // The minimum number of bits allowed to be specified in a block.
   const unsigned BlockMinBits;
+  // The set of block IDs for which abbreviation definitions have been
+  // omitted in the blockinfo block.
+  std::set<unsigned> BlocksWithOmittedAbbrevs;
 
   WriteState(const NaClMungedBitcode::WriteFlags &Flags)
       : Flags(Flags),
         BlockMinBits(NaClBitsNeededForValue(naclbitc::DEFAULT_MAX_ABBREV)) {
     BlockScope Scope(UnknownWriteBlockID, naclbitc::DEFAULT_MAX_ABBREV);
     ScopeStack.push_back(Scope);
   }
 
   // Returns stream to print error message to.
   raw_ostream &Error();
@@ skipping 14 matching lines @@
   unsigned getCurWriteBlockID() const {
     assert(!ScopeStack.empty());
     return ScopeStack.back().CurBlockID;
   }
 
   unsigned getCurAbbrevIndexLimit() const {
     assert(!ScopeStack.empty());
     return ScopeStack.back().AbbrevIndexLimit;
   }
 
+  // Returns whether any abbreviation definitions were not written to
+  // the bitcode buffer.
+  bool curBlockHasOmittedAbbreviations() const {
+    assert(!ScopeStack.empty());
+    return ScopeStack.back().OmittedAbbreviations
+        || BlocksWithOmittedAbbrevs.count(getCurWriteBlockID());
+  }
+
+  // Marks that an abbreviation definition is being omitted (i.e. not
+  // written) for the current block.
+  void markCurrentBlockWithOmittedAbbreviations() {
+    assert(!ScopeStack.empty());
+    ScopeStack.back().OmittedAbbreviations = true;
+    if (getCurWriteBlockID() == naclbitc::BLOCKINFO_BLOCK_ID)
+      BlocksWithOmittedAbbrevs.insert(SetBID);
+  }
+
   // Converts the abbreviation record to the corresponding abbreviation.
   // Returns nullptr if unable to build abbreviation.
   NaClBitCodeAbbrev *buildAbbrev(const NaClBitcodeAbbrevRecord &Record);
 
   // Emits the given record to the bitcode file. Returns true if
   // successful.
   bool LLVM_ATTRIBUTE_UNUSED_RESULT
-  emitRecord(NaClBitstreamWriter &Writer, const NaClBitcodeAbbrevRecord &Record);
+  emitRecord(NaClBitstreamWriter &Writer,
+             const NaClBitcodeAbbrevRecord &Record);
 
   // Enter the given block
   bool LLVM_ATTRIBUTE_UNUSED_RESULT
   enterBlock(NaClBitstreamWriter &Writer, uint64_t WriteBlockID,
              uint64_t NumBits, const NaClBitcodeAbbrevRecord &Record);
 
   // Exit current block and return to previous block. Silently recovers
   // if at outermost scope.
   bool LLVM_ATTRIBUTE_UNUSED_RESULT
   exitBlock(NaClBitstreamWriter &Writer,
             const NaClBitcodeAbbrevRecord *Record = nullptr) {
     if (atOutermostScope())
       return false;
     Writer.ExitBlock();
     ScopeStack.pop_back();
     if (DebugEmit)
       printScopeStack(errs());
     return true;
   }
 
+  void WriteRecord(NaClBitstreamWriter &Writer,
+                   const NaClBitcodeAbbrevRecord &Record,
+                   bool UsesDefaultAbbrev) {
+    if (UsesDefaultAbbrev)
+      Writer.EmitRecord(Record.Code, Record.Values);
+    else
+      Writer.EmitRecord(Record.Code, Record.Values, Record.Abbrev);
+  }
+
+  // Returns true if the abbreviation index defines an abbreviation
+  // that can be applied to the record.
+  bool LLVM_ATTRIBUTE_UNUSED_RESULT
+  canApplyAbbreviation(NaClBitstreamWriter &Writer,
+                       const NaClBitcodeAbbrevRecord &Record);
+
   // Completes the write.
   NaClMungedBitcode::WriteResults &finish(NaClBitstreamWriter &Writer,
                                           bool RecoverSilently);
 
   void printScopeStack(raw_ostream &Out) {
     Out << "Scope Stack:\n";
     for (auto &Scope : ScopeStack)
       Out << "  " << Scope << "\n";
   }
 };
@@ skipping 47 matching lines @@
         return false;
     }
     Writer.EnterBlockInfoBlock();
   } else {
     NaClBitcodeSelectorAbbrev CurCodeLen(MaxAbbrev);
     Writer.EnterSubblock(WriteBlockID, CurCodeLen);
   }
   return true;
 }
 
+bool WriteState::canApplyAbbreviation(
+    NaClBitstreamWriter &Writer, const NaClBitcodeAbbrevRecord &Record) {
+  const NaClBitCodeAbbrev *Abbrev = Writer.getAbbreviation(Record.Abbrev);
+  if (Abbrev == nullptr)
+    return false;
+
+  // Merge record code into values and then match abbreviation.
+  NaClBitcodeValues Values(Record);
+  size_t ValueIndex = 0;
+  size_t ValuesSize = Values.size();
+  size_t AbbrevIndex = 0;
+  size_t AbbrevSize = Abbrev->getNumOperandInfos();
+  bool FoundArray = false;
+  while (ValueIndex < ValuesSize && AbbrevIndex < AbbrevSize) {
+    const NaClBitCodeAbbrevOp *Op = &Abbrev->getOperandInfo(AbbrevIndex++);
+    uint64_t Value = Values[ValueIndex++];
+    if (Op->getEncoding() == NaClBitCodeAbbrevOp::Array) {
+      if (AbbrevIndex + 1 != AbbrevSize)
+        return false;
+      Op = &Abbrev->getOperandInfo(AbbrevIndex);
+      --AbbrevIndex; // i.e. don't advance to next abbreviation op.
+      FoundArray = true;
+    }
+    switch (Op->getEncoding()) {
+    case NaClBitCodeAbbrevOp::Literal:
+      if (Value != Op->getValue())
+        return false;
+      continue;
+    case NaClBitCodeAbbrevOp::Fixed:
+      if (Value >= (static_cast<uint64_t>(1)
+                    << NaClBitstreamWriter::MaxEmitNumBits)
+          || NaClBitsNeededForValue(Value) > Op->getValue())
+        return false;
+      continue;
+    case NaClBitCodeAbbrevOp::VBR:
+      if (Op->getValue() < 2)
+        return false;
+      continue;
+    case NaClBitCodeAbbrevOp::Array:
+      llvm_unreachable("Array(Array) abbreviation is not legal!");
+    case NaClBitCodeAbbrevOp::Char6:
+      if (!Op->isChar6(Value))
+        return false;
+      continue;
+    }
+  }
+  return ValueIndex == ValuesSize && (FoundArray || AbbrevIndex == AbbrevSize);
+}
+
 NaClMungedBitcode::WriteResults &WriteState::finish(
     NaClBitstreamWriter &Writer, bool RecoverSilently) {
   // Be sure blocks are balanced.
   while (!atOutermostScope()) {
     if (!RecoverSilently)
       RecoverableError() << "Missing close block.\n";
     if (!exitBlock(Writer)) {
       Error() << "Failed to add missing close block at end of file.\n";
       break;
     }
@@ skipping 66 matching lines @@
         return false;
     }
     if (!exitBlock(Writer)) {
       Error() << "Failed to write exit block, can't continue: "
               << Record << "\n";
       return false;
     }
     break;
   }
   case naclbitc::BLK_CODE_DEFINE_ABBREV: {
+    if (curBlockHasOmittedAbbreviations()) {
+      // If reached, a previous abbreviation for the block was omitted. Can't
+      // generate more abbreviations without having to fix abbreviation indices.
+      RecoverableError() << "Ignoring abbreviation: " << Record << "\n";
+      return Flags.getTryToRecover();
+    }
     if (Record.Abbrev != naclbitc::DEFINE_ABBREV) {
       RecoverableError()
           << "Uses illegal abbreviation index in define abbreviation record: "
           << Record << "\n";
       if (!Flags.getTryToRecover())
         return false;
     }
-    if (getCurWriteBlockID() != naclbitc::BLOCKINFO_BLOCK_ID
-        && Writer.getMaxCurAbbrevIndex() >= getCurAbbrevIndexLimit()) {
-      RecoverableError() << "Exceeds abbreviation index limit of "
-                         << getCurAbbrevIndexLimit() << ": " << Record << "\n";
-      // Recover by not writing.
+    NaClBitCodeAbbrev *Abbrev = buildAbbrev(Record);
+    if (Abbrev == nullptr) {
+      markCurrentBlockWithOmittedAbbreviations();
       return Flags.getTryToRecover();
     }
-    NaClBitCodeAbbrev *Abbrev = buildAbbrev(Record);
-    if (Abbrev == NULL)
-      return Flags.getTryToRecover();
     if (getCurWriteBlockID() == naclbitc::BLOCKINFO_BLOCK_ID) {
       Writer.EmitBlockInfoAbbrev(SetBID, Abbrev);
     } else {
       Writer.EmitAbbrev(Abbrev);
     }
     break;
   }
   case naclbitc::BLK_CODE_HEADER:
     // Note: There is no abbreviation index here. Ignore.
     for (uint64_t Value : Record.Values)
       Writer.Emit(Value, 8);
     break;
   default:
     bool UsesDefaultAbbrev = Record.Abbrev == naclbitc::UNABBREV_RECORD;
     if (atOutermostScope()) {
       RecoverableError() << "Record outside block: " << Record << "\n";
       if (!Flags.getTryToRecover())
         return false;
       // Create a dummy block to hold record.
       if (!enterBlock(Writer, UnknownWriteBlockID,
                       naclbitc::DEFAULT_MAX_ABBREV, Record)) {
         Error() << "Failed to recover from record outside block\n";
         return false;
       }
       UsesDefaultAbbrev = true;
     }
-    if (!UsesDefaultAbbrev
-        && !Writer.isUserRecordAbbreviation(Record.Abbrev)) {
-      // Illegal abbreviation index found.
+    if (!UsesDefaultAbbrev && !canApplyAbbreviation(Writer, Record)) {
+      if (Writer.getAbbreviation(Record.Abbrev) != nullptr) {
+        RecoverableError() << "Abbreviation doesn't apply to record: "
+                           << Record << "\n";
+        UsesDefaultAbbrev = true;
+        if (!Flags.getTryToRecover())
+          return false;
+        WriteRecord(Writer, Record, UsesDefaultAbbrev);
+        return true;
+      }
       if (Flags.getWriteBadAbbrevIndex()) {
+        // The abbreviation is not understood by the bitcode writer,
+        // and the flag value implies that we should still write it
+        // out so that unit tests for this error condition can be
+        // applied.
         Error() << "Uses illegal abbreviation index: " << Record << "\n";
         // Generate bad abbreviation index so that the bitcode reader
         // can be tested, and then quit.
         Results.WroteBadAbbrevIndex = true;
         Writer.EmitCode(Record.Abbrev);
         bool RecoverSilently = true;
         finish(Writer, RecoverSilently);
         return false;
       }
       RecoverableError() << "Uses illegal abbreviation index: "
                          << Record << "\n";
+      UsesDefaultAbbrev = true;
       if (!Flags.getTryToRecover())
         return false;
-      UsesDefaultAbbrev = true;
+      WriteRecord(Writer, Record, UsesDefaultAbbrev);
+      return true;
     }
     if (getCurWriteBlockID() == 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.
       // Rather just set SetBID and let call to Writer->EmitBlockInfoAbbrev
       // generate the SetBID record.
       if (NumValues != 1) {
         Error() << "SetBID record expects 1 value but found "
                 << NumValues << ": " << Record << "\n";
         return false;
       }
       SetBID = Record.Values[0];
       return true;
     }
-    if (UsesDefaultAbbrev)
-      Writer.EmitRecord(Record.Code, Record.Values);
-    else
-      Writer.EmitRecord(Record.Code, Record.Values, Record.Abbrev);
+    WriteRecord(Writer, Record, UsesDefaultAbbrev);
+    break;
   }
   return true;
 }
 
 static NaClBitCodeAbbrev *deleteAbbrev(NaClBitCodeAbbrev *Abbrev) {
   Abbrev->dropRef();
   return nullptr;
 }
 
 NaClBitCodeAbbrev *WriteState::buildAbbrev(
@@ skipping 10 matching lines @@
   size_t NumAbbreviations = Record.Values[Index++];
   if (NumAbbreviations == 0) {
     RecoverableError() << "Abbreviation must contain at least one operator: "
                        << Record << "\n";
     return deleteAbbrev(Abbrev);
   }
   for (size_t Count = 0; Count < NumAbbreviations; ++Count) {
     if (Index >= NumValues) {
       RecoverableError()
           << "Malformed abbreviation found. Expects " << NumAbbreviations
-          << " operands but ound " << Count << ": " << Record << "\n";
+          << " operands but found " << Count << ": " << Record << "\n";
       return deleteAbbrev(Abbrev);
     }
     switch (Record.Values[Index++]) {
     case 1:
       if (Index >= NumValues) {
         RecoverableError() << "Malformed literal abbreviation: "
                            << Record << "\n";
         return deleteAbbrev(Abbrev);
       }
       Abbrev->Add(NaClBitCodeAbbrevOp(Record.Values[Index++]));
@@ skipping 36 matching lines @@
         Abbrev->Add(NaClBitCodeAbbrevOp(NaClBitCodeAbbrevOp::Char6));
         break;
       default:
         RecoverableError() << "Unknown abbreviation kind " << Kind
                            << ": " << Record << "\n";
         return deleteAbbrev(Abbrev);
       }
       break;
     }
     default:
-      RecoverableError() << "Error: Bad literal flag " << Record.Values[Index]
-                         << ": " << Record << "\n";
+      RecoverableError() << "Error: Bad abbreviation operand encoding "
+                         << Record.Values[Index-1] << ": " << Record << "\n";
       return deleteAbbrev(Abbrev);
     }
   }
   return Abbrev;
 }
 
 } // end of anonymous namespace.
 
 NaClMungedBitcode::WriteResults NaClMungedBitcode::writeMaybeRepair(
     SmallVectorImpl<char> &Buffer, bool AddHeader,
     const WriteFlags &Flags) const {
   NaClBitstreamWriter Writer(Buffer);
   WriteState State(Flags);
   if (AddHeader) {
     NaClWriteHeader(Writer, true);
   }
   for (const NaClBitcodeAbbrevRecord &Record : *this) {
     if (!State.emitRecord(Writer, Record))
       break;
   }
   bool RecoverSilently =
       State.Results.NumErrors > 0 && !Flags.getTryToRecover();
   return State.finish(Writer, RecoverSilently);
 }