[Bitcode] Add BITCODE_SIZE_BLOCK_ID to encode the size of the bitcode

[MaskRay]

Summary

Currently a bitcode file has no terminator. If several bitcode files are
concatenated, there is no suitable tool splitting them. For example:

clang -fembed-bitcode=all -c a.c b.c
ld a.o b.o
objcopy --dump-section=.llvmbc=a.bc a.out /dev/null
llvm-dis a.bc  # LLVM ERROR: Invalid encoding

This is because llvm-dis does not terminate after reaching the end of the first
stream, and decodes the start of the second bitcode file as the continuation
instead, and fail.

This patch introduces an optional BITCODE_SIZE_BLOCK_ID block to encode the size
of the stream. The block can be placed anywhere in the stream. For a streaming
writer, BITCODE_SIZE_BLOCK_ID can be placed at the end as a terminator. If
random access is supported, BITCODE_SIZE_BLOCK_ID can be placed earlier to allow
fast filtering (by parsing some metadata records like MODULE_CODE_SOURCE_FILENAME).

llvm-dis is taught to handle concatenated streams.

[dougpuob]

llvm/lib/Bitcode/Reader/BitcodeAnalyzer.cpp

84

Seems like a typo.

llvm/lib/Bitcode/Reader/BitcodeReader.cpp

6504–6505

