Bitcode: Introduce initial multi-module reader API.

[pcc]

Summary

Implement getLazyBitcodeModule() and parseBitcodeFile() in terms of it.

[mehdi_amini]

llvm/include/llvm/Bitcode/BitcodeReader.h

48 ↗

(On Diff #78106)

Can you document these two fields?

llvm/lib/Bitcode/Reader/BitcodeReader.cpp

850 ↗

(On Diff #78106)

Why not in the initializer list?

6607 ↗

(On Diff #78106)

I think it is used only on error, isn't it?

We may be able to pass the IdentificationBit to the BitcodeReader and have it jump there to get the ProducerIdentification only if needed.

Might not worth the effort though.

llvm/test/Bitcode/invalid.test

38 ↗

(On Diff #78106)

I'm surprised by this change, is it really intended? I suspect that some of your change is losing coverage here.

(Same on most of the changes below)

[pcc]

llvm/include/llvm/Bitcode/BitcodeReader.h

48 ↗

(On Diff #78106)

Will do

llvm/lib/Bitcode/Reader/BitcodeReader.cpp

850 ↗

(On Diff #78106)

This field belongs to BitcodeReaderBase.

6607 ↗

(On Diff #78106)

We also need to check the epoch in the identification block.

I suppose we could have two functions, one that checks the epoch and the other that pulls the producer string, but again it doesn't seem worth the effort.

llvm/test/Bitcode/invalid.test

38 ↗

(On Diff #78106)

Unfortunately some of our invalid bitcode files also have an invalid block size in addition to the problem under test, so we get an error trying to skip past the module in getBitcodeModuleList().

I'll see if I can restore coverage for these, but it may be some work as it will require binary patching of our test files.

[mehdi_amini]

llvm/test/Bitcode/invalid.test

38 ↗

(On Diff #78106)

Yeah that's really annoying.
Maintaining invalid bitcode file across format change is not a great situation. I'd put this in the category of "unmaintainable test", especially in the absence of a code to easily generate them.

[mehdi_amini]

LGTM.

[pcc]

llvm/test/Bitcode/invalid.test

38 ↗

(On Diff #78106)

Sounds reasonable to me.

llvm/trunk/include/llvm/Bitcode/BitcodeReader.h

namespace llvm {
 
   template <typename T>
   ErrorOr<T> expectedToErrorOrAndEmitErrors(LLVMContext &Ctx, Expected<T> Val) {
     if (!Val)
       return errorToErrorCodeAndEmitErrors(Ctx, Val.takeError());
     return std::move(*Val);
   }
 
+  /// Represents a module in a bitcode file.
+  class BitcodeModule {
+    ArrayRef<uint8_t> Buffer;
+    StringRef ModuleIdentifier;
+
+    // The bitstream location of the IDENTIFICATION_BLOCK.
+    uint64_t IdentificationBit;
+
+    // The bitstream location of this module's MODULE_BLOCK.
+    uint64_t ModuleBit;
+
+    BitcodeModule(ArrayRef<uint8_t> Buffer, StringRef ModuleIdentifier,
+                  uint64_t IdentificationBit, uint64_t ModuleBit)
+        : Buffer(Buffer), ModuleIdentifier(ModuleIdentifier),
+          IdentificationBit(IdentificationBit), ModuleBit(ModuleBit) {}
+
+    // Calls the ctor.
+    friend Expected<std::vector<BitcodeModule>>
+    getBitcodeModuleList(MemoryBufferRef Buffer);
+
+    Expected<std::unique_ptr<Module>>
+    getModuleImpl(LLVMContext &Context, bool MaterializeAll,
+                  bool ShouldLazyLoadMetadata);
+
+  public:
+    /// Read the bitcode module and prepare for lazy deserialization of function
+    /// bodies. If ShouldLazyLoadMetadata is true, lazily load metadata as well.
+    Expected<std::unique_ptr<Module>>
+    getLazyModule(LLVMContext &Context, bool ShouldLazyLoadMetadata);
+
+    /// Read the entire bitcode module and return it.
+    Expected<std::unique_ptr<Module>> parseModule(LLVMContext &Context);
+  };
+
+  /// Returns a list of modules in the specified bitcode buffer.
+  Expected<std::vector<BitcodeModule>>
+  getBitcodeModuleList(MemoryBufferRef Buffer);
+
   /// Read the header of the specified bitcode buffer and prepare for lazy
   /// deserialization of function bodies. If ShouldLazyLoadMetadata is true,
   /// lazily load metadata as well.
   Expected<std::unique_ptr<Module>>
   getLazyBitcodeModule(MemoryBufferRef Buffer, LLVMContext &Context,
                        bool ShouldLazyLoadMetadata = false);
 
   /// Like getLazyBitcodeModule, except that the module takes ownership of

llvm/trunk/lib/Bitcode/Reader/BitcodeReader.cpp

class BitcodeReader : public BitcodeReaderBase, public GVMaterializer {
 
   /// Functions that need to be matched with subprograms when upgrading old
   /// metadata.
   SmallDenseMap<Function *, DISubprogram *, 16> FunctionsWithSPs;
 
   std::vector<std::string> BundleTags;
 
 public:
-  BitcodeReader(BitstreamCursor Stream, LLVMContext &Context);
+  BitcodeReader(BitstreamCursor Stream, StringRef ProducerIdentification,
+                LLVMContext &Context);
 
   Error materializeForwardReferencedFunctions();
 
   Error materialize(GlobalValue *GV) override;
   Error materializeModule() override;
   std::vector<StructType *> getIdentifiedStructTypes() const override;
 
   /// \brief Main interface to parsing a bitcode buffer.
handleAllErrors(std::move(Err), [&](ErrorInfoBase &EIB) {
       EC = EIB.convertToErrorCode();
       Ctx.emitError(EIB.message());
     });
     return EC;
   }
   return std::error_code();
 }
 
-BitcodeReader::BitcodeReader(BitstreamCursor Stream, LLVMContext &Context)
-    : BitcodeReaderBase(std::move(Stream)), Context(Context), ValueList(Context),
-      MetadataList(Context) {}
+BitcodeReader::BitcodeReader(BitstreamCursor Stream,
+                             StringRef ProducerIdentification,
+                             LLVMContext &Context)
+    : BitcodeReaderBase(std::move(Stream)), Context(Context),
+      ValueList(Context), MetadataList(Context) {
+  this->ProducerIdentification = ProducerIdentification;
+}
 
 Error BitcodeReader::materializeForwardReferencedFunctions() {
   if (WillMaterializeAllForwardRefs)
     return Error::success();
 
   // Prevent recursion.
   WillMaterializeAllForwardRefs = true;
 
case bitc::MODULE_CODE_SOURCE_FILENAME:
       break;
     }
     Record.clear();
   }
 }
 
 Error BitcodeReader::parseBitcodeInto(Module *M, bool ShouldLazyLoadMetadata) {
   TheModule = M;
-
-  // We expect a number of well-defined blocks, though we don't necessarily
-  // need to understand them all.
-  while (true) {
-    if (Stream.AtEndOfStream()) {
-      // We didn't really read a proper Module.
-      return error("Malformed IR file");
-    }
-
-    BitstreamEntry Entry =
-      Stream.advance(BitstreamCursor::AF_DontAutoprocessAbbrevs);
-
-    if (Entry.Kind != BitstreamEntry::SubBlock)
-      return error("Malformed block");
-
-    if (Entry.ID == bitc::IDENTIFICATION_BLOCK_ID) {
-      Expected<std::string> ProducerIdentificationOrErr =
-          readIdentificationBlock(Stream);
-      if (!ProducerIdentificationOrErr)
-        return ProducerIdentificationOrErr.takeError();
-      ProducerIdentification = *ProducerIdentificationOrErr;
-      continue;
-    }
-
-    if (Entry.ID == bitc::MODULE_BLOCK_ID)
-      return parseModule(0, ShouldLazyLoadMetadata);
-
-    if (Stream.SkipBlock())
-      return error("Invalid record");
-  }
 }
 
 Error BitcodeReader::parseGlobalObjectAttachment(GlobalObject &GO,
                                                  ArrayRef<uint64_t> Record) {
   assert(Record.size() % 2 == 0);
   for (unsigned I = 0, E = Record.size(); I != E; I += 2) {
     auto K = MDKindMap.find(Record[I]);
     if (K == MDKindMap.end())

 const std::error_category &llvm::BitcodeErrorCategory() {
   return *ErrorCategory;
 }
 
 //===----------------------------------------------------------------------===//
 // External interface
 //===----------------------------------------------------------------------===//
 
+Expected<std::vector<BitcodeModule>>
+llvm::getBitcodeModuleList(MemoryBufferRef Buffer) {
+  Expected<BitstreamCursor> StreamOrErr = initStream(Buffer);
+  if (!StreamOrErr)
+    return StreamOrErr.takeError();
+  BitstreamCursor &Stream = *StreamOrErr;
+
+  uint64_t IdentificationBit = -1ull;
+  std::vector<BitcodeModule> Modules;
+  while (true) {
+    // We may be consuming bitcode from a client that leaves garbage at the end
+    // of the bitcode stream (e.g. Apple's ar tool). If we are close enough to
+    // the end that there cannot possibly be another module, stop looking.
+    if (Stream.getCurrentByteNo() + 8 >= Stream.getBitcodeBytes().size())
+      return Modules;
+
+    BitstreamEntry Entry = Stream.advance();
+    switch (Entry.Kind) {
+    case BitstreamEntry::EndBlock:
+    case BitstreamEntry::Error:
+      return error("Malformed block");
+
+    case BitstreamEntry::SubBlock:
+      if (Entry.ID == bitc::IDENTIFICATION_BLOCK_ID)
+        IdentificationBit = Stream.GetCurrentBitNo();
+      else if (Entry.ID == bitc::MODULE_BLOCK_ID)
+        Modules.push_back({Stream.getBitcodeBytes(),
+                           Buffer.getBufferIdentifier(), IdentificationBit,
+                           Stream.GetCurrentBitNo()});
+
+      if (Stream.SkipBlock())
+        return error("Malformed block");
+      continue;
+    case BitstreamEntry::Record:
+      Stream.skipRecord(Entry.ID);
+      continue;
+    }
+  }
+}
+
 /// \brief Get a lazy one-at-time loading module from bitcode.
 ///
 /// This isn't always used in a lazy context.  In particular, it's also used by
-/// \a parseBitcodeFile().  If this is truly lazy, then we need to eagerly pull
+/// \a parseModule().  If this is truly lazy, then we need to eagerly pull
 /// in forward-referenced functions from block address references.
 ///
 /// \param[in] MaterializeAll Set to \c true if we should materialize
 /// everything.
-static Expected<std::unique_ptr<Module>>
-getLazyBitcodeModuleImpl(MemoryBufferRef Buffer, LLVMContext &Context,
-                         bool MaterializeAll,
-                         bool ShouldLazyLoadMetadata = false) {
-  Expected<BitstreamCursor> StreamOrErr = initStream(Buffer);
-  if (!StreamOrErr)
-    return StreamOrErr.takeError();
+Expected<std::unique_ptr<Module>>
+BitcodeModule::getModuleImpl(LLVMContext &Context, bool MaterializeAll,
+                             bool ShouldLazyLoadMetadata) {
+  BitstreamCursor Stream(Buffer);
+
+  std::string ProducerIdentification;
+  if (IdentificationBit != -1ull) {
+    Stream.JumpToBit(IdentificationBit);
+    Expected<std::string> ProducerIdentificationOrErr =
+        readIdentificationBlock(Stream);
+    if (!ProducerIdentificationOrErr)
+      return ProducerIdentificationOrErr.takeError();
+
+    ProducerIdentification = *ProducerIdentificationOrErr;
+  }
 
-  BitcodeReader *R = new BitcodeReader(std::move(*StreamOrErr), Context);
+  Stream.JumpToBit(ModuleBit);
+  auto *R =
+      new BitcodeReader(std::move(Stream), ProducerIdentification, Context);
 
   std::unique_ptr<Module> M =
-      llvm::make_unique<Module>(Buffer.getBufferIdentifier(), Context);
+      llvm::make_unique<Module>(ModuleIdentifier, Context);
   M->setMaterializer(R);
 
   // Delay parsing Metadata if ShouldLazyLoadMetadata is true.
   if (Error Err = R->parseBitcodeInto(M.get(), ShouldLazyLoadMetadata))
     return std::move(Err);
 
   if (MaterializeAll) {
     // Read in the entire module, and destroy the BitcodeReader.
     if (Error Err = M->materializeAll())
       return std::move(Err);
   } else {
     // Resolve forward references from blockaddresses.
     if (Error Err = R->materializeForwardReferencedFunctions())
       return std::move(Err);
   }
   return std::move(M);
 }
 
 Expected<std::unique_ptr<Module>>
+BitcodeModule::getLazyModule(LLVMContext &Context,
+                             bool ShouldLazyLoadMetadata) {
+  return getModuleImpl(Context, false, ShouldLazyLoadMetadata);
+}
+
+Expected<std::unique_ptr<Module>>
 llvm::getLazyBitcodeModule(MemoryBufferRef Buffer,
                            LLVMContext &Context, bool ShouldLazyLoadMetadata) {
-  return getLazyBitcodeModuleImpl(Buffer, Context, false,
-                                  ShouldLazyLoadMetadata);
+  Expected<std::vector<BitcodeModule>> MsOrErr = getBitcodeModuleList(Buffer);
+  if (!MsOrErr)
+    return MsOrErr.takeError();
+
+  if (MsOrErr->size() != 1)
+    return error("Expected a single module");
+
+  return (*MsOrErr)[0].getLazyModule(Context, ShouldLazyLoadMetadata);
 }
 
 Expected<std::unique_ptr<Module>>
 llvm::getOwningLazyBitcodeModule(std::unique_ptr<MemoryBuffer> &&Buffer,
                                  LLVMContext &Context,
                                  bool ShouldLazyLoadMetadata) {
   auto MOrErr = getLazyBitcodeModule(*Buffer, Context, ShouldLazyLoadMetadata);
   if (MOrErr)
     (*MOrErr)->setOwnedMemoryBuffer(std::move(Buffer));
   return MOrErr;
 }
 
-Expected<std::unique_ptr<Module>> llvm::parseBitcodeFile(MemoryBufferRef Buffer,
-                                                         LLVMContext &Context) {
-  return getLazyBitcodeModuleImpl(Buffer, Context, true);
+Expected<std::unique_ptr<Module>>
+BitcodeModule::parseModule(LLVMContext &Context) {
+  return getModuleImpl(Context, true, false);
   // TODO: Restore the use-lists to the in-memory state when the bitcode was
   // written.  We must defer until the Module has been fully materialized.
 }
 
+Expected<std::unique_ptr<Module>> llvm::parseBitcodeFile(MemoryBufferRef Buffer,
+                                                         LLVMContext &Context) {
+  Expected<std::vector<BitcodeModule>> MsOrErr = getBitcodeModuleList(Buffer);
+  if (!MsOrErr)
+    return MsOrErr.takeError();
+
+  if (MsOrErr->size() != 1)
+    return error("Expected a single module");
+
+  return (*MsOrErr)[0].parseModule(Context);
+}
+
 Expected<std::string> llvm::getBitcodeTargetTriple(MemoryBufferRef Buffer) {
   Expected<BitstreamCursor> StreamOrErr = initStream(Buffer);
   if (!StreamOrErr)
     return StreamOrErr.takeError();
 
   return readTriple(*StreamOrErr);
 }
 

llvm/trunk/test/Bitcode/invalid.test

 RUN:   FileCheck --check-prefix=NON-FUNCTION-EXPLICIT-CALL %s
 RUN: not llvm-dis -disable-output %p/Inputs/invalid-invoke-mismatched-explicit-type.bc 2>&1 | \
 RUN:   FileCheck --check-prefix=MISMATCHED-EXPLICIT-INVOKE %s
 RUN: not llvm-dis -disable-output %p/Inputs/invalid-invoke-non-function-explicit-type.bc 2>&1 | \
 RUN:   FileCheck --check-prefix=NON-FUNCTION-EXPLICIT-INVOKE %s
 
 INVALID-EMPTY: Invalid bitcode signature
 INVALID-ENCODING: Invalid encoding
-BAD-ABBREV: Abbreviation starts with an Array or a Blob
-UNEXPECTED-EOF: Unexpected end of file
-BAD-ABBREV-NUMBER: Invalid abbrev number
+BAD-ABBREV: Malformed block
+UNEXPECTED-EOF: Malformed block
+BAD-ABBREV-NUMBER: Malformed block
 BAD-TYPE-TABLE-FORWARD-REF: Invalid TYPE table: Only named structs can be forward referenced
-BAD-BITWIDTH: Bitwidth for integer type out of range
+BAD-BITWIDTH: Malformed block
 BAD-ALIGN: Invalid alignment value
 MISMATCHED-EXPLICIT-GEP: Explicit gep type does not match pointee type of pointer operand
 MISMATCHED-EXPLICIT-LOAD: Explicit load/store type does not match pointee type of pointer operand
 MISMATCHED-EXPLICIT-GEP-OPERATOR: Explicit gep operator type does not match pointee type of pointer operand
 MISMATCHED-EXPLICIT-CALL: Explicit call type does not match pointee type of callee operand
 NON-FUNCTION-EXPLICIT-CALL: Explicit call type is not a function type
 MISMATCHED-EXPLICIT-INVOKE: Explicit invoke type does not match pointee type of callee operand
 NON-FUNCTION-EXPLICIT-INVOKE: Explicit invoke type is not a function type

 EXTRACT-IDXS: EXTRACTVAL: Invalid type
 INSERT-ARRAY: INSERTVAL: Invalid array index
 INSERT-STRUCT: INSERTVAL: Invalid struct index
 INSERT-IDXS: INSERTVAL: Invalid type
 
 RUN: not llvm-dis -disable-output %p/Inputs/invalid-no-proper-module.bc 2>&1 | \
 RUN:   FileCheck --check-prefix=NO-MODULE %s
 
-NO-MODULE: Malformed IR file
+NO-MODULE: Expected a single module
 
 RUN: not llvm-dis -disable-output %p/Inputs/invalid-fp-shift.bc 2>&1 | \
 RUN:   FileCheck --check-prefix=FP-SHIFT %s
 
 FP-SHIFT: Invalid record
 
 RUN: not llvm-dis -disable-output %p/Inputs/invalid-abbrev-vbr-size-too-big.bc 2>&1 | \
 RUN:   FileCheck --check-prefix=HUGE-ABBREV-OP %s

 RUN: not llvm-dis -disable-output %p/Inputs/invalid-fwdref-type-mismatch.bc 2>&1 | \
 RUN:   FileCheck --check-prefix=FWDREF-TYPE %s
 
 FWDREF-TYPE: Invalid record
 
 RUN: not llvm-dis -disable-output %p/Inputs/invalid-fwdref-type-mismatch-2.bc 2>&1 | \
 RUN:   FileCheck --check-prefix=FWDREF-TYPE-MISMATCH %s
 
-FWDREF-TYPE-MISMATCH: Type mismatch in constant table!
+FWDREF-TYPE-MISMATCH: Malformed block
 
 RUN: not llvm-dis -disable-output %p/Inputs/invalid-array-element-type.bc 2>&1 | \
 RUN:   FileCheck --check-prefix=ELEMENT-TYPE %s
 RUN: not llvm-dis -disable-output %p/Inputs/invalid-vector-element-type.bc 2>&1 | \
 RUN:   FileCheck --check-prefix=ELEMENT-TYPE %s
 RUN: not llvm-dis -disable-output %p/Inputs/invalid-pointer-element-type.bc 2>&1 | \
 RUN:   FileCheck --check-prefix=ELEMENT-TYPE %s
 

 RUN: not llvm-dis -disable-output %p/Inputs/invalid-extract-0-indices.bc 2>&1 | \
 RUN:   FileCheck --check-prefix=EXTRACT-0-IDXS %s
 
 EXTRACT-0-IDXS: EXTRACTVAL: Invalid instruction with 0 indices
 
 RUN: not llvm-dis -disable-output %p/Inputs/invalid-load-ptr-type.bc 2>&1 | \
 RUN:   FileCheck --check-prefix=BAD-LOAD-PTR-TYPE %s
 
-BAD-LOAD-PTR-TYPE: Cannot load/store from pointer
+BAD-LOAD-PTR-TYPE: Malformed block
 
 RUN: not llvm-dis -disable-output %p/Inputs/invalid-inserted-value-type-mismatch.bc 2>&1 | \
 RUN:   FileCheck --check-prefix=INSERT-TYPE-MISMATCH %s
 
 INSERT-TYPE-MISMATCH: Inserted value type doesn't match aggregate type
 
 RUN: not llvm-dis -disable-output %p/Inputs/invalid-code-len-width.bc 2>&1 | \
 RUN:   FileCheck --check-prefix=INVALID-CODELENWIDTH %s
 
 INVALID-CODELENWIDTH: Malformed block
 
 RUN: not llvm-dis -disable-output %p/Inputs/invalid-function-argument-type.bc 2>&1 | \
 RUN:   FileCheck --check-prefix=INVALID-ARGUMENT-TYPE %s
 
 INVALID-ARGUMENT-TYPE: Invalid function argument type
 
 RUN: not llvm-dis -disable-output %p/Inputs/invalid-function-comdat-id.bc 2>&1 | \
 RUN:   FileCheck --check-prefix=INVALID-FCOMDAT-ID %s
 
-INVALID-FCOMDAT-ID: Invalid function comdat ID
+INVALID-FCOMDAT-ID: Malformed block
 
 RUN: not llvm-dis -disable-output %p/Inputs/invalid-global-var-comdat-id.bc 2>&1 | \
 RUN:   FileCheck --check-prefix=INVALID-GVCOMDAT-ID %s
 
 INVALID-GVCOMDAT-ID: Invalid global variable comdat ID
 
 RUN: not llvm-dis -disable-output %p/Inputs/invalid-abbrev-no-operands.bc 2>&1 | \
 RUN:   FileCheck --check-prefix=ABBREV-NO-OPS %s
 
 ABBREV-NO-OPS: Abbrev record with no operands
 
 RUN: not llvm-dis -disable-output %p/Inputs/invalid-array-operand-encoding.bc 2>&1 | \
 RUN:   FileCheck --check-prefix=ARRAY-OP-ENC %s
 
-ARRAY-OP-ENC: Array element type has to be an encoding of a type
+ARRAY-OP-ENC: Malformed block
 
 RUN: not llvm-dis -disable-output %p/Inputs/invalid-metadata-not-followed-named-node.bc 2>&1 | \
 RUN:   FileCheck --check-prefix=META-NOT-FOLLOWED-BY-NAMED-META %s
 
-META-NOT-FOLLOWED-BY-NAMED-META: METADATA_NAME not followed by METADATA_NAMED_NODE
+META-NOT-FOLLOWED-BY-NAMED-META: Malformed block
 
 RUN: not llvm-dis -disable-output %p/Inputs/invalid-vector-length.bc 2>&1 | \
 RUN:   FileCheck --check-prefix=VECTOR-LENGTH %s
 
 VECTOR-LENGTH: Invalid vector length
 
 RUN: not llvm-dis -disable-output %p/Inputs/invalid-alias-type-mismatch.bc 2>&1 | \
 RUN:   FileCheck --check-prefix=ALIAS-TYPE-MISMATCH %s
 
 ALIAS-TYPE-MISMATCH: Alias and aliasee types don't match
 
 RUN: not llvm-dis -disable-output %p/Inputs/invalid-no-function-block.bc 2>&1 | \
 RUN:   FileCheck --check-prefix=NO-FUNCTION-BLOCK %s
 
 NO-FUNCTION-BLOCK: Trying to materialize functions before seeing function blocks (Producer: 'LLVM3.8.0git' Reader: 'LLVM
 
 RUN: not llvm-dis -disable-output %p/Inputs/invalid-name-with-0-byte.bc 2>&1 | \
 RUN:   FileCheck --check-prefix=NAME-WITH-0 %s
 
-NAME-WITH-0: Invalid value name
+NAME-WITH-0: Malformed block
 
 RUN: not llvm-dis -disable-output %p/Inputs/invalid-void-constant.bc 2>&1 | \
 RUN:   FileCheck --check-prefix=VOID-CONSTANT-TYPE %s
 
 VOID-CONSTANT-TYPE: Invalid constant type
 
 RUN: not llvm-dis -disable-output %p/Inputs/invalid-gep-no-operands.bc 2>&1 | \
 RUN:   FileCheck --check-prefix=GEP-NO-OPERANDS %s

llvm/trunk/test/Bitcode/null-type.ll

 ; RUN: not llvm-dis < %s.bc 2>&1 | FileCheck %s
 ; PR8494
 
-; CHECK: Invalid record
+; CHECK: Malformed block