CAS.o/FlatV1: Adopt data::TargetInfoList in CompileUnit

Pull addends out of blocks and use data::TargetInfoList to combine with
the TargetIndex. This makes "block" encode exactly what is in
data::BlockData and stops encoding "target"-related data there, fixing a
data layering issue.

Unfortunately this bloats the compile-unit a bit, since the data has
nowhere else to go. Eventually we can/should outline data and references
in logical chunks.

This commit uniques by TargetInfoList, but there's 5-10% bloat in the
size of a CU which makes this hard to justify on its own. A follow-up
will optimize this down; this is not squashed so that it's easy to
review in isolation.

Author: dexonsmith

llvm/include/llvm/CASObjectFormats/FlatV1.h

@@ -9,6 +9,7 @@
 #ifndef LLVM_CASOBJECTFORMATS_FLATV1_H
 #define LLVM_CASOBJECTFORMATS_FLATV1_H
 
+#include "llvm/ADT/StringMap.h"
 #include "llvm/CAS/CASDB.h"
 #include "llvm/CASObjectFormats/Data.h"
 #include "llvm/CASObjectFormats/SchemaBase.h"
@@ -311,7 +312,8 @@ class CompileUnitBuilder {
 
   CompileUnitBuilder(const ObjectFileSchema &Schema, Triple Target,
                      raw_ostream *DebugOS)
-      : CAS(Schema.CAS), Schema(Schema), TT(Target), DebugOS(DebugOS) {}
+      : CAS(Schema.CAS), Schema(Schema), TT(Target), DebugOS(DebugOS),
+        TargetInfoPool(TargetInfoAlloc) {}
 
   Error createAndReferenceName(StringRef S);
   Error createAndReferenceContent(StringRef Content);
@@ -328,6 +330,9 @@ class CompileUnitBuilder {
   // Encode the index into compile unit.
   void encodeIndex(unsigned Index);
 
+  // Encode a target info list and return an offset.
+  Error encodeTargetInfo(ArrayRef<const jitlink::Edge *> Edges);
+
 private:
   friend class CompileUnitRef;
 
@@ -349,6 +354,11 @@ class CompileUnitBuilder {
   SmallVector<const jitlink::Block *, 16> Blocks;
   std::vector<unsigned> BlockIndexStarts;
 
+  // TargetInfoList storage.
+  SmallString<256> FlatTargetInfo;
+  BumpPtrAllocator TargetInfoAlloc;
+  StringMap<uint32_t, BumpPtrAllocator> TargetInfoPool;
+
   // Temp storage.
   SmallString<256> InlineBuffer;
 };
@@ -422,16 +432,19 @@ class LinkGraphBuilder {
   struct BlockInfo {
     jitlink::Block *Block;
     Optional<BlockRef> Ref;
-    Optional<data::BlockData> Data;
+    Optional<data::FixupList> Fixups;
     unsigned BlockIdx = 0;
-    unsigned Remaining = 0;
   };
 
   Expected<BlockInfo &> getBlockInfo(unsigned BlockIdx);
 
   Error addSymbol(unsigned SymbolIdx, jitlink::Symbol *S);
   Expected<jitlink::Symbol *> getSymbol(unsigned SymbolIdx);
 
+  /// Returns a list starting from \p Start, but that iterates past the end of
+  /// the edges.
+  data::TargetInfoList getTargetInfoFrom(unsigned Start);
+
 private:
   friend class CompileUnitRef;
   std::string Name;
@@ -447,6 +460,9 @@ class LinkGraphBuilder {
   unsigned SymbolsSize;
   unsigned BlocksSize;
 
+  // Target info.
+  StringRef FlatTargetInfo;
+
   // Lookup Cache.
   SmallVector<SectionInfo, 16> Sections;
   SmallVector<BlockInfo, 16> Blocks;

llvm/lib/CASObjectFormats/FlatV1.cpp

@@ -293,42 +293,69 @@ Expected<SectionRef> SectionRef::get(Expected<ObjectFormatNodeRef> Ref) {
   return SectionRef(*Specific);
 }
 
-static Error encodeEdge(CompileUnitBuilder &CUB, SmallVectorImpl<char> &Data,
+static Error encodeEdge(CompileUnitBuilder &CUB,
+                        SmallVectorImpl<data::TargetInfo> &List,
                         const jitlink::Edge *E) {
-
   auto SymbolIndex = CUB.getSymbolIndex(E->getTarget());
   if (!SymbolIndex)
     return SymbolIndex.takeError();
+  List.push_back({E->getAddend(), *SymbolIndex});
+  return Error::success();
+}
 
-  unsigned IdxAndHasAddend = *SymbolIndex << 1 | (E->getAddend() != 0);
-  CUB.encodeIndex(IdxAndHasAddend);
+Error CompileUnitBuilder::encodeTargetInfo(
+    ArrayRef<const jitlink::Edge *> Edges) {
+  SmallVector<data::TargetInfo> List;
+  for (const jitlink::Edge *E : Edges)
+    if (Error Err = encodeEdge(*this, List, E))
+      return Err;
+
+  size_t Start = FlatTargetInfo.size();
+  data::TargetInfoList::encode(List, FlatTargetInfo);
+  size_t Length = FlatTargetInfo.size() - Start;
+  auto &MaybeExisting =
+      TargetInfoPool[StringRef(FlatTargetInfo).take_back(Length)];
+  if (MaybeExisting) {
+    // Truncate FlatTargetInfo again since this is already in the pool.
+    FlatTargetInfo.resize(Start);
+    Start = MaybeExisting - 1;
+  } else {
+    MaybeExisting = Start + 1;
+  }
 
-  if (E->getAddend() != 0)
-    encoding::writeVBR8(E->getAddend(), Data);
+  // FIXME: Maybe should use uint64_t, or be a different range from CAS object
+  // indexes.
+  assert(Start < UINT_MAX && "Index out of range");
+  encodeIndex(Start);
   return Error::success();
 }
 
-static Error decodeEdge(LinkGraphBuilder &LGB, StringRef &Data,
-                        jitlink::Block &Parent, unsigned BlockIdx,
-                        const data::Fixup &Fixup) {
-  unsigned SymbolIdx = LGB.nextIdxForBlock(BlockIdx);
+data::TargetInfoList LinkGraphBuilder::getTargetInfoFrom(unsigned Start) {
+  assert(Start < FlatTargetInfo.size() && "Start out of range");
+  return data::TargetInfoList(StringRef(FlatTargetInfo).drop_front(Start));
+}
 
-  bool HasAddend = SymbolIdx & 1U;
-  auto Symbol = LGB.getSymbol(SymbolIdx >> 1);
+static Error decodeEdge(LinkGraphBuilder &LGB, jitlink::Block &Parent,
+                        const data::TargetInfo &TI, const data::Fixup &Fixup) {
+  auto Symbol = LGB.getSymbol(TI.Index);
   if (!Symbol)
     return Symbol.takeError();
 
-  jitlink::Edge::AddendT Addend = 0;
-  if (HasAddend) {
-    if (auto E = encoding::consumeVBR8(Data, Addend))
-      return E;
-  }
-  auto BlockInfo = LGB.getBlockInfo(BlockIdx);
-  if (!BlockInfo)
-    return BlockInfo.takeError();
+  Parent.addEdge(Fixup.Kind, Fixup.Offset, **Symbol, TI.Addend);
+  return Error::success();
+}
 
-  Parent.addEdge(Fixup.Kind, Fixup.Offset, **Symbol, Addend);
+static Error decodeEdges(LinkGraphBuilder &LGB, jitlink::Block &Parent,
+                         unsigned BlockIdx, const data::FixupList &FL) {
+  if (FL.empty())
+    return Error::success();
 
+  unsigned TargetInfoStart = LGB.nextIdxForBlock(BlockIdx);
+  data::TargetInfoList TIL(LGB.getTargetInfoFrom(TargetInfoStart));
+  data::TargetInfoList::iterator TI = TIL.begin();
+  for (const data::Fixup &Fixup : FL)
+    if (Error E = decodeEdge(LGB, Parent, *TI++, Fixup))
+      return E;
   return Error::success();
 }
 
@@ -390,23 +417,19 @@ Expected<BlockRef> BlockRef::create(CompileUnitBuilder &CUB,
     if (auto E = CUB.createAndReferenceContent(BlockData))
       return std::move(E);
   } else {
-    encoding::writeVBR8(BlockData.size(), B->Data);
     B->Data.append(BlockData);
   }
 
-  for (const auto *E : Edges) {
-    // Nest the Edge in block.
-    if (auto Err = encodeEdge(CUB, B->Data, E))
+  if (!Edges.empty())
+    if (auto Err = CUB.encodeTargetInfo(Edges))
       return std::move(Err);
-  }
 
   return get(B->build());
 }
 
 Error BlockRef::materializeBlock(LinkGraphBuilder &LGB,
                                  unsigned BlockIdx) const {
   uint64_t SectionIdx;
-  auto Remaining = getData();
   SectionIdx = LGB.nextIdxForBlock(BlockIdx);
 
   auto SectionInfo = LGB.getSectionInfo(SectionIdx);
@@ -420,21 +443,14 @@ Error BlockRef::materializeBlock(LinkGraphBuilder &LGB,
       return ContentRef.takeError();
     ContentData = ContentRef->getData();
   } else {
-    unsigned ContentSize;
-    if (auto E = encoding::consumeVBR8(Remaining, ContentSize))
-      return E;
-    auto Data = consumeDataOfSize(Remaining, ContentSize);
-    if (!Data)
-      return Data.takeError();
-    ContentData = *Data;
+    ContentData = getData();
   }
 
   auto BlockInfo = LGB.getBlockInfo(BlockIdx);
   if (!BlockInfo)
     return BlockInfo.takeError();
 
-  BlockInfo->Data.emplace(ContentData);
-  auto &Block = *BlockInfo->Data;
+  data::BlockData Block(ContentData);
   auto Address =
       getAlignedAddress(*SectionInfo, Block.getSize(), Block.getAlignment(),
                         Block.getAlignmentOffset());
@@ -447,7 +463,7 @@ Error BlockRef::materializeBlock(LinkGraphBuilder &LGB,
                       Block.getAlignment(), Block.getAlignmentOffset());
 
   BlockInfo->Block = &B;
-  BlockInfo->Remaining = Remaining.size();
+  BlockInfo->Fixups = Block.getFixups();
   return Error::success();
 }
 
@@ -458,17 +474,10 @@ Error BlockRef::materializeEdges(LinkGraphBuilder &LGB,
     return BlockInfo.takeError();
 
   // Nothing remains, no edges.
-  if (!BlockInfo->Remaining)
+  if (BlockInfo->Fixups->empty())
     return Error::success();
 
-  auto Remaining = getData().take_back(BlockInfo->Remaining);
-  for (const data::Fixup &Fixup : BlockInfo->Data->getFixups()) {
-    if (Error Err =
-            decodeEdge(LGB, Remaining, *BlockInfo->Block, BlockIdx, Fixup))
-      return Err;
-  }
-
-  return Error::success();
+  return decodeEdges(LGB, *BlockInfo->Block, BlockIdx, *BlockInfo->Fixups);
 }
 
 Expected<BlockRef> BlockRef::get(Expected<ObjectFormatNodeRef> Ref) {
@@ -844,6 +853,10 @@ Expected<CompileUnitRef> CompileUnitRef::create(const ObjectFileSchema &Schema,
   for (auto Idx : Builder.BlockIndexStarts)
     encoding::writeVBR8(Idx, B->Data);
 
+  /// Write the flat TargetInfoList data.
+  encoding::writeVBR8(Builder.FlatTargetInfo.size(), B->Data);
+  B->Data.append(Builder.FlatTargetInfo);
+
   // Inlined symbols if set.
   if (InlineSymbols)
     B->Data.append(Builder.InlineBuffer);
@@ -914,6 +927,14 @@ Error CompileUnitRef::materialize(LinkGraphBuilder &LGB) const {
     LGB.Blocks[I].BlockIdx = Idx;
   }
 
+  /// Write the flat TargetInfoList data.
+  size_t FlatTargetInfoSize;
+  if (Error Err = encoding::consumeVBR8(Remaining, FlatTargetInfoSize))
+    return Err;
+  if (Error Err = consumeDataOfSize(Remaining, FlatTargetInfoSize)
+                      .moveInto(LGB.FlatTargetInfo))
+    return Err;
+
   if (InlineSymbols)
     LGB.InlineBuffer = Remaining;
   else if (!Remaining.empty())