About the boundary checking issue, unsigned integer wrapping. Is it possible that the BCBegin is greater than Size? if so, is it possible that the `return error() path will be executed?

I'm a naive to compiler, just curious about it.

[MaskRay]

llvm/lib/Bitcode/Reader/BitcodeReader.cpp

6504–6505

The check is to catch erroneous cases (if Size in BITCODE_SIZE_BLOCK_ID block is incorrectly set). Otherwise it is difficult to catch such errors.

[dougpuob]

llvm/lib/Bitcode/Reader/BitcodeReader.cpp

6504–6505

Got it, thank you :)

[steven_wu]

I am a bit worry that linker might concatenate bitcode file with padding to achieve alignment requirement, etc. I guess you can create a termination block to mark the end but it is hard to seek the next start.

[MaskRay]

In D86847#2248353, @steven_wu wrote:

I am a bit worry that linker might concatenate bitcode file with padding to achieve alignment requirement, etc. I guess you can create a termination block to mark the end but it is hard to seek the next start.

We can set the section alignment to 1 to avoid padding: (I rushed a bit, sorry: 6ae7b403c3e1aebcb825d3dd4777d3c1149d6d67)

[dblaikie]

@bartell mentioned something about a bitcode wrapper that does provide the length - how does this patch's strategy compare to that approach? Are they redundant? Should the wrapper be replaced with this patch's approach? (or should the bitcode wrapper approach be used instead of adding this patch?)

[steven_wu]

In D86847#2248354, @MaskRay wrote:

In D86847#2248353, @steven_wu wrote:

I am a bit worry that linker might concatenate bitcode file with padding to achieve alignment requirement, etc. I guess you can create a termination block to mark the end but it is hard to seek the next start.

We can set the section alignment to 1 to avoid padding: (I rushed a bit, sorry: 6ae7b403c3e1aebcb825d3dd4777d3c1149d6d67)

I didn't see that. I don't really have concerns then.

In D86847#2248516, @dblaikie wrote:

@bartell mentioned something about a bitcode wrapper that does provide the length - how does this patch's strategy compare to that approach? Are they redundant? Should the wrapper be replaced with this patch's approach? (or should the bitcode wrapper approach be used instead of adding this patch?)

I guess that is an implementation choice. This approach prefers a simple linker implementation, while the bitcode wrapper requires linker know to treat this section differently and write bitcode wrapper into section.
@MaskRay Do you have any specific use case in mind for this? If you are tied to a linker like lld, it might be better just teach lld to treat this section differently so we don't need worry about concatenated bitcode file.

[bartell]

In D86847#2248516, @dblaikie wrote:

@bartell mentioned something about a bitcode wrapper that does provide the length - how does this patch's strategy compare to that approach? Are they redundant? Should the wrapper be replaced with this patch's approach? (or should the bitcode wrapper approach be used instead of adding this patch?)

I guess that is an implementation choice. This approach prefers a simple linker implementation, while the bitcode wrapper requires linker know to treat this section differently and write bitcode wrapper into section.
@MaskRay Do you have any specific use case in mind for this? If you are tied to a linker like lld, it might be better just teach lld to treat this section differently so we don't need worry about concatenated bitcode file.

The linker wouldn't need to know about the wrapper. We would just change Clang to always emit the wrapper in .llvmbc.

The "bitcode wrapper" I'm referring to is the simple header written by emitDarwinBCHeaderAndTrailer and read by SkipBitcodeWrapperHeader. It's currently only written for bitcode files that use Darwin/Mach-O targets, but it should be read successfully regardless of target.

Rather than adding new fields to the bitcode itself as this patch does, I would recommend using the "bitcode wrapper" instead, just for the sake of reusing code. We could either enable it for all uses of WriteBitcodeToFile, or just use it in Clang when generating the .llvmbc section.

[MaskRay]

In D86847#2250027, @steven_wu wrote:

In D86847#2248354, @MaskRay wrote:

In D86847#2248353, @steven_wu wrote:

I am a bit worry that linker might concatenate bitcode file with padding to achieve alignment requirement, etc. I guess you can create a termination block to mark the end but it is hard to seek the next start.

We can set the section alignment to 1 to avoid padding: (I rushed a bit, sorry: 6ae7b403c3e1aebcb825d3dd4777d3c1149d6d67)

I didn't see that. I don't really have concerns then.

In D86847#2248516, @dblaikie wrote:

@bartell mentioned something about a bitcode wrapper that does provide the length - how does this patch's strategy compare to that approach? Are they redundant? Should the wrapper be replaced with this patch's approach? (or should the bitcode wrapper approach be used instead of adding this patch?)

I guess that is an implementation choice. This approach prefers a simple linker implementation, while the bitcode wrapper requires linker know to treat this section differently and write bitcode wrapper into section.
@MaskRay Do you have any specific use case in mind for this? If you are tied to a linker like lld, it might be better just teach lld to treat this section differently so we don't need worry about concatenated bitcode file.

It is not tied to a linker like lld. I believe in most binary formats, if you specify an alignment of 1, their linkers will concatenate input sections in the output without padding. Actually, the section content of .llvmbc is entirely opaque to lld/GNU ld/gold. ("dumb linker, smart format" design).

If the bitcode wrapper requires the linker to understand its format, I'd vote against that solution. (It will not work with GNU ld/gold.)

[tejohnson]

In D86847#2250194, @MaskRay wrote:

In D86847#2250027, @steven_wu wrote:

In D86847#2248354, @MaskRay wrote:

In D86847#2248353, @steven_wu wrote:

I am a bit worry that linker might concatenate bitcode file with padding to achieve alignment requirement, etc. I guess you can create a termination block to mark the end but it is hard to seek the next start.

We can set the section alignment to 1 to avoid padding: (I rushed a bit, sorry: 6ae7b403c3e1aebcb825d3dd4777d3c1149d6d67)

I didn't see that. I don't really have concerns then.

In D86847#2248516, @dblaikie wrote:

@bartell mentioned something about a bitcode wrapper that does provide the length - how does this patch's strategy compare to that approach? Are they redundant? Should the wrapper be replaced with this patch's approach? (or should the bitcode wrapper approach be used instead of adding this patch?)

I guess that is an implementation choice. This approach prefers a simple linker implementation, while the bitcode wrapper requires linker know to treat this section differently and write bitcode wrapper into section.
@MaskRay Do you have any specific use case in mind for this? If you are tied to a linker like lld, it might be better just teach lld to treat this section differently so we don't need worry about concatenated bitcode file.

It is not tied to a linker like lld. I believe in most binary formats, if you specify an alignment of 1, their linkers will concatenate input sections in the output without padding. Actually, the section content of .llvmbc is entirely opaque to lld/GNU ld/gold. ("dumb linker, smart format" design).

If the bitcode wrapper requires the linker to understand its format, I'd vote against that solution. (It will not work with GNU ld/gold.)

Per @bartell 's last comment, it sounds like it doesn't require the linker to understand the format (the handling is on the bitcode reader/writer side). Do we know whether that would work here instead?

[MaskRay]

The idea is to add a terminator block, but a terminator carries less information so this patch intends to add a size block instead. With a terminator, the following problem can be better addressed

// 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 the container does not expect more than one bitcode stream. Then after passing one stream, it can bail out.
For users who want to have multiple streams, they can strip any padding. The scheme will automatically work with most binary formats (alignment=1).

[dblaikie]

In D86847#2278056, @MaskRay wrote:

The idea is to add a terminator block, but a terminator carries less information so this patch intends to add a size block instead. With a terminator, the following problem can be better addressed

// 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 the container does not expect more than one bitcode stream. Then after passing one stream, it can bail out.
For users who want to have multiple streams, they can strip any padding. The scheme will automatically work with most binary formats (alignment=1).

It'd be nice to not have two different tools to solve what seem to be fairly similar problems. Perhaps the MachO/bitcode wrapper could be simplified if the alignment issue was fixed there? (not sure if it needs the full header, etc) & be a simple prefix length, perhaps. I'm probably a bit biased by DWARF, but it seems nice to be able to read a length fairly quickly/easily and skip the rest, rather than having to parse the contents to find out where it ends.

[MaskRay]

BitcodeReader.h has the following code

/// SkipBitcodeWrapperHeader - Some systems wrap bc files with a special
/// header for padding or other reasons.  The format of this header is:
///
/// struct bc_header {
///   uint32_t Magic;         // 0x0B17C0DE
///   uint32_t Version;       // Version, currently always 0.
///   uint32_t BitcodeOffset; // Offset to traditional bitcode file.
///   uint32_t BitcodeSize;   // Size of traditional bitcode file.
///   ... potentially other gunk ...
/// };

The header contains the size so I assume this patch is not needed.

llvm/include/llvm/Bitcode/BitcodeAnalyzer.h

 };
 
 class BitcodeAnalyzer {
   BitstreamCursor Stream;
   BitstreamBlockInfo BlockInfo;
   CurStreamTypeType CurStreamType;
   Optional<BitstreamCursor> BlockInfoStream;
   unsigned NumTopBlocks = 0;
+  uint64_t Size = UINT64_MAX;
 
   struct PerRecordStats {
     unsigned NumInstances;
     unsigned NumAbbrev;
     uint64_t TotalBits;
     PerRecordStats() : NumInstances(0), NumAbbrev(0), TotalBits(0) {}
   };
 

llvm/include/llvm/Bitcode/BitcodeReader.h

class BitcodeModule {
 
     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<BitcodeFileContents>
-    getBitcodeFileContents(MemoryBufferRef Buffer);
+    getBitcodeFileContentsFromPrefix(MemoryBufferRef &Buffer);
 
     Expected<std::unique_ptr<Module>>
     getModuleImpl(LLVMContext &Context, bool MaterializeAll,
                   bool ShouldLazyLoadMetadata, bool IsImporting,
                   DataLayoutCallbackTy DataLayoutCallback);
 
   public:
     StringRef getBuffer() const {
struct BitcodeFileContents {
     std::vector<BitcodeModule> Mods;
     StringRef Symtab, StrtabForSymtab;
   };
 
   /// Returns the contents of a bitcode file. This includes the raw contents of
   /// the symbol table embedded in the bitcode file. Clients which require a
   /// symbol table should prefer to use irsymtab::read instead of this function
   /// because it creates a reader for the irsymtab and handles upgrading bitcode
-  /// files without a symbol table or with an old symbol table.
+  /// files without a symbol table or with an old symbol table.  Buffer is
+  /// modified to represent the unparsed part.
+  Expected<BitcodeFileContents>
+  getBitcodeFileContentsFromPrefix(MemoryBufferRef &Buffer);
+
+  /// Returns the contents of a bitcode file, assuming there is only one stream.
   Expected<BitcodeFileContents> getBitcodeFileContents(MemoryBufferRef Buffer);
 
   /// 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,

llvm/include/llvm/Bitcode/BitcodeWriter.h

class BitcodeWriter {
     std::vector<Module *> Mods;
 
   public:
     /// Create a BitcodeWriter that writes to Buffer.
     BitcodeWriter(SmallVectorImpl<char> &Buffer);
 
     ~BitcodeWriter();
 
+    /// Emit a BITCODE_SIZE_BLOCK_ID block which encodes a 64-bit integer
+    /// representing the size (in bytes) of the bitstream.  This is used to mark
+    /// the end of the bitstream. At the end of the bitstream, call
+    /// backfillSizeBlock to fill in the size field.
+    uint64_t reserveSizeBlock();
+    void backfillSizeBlock(uint64_t SizeBitPos);
+
     /// Attempt to write a symbol table to the bitcode file. This must be called
     /// at most once after all modules have been written.
     ///
     /// A reader does not require a symbol table to interpret a bitcode file;
     /// the symbol table is needed only to improve link-time performance. So
     /// this function may decide not to write a symbol table. It may so decide
     /// if, for example, the target is unregistered or the IR is malformed.
     void writeSymtab();

llvm/include/llvm/Bitcode/LLVMBitCodes.h

enum BlockIDs {
 
   STRTAB_BLOCK_ID,
 
   FULL_LTO_GLOBALVAL_SUMMARY_BLOCK_ID,
 
   SYMTAB_BLOCK_ID,
 
   SYNC_SCOPE_NAMES_BLOCK_ID,
+
+  BITCODE_SIZE_BLOCK_ID,
 };
 
 /// Identification block contains a string that describes the producer details,
 /// and an epoch that defines the auto-upgrade capability.
 enum IdentificationCodes {
   IDENTIFICATION_CODE_STRING = 1, // IDENTIFICATION:      [strchr x N]
   IDENTIFICATION_CODE_EPOCH = 2,  // EPOCH:               [epoch#]
 };

llvm/include/llvm/Bitstream/BitstreamWriter.h

class BitstreamWriter {
   }
 
   void WriteWord(unsigned Value) {
     Value = support::endian::byte_swap<uint32_t, support::little>(Value);
     Out.append(reinterpret_cast<const char *>(&Value),
                reinterpret_cast<const char *>(&Value + 1));
   }
 
-  size_t GetBufferOffset() const { return Out.size(); }
-
   size_t GetWordIndex() const {
     size_t Offset = GetBufferOffset();
     assert((Offset & 3) == 0 && "Not 32-bit aligned");
     return Offset / 4;
   }
 
 public:
   explicit BitstreamWriter(SmallVectorImpl<char> &O)
     : Out(O), CurBit(0), CurValue(0), CurCodeSize(2) {}
 
   ~BitstreamWriter() {
     assert(CurBit == 0 && "Unflushed data remaining");
     assert(BlockScope.empty() && CurAbbrevs.empty() && "Block imbalance");
   }
 
+  size_t GetBufferOffset() const { return Out.size(); }
+
   /// Retrieve the current position in the stream, in bits.
   uint64_t GetCurrentBitNo() const { return GetBufferOffset() * 8 + CurBit; }
 
   /// Retrieve the number of bits currently used to encode an abbrev ID.
   unsigned GetAbbrevIDWidth() const { return CurCodeSize; }
 
   //===--------------------------------------------------------------------===//
   // Basic Primitives for emitting bits to the stream.

llvm/lib/Bitcode/Reader/BitcodeAnalyzer.cpp

static Optional<const char *> GetBlockName(unsigned BlockID,
   case bitc::FULL_LTO_GLOBALVAL_SUMMARY_BLOCK_ID:
     return "FULL_LTO_GLOBALVAL_SUMMARY_BLOCK";
   case bitc::MODULE_STRTAB_BLOCK_ID:
     return "MODULE_STRTAB_BLOCK";
   case bitc::STRTAB_BLOCK_ID:
     return "STRTAB_BLOCK";
   case bitc::SYMTAB_BLOCK_ID:
     return "SYMTAB_BLOCK";
+  case bitc::BITCODE_SIZE_BLOCK_ID:
+    return "BITCODE_SIZE_BLOCK_ID";
   }
 }
 
 /// Return a symbolic code name if known, otherwise return null.
 static Optional<const char *> GetCodeName(unsigned CodeID, unsigned BlockID,
                                           const BitstreamBlockInfo &BlockInfo,
                                           CurStreamTypeType CurStreamType) {
   // Standard blocks for all bitcode files.
while (!BlockInfoCursor.AtEndOfStream()) {
       }
 
       if (Error Err = BlockInfoCursor.SkipBlock())
         return Err;
     }
   }
 
   // Parse the top-level structure.  We only allow blocks at the top-level.
-  while (!Stream.AtEndOfStream()) {
+  // Stop if BITCODE_SIZE_BLOCK_ID exists and the cursor is at the end.
+  while (!Stream.AtEndOfStream() && Stream.getCurrentByteNo() != Size) {
     Expected<unsigned> MaybeCode = Stream.ReadCode();
     if (!MaybeCode)
       return MaybeCode.takeError();
     if (MaybeCode.get() != bitc::ENTER_SUBBLOCK)
       return reportError("Invalid record at top-level");
 
     Expected<unsigned> MaybeBlockID = Stream.ReadSubBlockID();
     if (!MaybeBlockID)
if (DumpRecords) {
 
     if (!O->Symbolic && BlockName)
       O->OS << " BlockID=" << BlockID;
 
     O->OS << " NumWords=" << NumWords
           << " BlockCodeSize=" << Stream.getAbbrevIDWidth() << ">\n";
   }
 
+  if (BlockID == bitc::BITCODE_SIZE_BLOCK_ID) {
+    Expected<uint32_t> MaybeSize = Stream.Read(32);
+    if (!MaybeSize)
+      return MaybeSize.takeError();
+    Size = *MaybeSize;
+    MaybeSize = Stream.Read(32);
+    if (!MaybeSize)
+      return MaybeSize.takeError();
+    Size |= static_cast<uint64_t>(*MaybeSize) << 32;
+
+    Expected<BitstreamEntry> MaybeEntry = Stream.advance();
+    if (!MaybeEntry)
+      return MaybeEntry.takeError();
+    if (MaybeEntry.get().Kind != BitstreamEntry::EndBlock)
+      return reportError("malformed BITCODE_SIZE_BLOCK_ID");
+    if (DumpRecords) {
+      O->OS << Indent << "  <SIZE " << Size << "/>\n";
+      O->OS << Indent << "</BITCODE_SIZE_BLOCK_ID>\n";
+    }
+    return Error::success();
+  }
+
   SmallVector<uint64_t, 64> Record;
 
   // Keep the offset to the metadata index if seen.
   uint64_t MetadataIndexOffset = 0;
 
   // Read all the records for this block.
   while (1) {
     if (Stream.AtEndOfStream())

llvm/lib/Bitcode/Reader/BitcodeReader.cpp

 llvm::getBitcodeModuleList(MemoryBufferRef Buffer) {
   auto FOrErr = getBitcodeFileContents(Buffer);
   if (!FOrErr)
     return FOrErr.takeError();
   return std::move(FOrErr->Mods);
 }
 
 Expected<BitcodeFileContents>
-llvm::getBitcodeFileContents(MemoryBufferRef Buffer) {
+llvm::getBitcodeFileContentsFromPrefix(MemoryBufferRef &Buffer) {
   Expected<BitstreamCursor> StreamOrErr = initStream(Buffer);
   if (!StreamOrErr)
     return StreamOrErr.takeError();
   BitstreamCursor &Stream = *StreamOrErr;
 
   BitcodeFileContents F;
+  uint64_t Size = UINT64_MAX;
   while (true) {
     uint64_t BCBegin = Stream.getCurrentByteNo();
+    if (BCBegin >= Size) {
+      if (BCBegin == Size) {
+        Buffer = MemoryBufferRef(Buffer.getBuffer().substr(BCBegin),
+                                 Buffer.getBufferIdentifier());
+        return F;
+      }
+      return error("byte position 0x" + Twine::utohexstr(BCBegin) +
+                   " is greater than bitcode size 0x" + Twine::utohexstr(Size));
+    }
 
     // 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 (BCBegin + 8 >= Stream.getBitcodeBytes().size())
+    if (BCBegin + 8 >= Stream.getBitcodeBytes().size()) {
+      Buffer = MemoryBufferRef("", Buffer.getBufferIdentifier());
       return F;
+    }
 
     Expected<llvm::BitstreamEntry> MaybeEntry = Stream.advance();
     if (!MaybeEntry)
       return MaybeEntry.takeError();
     llvm::BitstreamEntry Entry = MaybeEntry.get();
 
     switch (Entry.Kind) {
     case BitstreamEntry::EndBlock:
     case BitstreamEntry::Error:
       return error("Malformed block");
 
     case BitstreamEntry::SubBlock: {
+      // A BITCODE_SIZE_BLOCK_ID block encodes a 64-bit integer representing the
+      // size (in bytes) of the bitstream. This is used to mark the end of the
+      // bitstream.
+      if (Entry.ID == bitc::BITCODE_SIZE_BLOCK_ID) {
+        if (Error Err = Stream.EnterSubBlock(Entry.ID))
+          return std::move(Err);
+        Expected<uint32_t> MaybeSize = Stream.Read(32);
+        if (!MaybeSize)
+          return MaybeSize.takeError();
+        Size = *MaybeSize;
+        MaybeSize = Stream.Read(32);
+        if (!MaybeSize)
+          return MaybeSize.takeError();
+        Size |= static_cast<uint64_t>(*MaybeSize) << 32;
+
+        MaybeEntry = Stream.advance();
+        if (!MaybeEntry)
+          return MaybeEntry.takeError();
+        if (MaybeEntry.get().Kind != BitstreamEntry::EndBlock)
+          return error("Malformed block");
+        continue;
+      }
+
       uint64_t IdentificationBit = -1ull;
       if (Entry.ID == bitc::IDENTIFICATION_BLOCK_ID) {
         IdentificationBit = Stream.GetCurrentBitNo() - BCBegin * 8;
         if (Error Err = Stream.SkipBlock())
           return std::move(Err);
 
         {
           Expected<llvm::BitstreamEntry> MaybeEntry = Stream.advance();
case BitstreamEntry::Record:
       if (Expected<unsigned> StreamFailed = Stream.skipRecord(Entry.ID))
         continue;
       else
         return StreamFailed.takeError();
     }
   }
 }
 
+Expected<BitcodeFileContents>
+llvm::getBitcodeFileContents(MemoryBufferRef Buffer) {
+  return getBitcodeFileContentsFromPrefix(Buffer);
+}
+
 /// 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 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.

llvm/lib/Bitcode/Writer/BitcodeWriter.cpp

void BitcodeWriter::writeBlob(unsigned Block, unsigned Record, StringRef Blob) {
   Abbv->Add(BitCodeAbbrevOp(BitCodeAbbrevOp::Blob));
   auto AbbrevNo = Stream->EmitAbbrev(std::move(Abbv));
 
   Stream->EmitRecordWithBlob(AbbrevNo, ArrayRef<uint64_t>{Record}, Blob);
 
   Stream->ExitBlock();
 }
 
+uint64_t BitcodeWriter::reserveSizeBlock() {
+  Stream->EnterSubblock(bitc::BITCODE_SIZE_BLOCK_ID, 2);
+  const uint64_t Pos = Stream->GetCurrentBitNo();
+  Stream->Emit(0, 32);
+  Stream->Emit(0, 32);
+  Stream->ExitBlock();
+  return Pos;
+}
+
+void BitcodeWriter::backfillSizeBlock(uint64_t SizeBitPos) {
+  Stream->BackpatchWord64(SizeBitPos, Stream->GetBufferOffset());
+}
+
 void BitcodeWriter::writeSymtab() {
   assert(!WroteStrtab && !WroteSymtab);
 
   // If any module has module-level inline asm, we will require a registered asm
   // parser for the target so that we can create an accurate symbol table for
   // the module.
   for (Module *M : Mods) {
     if (M->getModuleInlineAsm().empty())
void llvm::WriteBitcodeToFile(const Module &M, raw_ostream &Out,
 
   // If this is darwin or another generic macho target, reserve space for the
   // header.
   Triple TT(M.getTargetTriple());
   if (TT.isOSDarwin() || TT.isOSBinFormatMachO())
     Buffer.insert(Buffer.begin(), BWH_HeaderSize, 0);
 
   BitcodeWriter Writer(Buffer);
+  const uint64_t SizeBitPos = Writer.reserveSizeBlock();
   Writer.writeModule(M, ShouldPreserveUseListOrder, Index, GenerateHash,
                      ModHash);
   Writer.writeSymtab();
   Writer.writeStrtab();
+  // Backfill the BITCODE_SIZE_BLOCK_ID block after the stream is complete.
+  Writer.backfillSizeBlock(SizeBitPos);
 
   if (TT.isOSDarwin() || TT.isOSBinFormatMachO())
     emitDarwinBCHeaderAndTrailer(Buffer, TT);
 
   // Write the generated bitstream to "Out".
   Out.write((char*)&Buffer.front(), Buffer.size());
 }
 

llvm/test/Assembler/multi-mod-disassemble.ll

-; RUN: llvm-as %s -o %t.o
-; RUN: llvm-cat -b -o %t2.o %t.o %t.o
-; RUN: llvm-dis -o %t3 %t2.o
-; RUN: llvm-as -o %t4.o %t3.0
-; RUN: llvm-as -o %t5.o %t3.1
-; RUN: cmp %t4.o %t5.o
-; RUN: llvm-cat -b -o %t6.o %t5.o %t4.o
-; RUN: llvm-dis -o %t7.o %t6.o
-; RUN: diff %t7.o.0 %t7.o.1
-; RUN: FileCheck < %t7.o.0 %s
-; RUN: FileCheck < %t7.o.1 %s
+; RUN: rm -rf %t && mkdir %t && cd %t
+
+;; A single bitstream with multiple modules can be disassembled.
+; RUN: llvm-as %s -o a.bc
+; RUN: llvm-cat -b -o a2.bc a.bc a.bc
+; RUN: llvm-dis -o a3 a2.bc
+; RUN: llvm-as -o a3.0.bc a3.0
+; RUN: llvm-as -o a3.1.bc a3.1
+; RUN: cmp a3.0.bc a3.1.bc
+
+; RUN: llvm-cat -b -o a4.bc a3.1.bc a3.0.bc
+; RUN: llvm-dis -o a5.bc a4.bc
+; RUN: diff a5.bc.0 a5.bc.1
+; RUN: FileCheck < a5.bc.0 %s
+; RUN: FileCheck < a5.bc.1 %s
 
 ; CHECK: source_filename = "{{.*}}multi-mod-disassemble.ll"
+
+;; Concatenated bitstreams can be dumped. For a component bitstream with multiple modules,
+;; the output filenames get an additional `.N` suffixes.
+; RUN: cat a.bc a2.bc > a6.bc
+; RUN: llvm-dis -o a7 a6.bc
+; RUN: ls a7.0 a7.1.0 a7.1.1
+; RUN: llvm-dis < a6.bc | FileCheck %s --check-prefix=CONCAT
+
+; CONCAT: source_filename = "{{.*}}multi-mod-disassemble.ll"
+; CONCAT: source_filename = "{{.*}}multi-mod-disassemble.ll"
+; CONCAT: source_filename = "{{.*}}multi-mod-disassemble.ll"
+
+;; llvm-bcanalyzer only dumps the prefix.
+; RUN: llvm-bcanalyzer --dump a6.bc | FileCheck %s --check-prefix=CONCAT-DUMP --match-full-lines --strict-whitespace
+
+;      CONCAT-DUMP:<BITCODE_SIZE_BLOCK_ID NumWords=3 BlockCodeSize=2>
+; CONCAT-DUMP-NEXT:  <SIZE [[# SIZE:]]/>
+; CONCAT-DUMP-NEXT:</BITCODE_SIZE_BLOCK_ID>
+; CONCAT-DUMP-NEXT:<IDENTIFICATION_BLOCK_ID {{.*}}
+;  CONCAT-DUMP-NOT:<BITCODE_SIZE_BLOCK_ID
+
 target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
 target triple = "x86_64-unknown-linux-gnu"

llvm/tools/llvm-dis/llvm-dis.cpp

if (DI.getSeverity() == DS_Error)
       exit(1);
     return true;
   }
 };
 } // end anon namespace
 
 static ExitOnError ExitOnErr;
 
-int main(int argc, char **argv) {
-  InitLLVM X(argc, argv);
-
-  ExitOnErr.setBanner(std::string(argv[0]) + ": error: ");
-
-  LLVMContext Context;
-  Context.setDiagnosticHandler(
-      std::make_unique<LLVMDisDiagnosticHandler>(argv[0]));
-  cl::ParseCommandLineOptions(argc, argv, "llvm .bc -> .ll disassembler\n");
-
-  std::unique_ptr<MemoryBuffer> MB =
-      ExitOnErr(errorOrToExpected(MemoryBuffer::getFileOrSTDIN(InputFilename)));
-
-  BitcodeFileContents IF = ExitOnErr(llvm::getBitcodeFileContents(*MB));
-
-  const size_t N = IF.Mods.size();
-
-  if (OutputFilename == "-" && N > 1)
-      errs() << "only single module bitcode files can be written to stdout\n";
-
-  for (size_t i = 0; i < N; ++i) {
+static bool disassemble(int Nth, const BitcodeFileContents &IF,
+                        LLVMContext &Ctx) {
+  for (size_t i = 0, N = IF.Mods.size(); i != N; ++i) {
     BitcodeModule MB = IF.Mods[i];
-    std::unique_ptr<Module> M = ExitOnErr(MB.getLazyModule(Context, MaterializeMetadata,
-                                          SetImporting));
+    std::unique_ptr<Module> M =
+        ExitOnErr(MB.getLazyModule(Ctx, MaterializeMetadata, SetImporting));
     if (MaterializeMetadata)
       ExitOnErr(M->materializeMetadata());
     else
       ExitOnErr(M->materializeAll());
 
     BitcodeLTOInfo LTOInfo = ExitOnErr(MB.getLTOInfo());
     std::unique_ptr<ModuleSummaryIndex> Index;
     if (LTOInfo.HasSummary)
       Index = ExitOnErr(MB.getSummary());
 
-    std::string FinalFilename(OutputFilename);
+    std::string OutName(OutputFilename);
 
     // Just use stdout.  We won't actually print anything on it.
-    if (DontPrint)
-      FinalFilename = "-";
-
-    if (FinalFilename.empty()) { // Unspecified output, infer it.
-      if (InputFilename == "-") {
-        FinalFilename = "-";
-      } else {
+    if (DontPrint || (OutName.empty() && InputFilename == "-")) {
+      OutName = "-";
+      if (OutName.empty()) {
         StringRef IFN = InputFilename;
-        FinalFilename = (IFN.endswith(".bc") ? IFN.drop_back(3) : IFN).str();
-        if (N > 1)
-          FinalFilename += std::string(".") + std::to_string(i);
-        FinalFilename += ".ll";
+        OutName = (IFN.endswith(".bc") ? IFN.drop_back(3) : IFN).str();
       }
-    } else {
+      if (Nth >= 0)
+        OutName += ("." + Twine(Nth)).str();
       if (N > 1)
-        FinalFilename += std::string(".") + std::to_string(i);
+        OutName += ("." + Twine(i)).str();
+      if (OutputFilename.empty())
+        OutName += ".ll";
     }
 
     std::error_code EC;
     std::unique_ptr<ToolOutputFile> Out(
-        new ToolOutputFile(FinalFilename, EC, sys::fs::OF_Text));
+        new ToolOutputFile(OutName, EC, sys::fs::OF_Text));
     if (EC) {
       errs() << EC.message() << '\n';
-      return 1;
+      return false;
     }
 
     std::unique_ptr<AssemblyAnnotationWriter> Annotator;
     if (ShowAnnotations)
       Annotator.reset(new CommentWriter());
 
     // All that llvm-dis does is write the assembly to a file.
     if (!DontPrint) {
       M->print(Out->os(), Annotator.get(), PreserveAssemblyUseListOrder);
       if (Index)
         Index->print(Out->os());
     }
 
     // Declare success.
     Out->keep();
   }
+  return true;
+}
+
+int main(int argc, char **argv) {
+  InitLLVM X(argc, argv);
+
+  ExitOnErr.setBanner(std::string(argv[0]) + ": error: ");
+
+  LLVMContext Context;
+  Context.setDiagnosticHandler(
+      std::make_unique<LLVMDisDiagnosticHandler>(argv[0]));
+  cl::ParseCommandLineOptions(argc, argv, "llvm .bc -> .ll disassembler\n");
+
+  std::unique_ptr<MemoryBuffer> Buffer =
+      ExitOnErr(errorOrToExpected(MemoryBuffer::getFileOrSTDIN(InputFilename)));
+  MemoryBufferRef MB(*Buffer);
+
+  int Nth = 0;
+  bool Stop;
+  do {
+    BitcodeFileContents IF =
+        ExitOnErr(llvm::getBitcodeFileContentsFromPrefix(MB));
+    Stop = MB.getBuffer().empty();
+    if (!disassemble((Nth == 0 && Stop ? -1 : Nth), IF, Context))
+      return 1;
+    ++Nth;
+
+    // After parsing the prefix of MB as a complete bitstream, if MB is still
+    // non-empty, continue parsing succeeding bitstreams.
+  } while (!Stop);
 
   return 0;
 }