1 //===- SampleProfWriter.cpp - Write LLVM sample profile data --------------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 //
9 // This file implements the class that writes LLVM sample profiles. It
10 // supports two file formats: text and binary. The textual representation
11 // is useful for debugging and testing purposes. The binary representation
12 // is more compact, resulting in smaller file sizes. However, they can
13 // both be used interchangeably.
14 //
15 // See lib/ProfileData/SampleProfReader.cpp for documentation on each of the
16 // supported formats.
17 //
18 //===----------------------------------------------------------------------===//
19 
20 #include "llvm/ProfileData/SampleProfWriter.h"
21 #include "llvm/ADT/StringRef.h"
22 #include "llvm/ProfileData/ProfileCommon.h"
23 #include "llvm/ProfileData/SampleProf.h"
24 #include "llvm/Support/Compression.h"
25 #include "llvm/Support/Endian.h"
26 #include "llvm/Support/EndianStream.h"
27 #include "llvm/Support/ErrorOr.h"
28 #include "llvm/Support/FileSystem.h"
29 #include "llvm/Support/LEB128.h"
30 #include "llvm/Support/MD5.h"
31 #include "llvm/Support/raw_ostream.h"
32 #include <algorithm>
33 #include <cstdint>
34 #include <memory>
35 #include <set>
36 #include <system_error>
37 #include <utility>
38 #include <vector>
39 
40 using namespace llvm;
41 using namespace sampleprof;
42 
43 std::error_code
44 SampleProfileWriter::writeFuncProfiles(const SampleProfileMap &ProfileMap) {
45   std::vector<NameFunctionSamples> V;
46   sortFuncProfiles(ProfileMap, V);
47   for (const auto &I : V) {
48     if (std::error_code EC = writeSample(*I.second))
49       return EC;
50   }
51   return sampleprof_error::success;
52 }
53 
54 std::error_code SampleProfileWriter::write(const SampleProfileMap &ProfileMap) {
55   if (std::error_code EC = writeHeader(ProfileMap))
56     return EC;
57 
58   if (std::error_code EC = writeFuncProfiles(ProfileMap))
59     return EC;
60 
61   return sampleprof_error::success;
62 }
63 
64 /// Return the current position and prepare to use it as the start
65 /// position of a section given the section type \p Type and its position
66 /// \p LayoutIdx in SectionHdrLayout.
67 uint64_t
68 SampleProfileWriterExtBinaryBase::markSectionStart(SecType Type,
69                                                    uint32_t LayoutIdx) {
70   uint64_t SectionStart = OutputStream->tell();
71   assert(LayoutIdx < SectionHdrLayout.size() && "LayoutIdx out of range");
72   const auto &Entry = SectionHdrLayout[LayoutIdx];
73   assert(Entry.Type == Type && "Unexpected section type");
74   // Use LocalBuf as a temporary output for writting data.
75   if (hasSecFlag(Entry, SecCommonFlags::SecFlagCompress))
76     LocalBufStream.swap(OutputStream);
77   return SectionStart;
78 }
79 
80 std::error_code SampleProfileWriterExtBinaryBase::compressAndOutput() {
81   if (!llvm::compression::zlib::isAvailable())
82     return sampleprof_error::zlib_unavailable;
83   std::string &UncompressedStrings =
84       static_cast<raw_string_ostream *>(LocalBufStream.get())->str();
85   if (UncompressedStrings.size() == 0)
86     return sampleprof_error::success;
87   auto &OS = *OutputStream;
88   SmallVector<uint8_t, 128> CompressedStrings;
89   compression::zlib::compress(arrayRefFromStringRef(UncompressedStrings),
90                               CompressedStrings,
91                               compression::zlib::BestSizeCompression);
92   encodeULEB128(UncompressedStrings.size(), OS);
93   encodeULEB128(CompressedStrings.size(), OS);
94   OS << toStringRef(CompressedStrings);
95   UncompressedStrings.clear();
96   return sampleprof_error::success;
97 }
98 
99 /// Add a new section into section header table given the section type
100 /// \p Type, its position \p LayoutIdx in SectionHdrLayout and the
101 /// location \p SectionStart where the section should be written to.
102 std::error_code SampleProfileWriterExtBinaryBase::addNewSection(
103     SecType Type, uint32_t LayoutIdx, uint64_t SectionStart) {
104   assert(LayoutIdx < SectionHdrLayout.size() && "LayoutIdx out of range");
105   const auto &Entry = SectionHdrLayout[LayoutIdx];
106   assert(Entry.Type == Type && "Unexpected section type");
107   if (hasSecFlag(Entry, SecCommonFlags::SecFlagCompress)) {
108     LocalBufStream.swap(OutputStream);
109     if (std::error_code EC = compressAndOutput())
110       return EC;
111   }
112   SecHdrTable.push_back({Type, Entry.Flags, SectionStart - FileStart,
113                          OutputStream->tell() - SectionStart, LayoutIdx});
114   return sampleprof_error::success;
115 }
116 
117 std::error_code
118 SampleProfileWriterExtBinaryBase::write(const SampleProfileMap &ProfileMap) {
119   if (std::error_code EC = writeHeader(ProfileMap))
120     return EC;
121 
122   std::string LocalBuf;
123   LocalBufStream = std::make_unique<raw_string_ostream>(LocalBuf);
124   if (std::error_code EC = writeSections(ProfileMap))
125     return EC;
126 
127   if (std::error_code EC = writeSecHdrTable())
128     return EC;
129 
130   return sampleprof_error::success;
131 }
132 
133 std::error_code SampleProfileWriterExtBinaryBase::writeContextIdx(
134     const SampleContext &Context) {
135   if (Context.hasContext())
136     return writeCSNameIdx(Context);
137   else
138     return SampleProfileWriterBinary::writeNameIdx(Context.getName());
139 }
140 
141 std::error_code
142 SampleProfileWriterExtBinaryBase::writeCSNameIdx(const SampleContext &Context) {
143   const auto &Ret = CSNameTable.find(Context);
144   if (Ret == CSNameTable.end())
145     return sampleprof_error::truncated_name_table;
146   encodeULEB128(Ret->second, *OutputStream);
147   return sampleprof_error::success;
148 }
149 
150 std::error_code
151 SampleProfileWriterExtBinaryBase::writeSample(const FunctionSamples &S) {
152   uint64_t Offset = OutputStream->tell();
153   auto &Context = S.getContext();
154   FuncOffsetTable[Context] = Offset - SecLBRProfileStart;
155   encodeULEB128(S.getHeadSamples(), *OutputStream);
156   return writeBody(S);
157 }
158 
159 std::error_code SampleProfileWriterExtBinaryBase::writeFuncOffsetTable() {
160   auto &OS = *OutputStream;
161 
162   // Write out the table size.
163   encodeULEB128(FuncOffsetTable.size(), OS);
164 
165   // Write out FuncOffsetTable.
166   auto WriteItem = [&](const SampleContext &Context, uint64_t Offset) {
167     if (std::error_code EC = writeContextIdx(Context))
168       return EC;
169     encodeULEB128(Offset, OS);
170     return (std::error_code)sampleprof_error::success;
171   };
172 
173   if (FunctionSamples::ProfileIsCS) {
174     // Sort the contexts before writing them out. This is to help fast load all
175     // context profiles for a function as well as their callee contexts which
176     // can help profile-guided importing for ThinLTO.
177     std::map<SampleContext, uint64_t> OrderedFuncOffsetTable(
178         FuncOffsetTable.begin(), FuncOffsetTable.end());
179     for (const auto &Entry : OrderedFuncOffsetTable) {
180       if (std::error_code EC = WriteItem(Entry.first, Entry.second))
181         return EC;
182     }
183     addSectionFlag(SecFuncOffsetTable, SecFuncOffsetFlags::SecFlagOrdered);
184   } else {
185     for (const auto &Entry : FuncOffsetTable) {
186       if (std::error_code EC = WriteItem(Entry.first, Entry.second))
187         return EC;
188     }
189   }
190 
191   FuncOffsetTable.clear();
192   return sampleprof_error::success;
193 }
194 
195 std::error_code SampleProfileWriterExtBinaryBase::writeFuncMetadata(
196     const FunctionSamples &FunctionProfile) {
197   auto &OS = *OutputStream;
198   if (std::error_code EC = writeContextIdx(FunctionProfile.getContext()))
199     return EC;
200 
201   if (FunctionSamples::ProfileIsProbeBased)
202     encodeULEB128(FunctionProfile.getFunctionHash(), OS);
203   if (FunctionSamples::ProfileIsCS || FunctionSamples::ProfileIsPreInlined) {
204     encodeULEB128(FunctionProfile.getContext().getAllAttributes(), OS);
205   }
206 
207   if (!FunctionSamples::ProfileIsCS) {
208     // Recursively emit attributes for all callee samples.
209     uint64_t NumCallsites = 0;
210     for (const auto &J : FunctionProfile.getCallsiteSamples())
211       NumCallsites += J.second.size();
212     encodeULEB128(NumCallsites, OS);
213     for (const auto &J : FunctionProfile.getCallsiteSamples()) {
214       for (const auto &FS : J.second) {
215         LineLocation Loc = J.first;
216         encodeULEB128(Loc.LineOffset, OS);
217         encodeULEB128(Loc.Discriminator, OS);
218         if (std::error_code EC = writeFuncMetadata(FS.second))
219           return EC;
220       }
221     }
222   }
223 
224   return sampleprof_error::success;
225 }
226 
227 std::error_code SampleProfileWriterExtBinaryBase::writeFuncMetadata(
228     const SampleProfileMap &Profiles) {
229   if (!FunctionSamples::ProfileIsProbeBased && !FunctionSamples::ProfileIsCS &&
230       !FunctionSamples::ProfileIsPreInlined)
231     return sampleprof_error::success;
232   for (const auto &Entry : Profiles) {
233     if (std::error_code EC = writeFuncMetadata(Entry.second))
234       return EC;
235   }
236   return sampleprof_error::success;
237 }
238 
239 std::error_code SampleProfileWriterExtBinaryBase::writeNameTable() {
240   if (!UseMD5)
241     return SampleProfileWriterBinary::writeNameTable();
242 
243   auto &OS = *OutputStream;
244   std::set<StringRef> V;
245   stablizeNameTable(NameTable, V);
246 
247   // Write out the MD5 name table. We wrote unencoded MD5 so reader can
248   // retrieve the name using the name index without having to read the
249   // whole name table.
250   encodeULEB128(NameTable.size(), OS);
251   support::endian::Writer Writer(OS, support::little);
252   for (auto N : V)
253     Writer.write(MD5Hash(N));
254   return sampleprof_error::success;
255 }
256 
257 std::error_code SampleProfileWriterExtBinaryBase::writeNameTableSection(
258     const SampleProfileMap &ProfileMap) {
259   for (const auto &I : ProfileMap) {
260     assert(I.first == I.second.getContext() && "Inconsistent profile map");
261     addContext(I.second.getContext());
262     addNames(I.second);
263   }
264 
265   // If NameTable contains ".__uniq." suffix, set SecFlagUniqSuffix flag
266   // so compiler won't strip the suffix during profile matching after
267   // seeing the flag in the profile.
268   for (const auto &I : NameTable) {
269     if (I.first.contains(FunctionSamples::UniqSuffix)) {
270       addSectionFlag(SecNameTable, SecNameTableFlags::SecFlagUniqSuffix);
271       break;
272     }
273   }
274 
275   if (auto EC = writeNameTable())
276     return EC;
277   return sampleprof_error::success;
278 }
279 
280 std::error_code SampleProfileWriterExtBinaryBase::writeCSNameTableSection() {
281   // Sort the names to make CSNameTable deterministic.
282   std::set<SampleContext> OrderedContexts;
283   for (const auto &I : CSNameTable)
284     OrderedContexts.insert(I.first);
285   assert(OrderedContexts.size() == CSNameTable.size() &&
286          "Unmatched ordered and unordered contexts");
287   uint64_t I = 0;
288   for (auto &Context : OrderedContexts)
289     CSNameTable[Context] = I++;
290 
291   auto &OS = *OutputStream;
292   encodeULEB128(OrderedContexts.size(), OS);
293   support::endian::Writer Writer(OS, support::little);
294   for (auto Context : OrderedContexts) {
295     auto Frames = Context.getContextFrames();
296     encodeULEB128(Frames.size(), OS);
297     for (auto &Callsite : Frames) {
298       if (std::error_code EC = writeNameIdx(Callsite.FuncName))
299         return EC;
300       encodeULEB128(Callsite.Location.LineOffset, OS);
301       encodeULEB128(Callsite.Location.Discriminator, OS);
302     }
303   }
304 
305   return sampleprof_error::success;
306 }
307 
308 std::error_code
309 SampleProfileWriterExtBinaryBase::writeProfileSymbolListSection() {
310   if (ProfSymList && ProfSymList->size() > 0)
311     if (std::error_code EC = ProfSymList->write(*OutputStream))
312       return EC;
313 
314   return sampleprof_error::success;
315 }
316 
317 std::error_code SampleProfileWriterExtBinaryBase::writeOneSection(
318     SecType Type, uint32_t LayoutIdx, const SampleProfileMap &ProfileMap) {
319   // The setting of SecFlagCompress should happen before markSectionStart.
320   if (Type == SecProfileSymbolList && ProfSymList && ProfSymList->toCompress())
321     setToCompressSection(SecProfileSymbolList);
322   if (Type == SecFuncMetadata && FunctionSamples::ProfileIsProbeBased)
323     addSectionFlag(SecFuncMetadata, SecFuncMetadataFlags::SecFlagIsProbeBased);
324   if (Type == SecFuncMetadata &&
325       (FunctionSamples::ProfileIsCS || FunctionSamples::ProfileIsPreInlined))
326     addSectionFlag(SecFuncMetadata, SecFuncMetadataFlags::SecFlagHasAttribute);
327   if (Type == SecProfSummary && FunctionSamples::ProfileIsCS)
328     addSectionFlag(SecProfSummary, SecProfSummaryFlags::SecFlagFullContext);
329   if (Type == SecProfSummary && FunctionSamples::ProfileIsPreInlined)
330     addSectionFlag(SecProfSummary, SecProfSummaryFlags::SecFlagIsPreInlined);
331   if (Type == SecProfSummary && FunctionSamples::ProfileIsFS)
332     addSectionFlag(SecProfSummary, SecProfSummaryFlags::SecFlagFSDiscriminator);
333 
334   uint64_t SectionStart = markSectionStart(Type, LayoutIdx);
335   switch (Type) {
336   case SecProfSummary:
337     computeSummary(ProfileMap);
338     if (auto EC = writeSummary())
339       return EC;
340     break;
341   case SecNameTable:
342     if (auto EC = writeNameTableSection(ProfileMap))
343       return EC;
344     break;
345   case SecCSNameTable:
346     if (auto EC = writeCSNameTableSection())
347       return EC;
348     break;
349   case SecLBRProfile:
350     SecLBRProfileStart = OutputStream->tell();
351     if (std::error_code EC = writeFuncProfiles(ProfileMap))
352       return EC;
353     break;
354   case SecFuncOffsetTable:
355     if (auto EC = writeFuncOffsetTable())
356       return EC;
357     break;
358   case SecFuncMetadata:
359     if (std::error_code EC = writeFuncMetadata(ProfileMap))
360       return EC;
361     break;
362   case SecProfileSymbolList:
363     if (auto EC = writeProfileSymbolListSection())
364       return EC;
365     break;
366   default:
367     if (auto EC = writeCustomSection(Type))
368       return EC;
369     break;
370   }
371   if (std::error_code EC = addNewSection(Type, LayoutIdx, SectionStart))
372     return EC;
373   return sampleprof_error::success;
374 }
375 
376 std::error_code SampleProfileWriterExtBinary::writeDefaultLayout(
377     const SampleProfileMap &ProfileMap) {
378   // The const indices passed to writeOneSection below are specifying the
379   // positions of the sections in SectionHdrLayout. Look at
380   // initSectionHdrLayout to find out where each section is located in
381   // SectionHdrLayout.
382   if (auto EC = writeOneSection(SecProfSummary, 0, ProfileMap))
383     return EC;
384   if (auto EC = writeOneSection(SecNameTable, 1, ProfileMap))
385     return EC;
386   if (auto EC = writeOneSection(SecCSNameTable, 2, ProfileMap))
387     return EC;
388   if (auto EC = writeOneSection(SecLBRProfile, 4, ProfileMap))
389     return EC;
390   if (auto EC = writeOneSection(SecProfileSymbolList, 5, ProfileMap))
391     return EC;
392   if (auto EC = writeOneSection(SecFuncOffsetTable, 3, ProfileMap))
393     return EC;
394   if (auto EC = writeOneSection(SecFuncMetadata, 6, ProfileMap))
395     return EC;
396   return sampleprof_error::success;
397 }
398 
399 static void splitProfileMapToTwo(const SampleProfileMap &ProfileMap,
400                                  SampleProfileMap &ContextProfileMap,
401                                  SampleProfileMap &NoContextProfileMap) {
402   for (const auto &I : ProfileMap) {
403     if (I.second.getCallsiteSamples().size())
404       ContextProfileMap.insert({I.first, I.second});
405     else
406       NoContextProfileMap.insert({I.first, I.second});
407   }
408 }
409 
410 std::error_code SampleProfileWriterExtBinary::writeCtxSplitLayout(
411     const SampleProfileMap &ProfileMap) {
412   SampleProfileMap ContextProfileMap, NoContextProfileMap;
413   splitProfileMapToTwo(ProfileMap, ContextProfileMap, NoContextProfileMap);
414 
415   if (auto EC = writeOneSection(SecProfSummary, 0, ProfileMap))
416     return EC;
417   if (auto EC = writeOneSection(SecNameTable, 1, ProfileMap))
418     return EC;
419   if (auto EC = writeOneSection(SecLBRProfile, 3, ContextProfileMap))
420     return EC;
421   if (auto EC = writeOneSection(SecFuncOffsetTable, 2, ContextProfileMap))
422     return EC;
423   // Mark the section to have no context. Note section flag needs to be set
424   // before writing the section.
425   addSectionFlag(5, SecCommonFlags::SecFlagFlat);
426   if (auto EC = writeOneSection(SecLBRProfile, 5, NoContextProfileMap))
427     return EC;
428   // Mark the section to have no context. Note section flag needs to be set
429   // before writing the section.
430   addSectionFlag(4, SecCommonFlags::SecFlagFlat);
431   if (auto EC = writeOneSection(SecFuncOffsetTable, 4, NoContextProfileMap))
432     return EC;
433   if (auto EC = writeOneSection(SecProfileSymbolList, 6, ProfileMap))
434     return EC;
435   if (auto EC = writeOneSection(SecFuncMetadata, 7, ProfileMap))
436     return EC;
437 
438   return sampleprof_error::success;
439 }
440 
441 std::error_code SampleProfileWriterExtBinary::writeSections(
442     const SampleProfileMap &ProfileMap) {
443   std::error_code EC;
444   if (SecLayout == DefaultLayout)
445     EC = writeDefaultLayout(ProfileMap);
446   else if (SecLayout == CtxSplitLayout)
447     EC = writeCtxSplitLayout(ProfileMap);
448   else
449     llvm_unreachable("Unsupported layout");
450   return EC;
451 }
452 
453 std::error_code
454 SampleProfileWriterCompactBinary::write(const SampleProfileMap &ProfileMap) {
455   if (std::error_code EC = SampleProfileWriter::write(ProfileMap))
456     return EC;
457   if (std::error_code EC = writeFuncOffsetTable())
458     return EC;
459   return sampleprof_error::success;
460 }
461 
462 /// Write samples to a text file.
463 ///
464 /// Note: it may be tempting to implement this in terms of
465 /// FunctionSamples::print().  Please don't.  The dump functionality is intended
466 /// for debugging and has no specified form.
467 ///
468 /// The format used here is more structured and deliberate because
469 /// it needs to be parsed by the SampleProfileReaderText class.
470 std::error_code SampleProfileWriterText::writeSample(const FunctionSamples &S) {
471   auto &OS = *OutputStream;
472   if (FunctionSamples::ProfileIsCS)
473     OS << "[" << S.getContext().toString() << "]:" << S.getTotalSamples();
474   else
475     OS << S.getName() << ":" << S.getTotalSamples();
476 
477   if (Indent == 0)
478     OS << ":" << S.getHeadSamples();
479   OS << "\n";
480 
481   SampleSorter<LineLocation, SampleRecord> SortedSamples(S.getBodySamples());
482   for (const auto &I : SortedSamples.get()) {
483     LineLocation Loc = I->first;
484     const SampleRecord &Sample = I->second;
485     OS.indent(Indent + 1);
486     if (Loc.Discriminator == 0)
487       OS << Loc.LineOffset << ": ";
488     else
489       OS << Loc.LineOffset << "." << Loc.Discriminator << ": ";
490 
491     OS << Sample.getSamples();
492 
493     for (const auto &J : Sample.getSortedCallTargets())
494       OS << " " << J.first << ":" << J.second;
495     OS << "\n";
496   }
497 
498   SampleSorter<LineLocation, FunctionSamplesMap> SortedCallsiteSamples(
499       S.getCallsiteSamples());
500   Indent += 1;
501   for (const auto &I : SortedCallsiteSamples.get())
502     for (const auto &FS : I->second) {
503       LineLocation Loc = I->first;
504       const FunctionSamples &CalleeSamples = FS.second;
505       OS.indent(Indent);
506       if (Loc.Discriminator == 0)
507         OS << Loc.LineOffset << ": ";
508       else
509         OS << Loc.LineOffset << "." << Loc.Discriminator << ": ";
510       if (std::error_code EC = writeSample(CalleeSamples))
511         return EC;
512     }
513   Indent -= 1;
514 
515   if (FunctionSamples::ProfileIsProbeBased) {
516     OS.indent(Indent + 1);
517     OS << "!CFGChecksum: " << S.getFunctionHash() << "\n";
518   }
519 
520   if (S.getContext().getAllAttributes()) {
521     OS.indent(Indent + 1);
522     OS << "!Attributes: " << S.getContext().getAllAttributes() << "\n";
523   }
524 
525   return sampleprof_error::success;
526 }
527 
528 std::error_code
529 SampleProfileWriterBinary::writeContextIdx(const SampleContext &Context) {
530   assert(!Context.hasContext() && "cs profile is not supported");
531   return writeNameIdx(Context.getName());
532 }
533 
534 std::error_code SampleProfileWriterBinary::writeNameIdx(StringRef FName) {
535   auto &NTable = getNameTable();
536   const auto &Ret = NTable.find(FName);
537   if (Ret == NTable.end())
538     return sampleprof_error::truncated_name_table;
539   encodeULEB128(Ret->second, *OutputStream);
540   return sampleprof_error::success;
541 }
542 
543 void SampleProfileWriterBinary::addName(StringRef FName) {
544   auto &NTable = getNameTable();
545   NTable.insert(std::make_pair(FName, 0));
546 }
547 
548 void SampleProfileWriterBinary::addContext(const SampleContext &Context) {
549   addName(Context.getName());
550 }
551 
552 void SampleProfileWriterBinary::addNames(const FunctionSamples &S) {
553   // Add all the names in indirect call targets.
554   for (const auto &I : S.getBodySamples()) {
555     const SampleRecord &Sample = I.second;
556     for (const auto &J : Sample.getCallTargets())
557       addName(J.first());
558   }
559 
560   // Recursively add all the names for inlined callsites.
561   for (const auto &J : S.getCallsiteSamples())
562     for (const auto &FS : J.second) {
563       const FunctionSamples &CalleeSamples = FS.second;
564       addName(CalleeSamples.getName());
565       addNames(CalleeSamples);
566     }
567 }
568 
569 void SampleProfileWriterExtBinaryBase::addContext(
570     const SampleContext &Context) {
571   if (Context.hasContext()) {
572     for (auto &Callsite : Context.getContextFrames())
573       SampleProfileWriterBinary::addName(Callsite.FuncName);
574     CSNameTable.insert(std::make_pair(Context, 0));
575   } else {
576     SampleProfileWriterBinary::addName(Context.getName());
577   }
578 }
579 
580 void SampleProfileWriterBinary::stablizeNameTable(
581     MapVector<StringRef, uint32_t> &NameTable, std::set<StringRef> &V) {
582   // Sort the names to make NameTable deterministic.
583   for (const auto &I : NameTable)
584     V.insert(I.first);
585   int i = 0;
586   for (const StringRef &N : V)
587     NameTable[N] = i++;
588 }
589 
590 std::error_code SampleProfileWriterBinary::writeNameTable() {
591   auto &OS = *OutputStream;
592   std::set<StringRef> V;
593   stablizeNameTable(NameTable, V);
594 
595   // Write out the name table.
596   encodeULEB128(NameTable.size(), OS);
597   for (auto N : V) {
598     OS << N;
599     encodeULEB128(0, OS);
600   }
601   return sampleprof_error::success;
602 }
603 
604 std::error_code SampleProfileWriterCompactBinary::writeFuncOffsetTable() {
605   auto &OS = *OutputStream;
606 
607   // Fill the slot remembered by TableOffset with the offset of FuncOffsetTable.
608   auto &OFS = static_cast<raw_fd_ostream &>(OS);
609   uint64_t FuncOffsetTableStart = OS.tell();
610   if (OFS.seek(TableOffset) == (uint64_t)-1)
611     return sampleprof_error::ostream_seek_unsupported;
612   support::endian::Writer Writer(*OutputStream, support::little);
613   Writer.write(FuncOffsetTableStart);
614   if (OFS.seek(FuncOffsetTableStart) == (uint64_t)-1)
615     return sampleprof_error::ostream_seek_unsupported;
616 
617   // Write out the table size.
618   encodeULEB128(FuncOffsetTable.size(), OS);
619 
620   // Write out FuncOffsetTable.
621   for (auto Entry : FuncOffsetTable) {
622     if (std::error_code EC = writeNameIdx(Entry.first))
623       return EC;
624     encodeULEB128(Entry.second, OS);
625   }
626   return sampleprof_error::success;
627 }
628 
629 std::error_code SampleProfileWriterCompactBinary::writeNameTable() {
630   auto &OS = *OutputStream;
631   std::set<StringRef> V;
632   stablizeNameTable(NameTable, V);
633 
634   // Write out the name table.
635   encodeULEB128(NameTable.size(), OS);
636   for (auto N : V) {
637     encodeULEB128(MD5Hash(N), OS);
638   }
639   return sampleprof_error::success;
640 }
641 
642 std::error_code
643 SampleProfileWriterBinary::writeMagicIdent(SampleProfileFormat Format) {
644   auto &OS = *OutputStream;
645   // Write file magic identifier.
646   encodeULEB128(SPMagic(Format), OS);
647   encodeULEB128(SPVersion(), OS);
648   return sampleprof_error::success;
649 }
650 
651 std::error_code
652 SampleProfileWriterBinary::writeHeader(const SampleProfileMap &ProfileMap) {
653   writeMagicIdent(Format);
654 
655   computeSummary(ProfileMap);
656   if (auto EC = writeSummary())
657     return EC;
658 
659   // Generate the name table for all the functions referenced in the profile.
660   for (const auto &I : ProfileMap) {
661     assert(I.first == I.second.getContext() && "Inconsistent profile map");
662     addContext(I.first);
663     addNames(I.second);
664   }
665 
666   writeNameTable();
667   return sampleprof_error::success;
668 }
669 
670 void SampleProfileWriterExtBinaryBase::setToCompressAllSections() {
671   for (auto &Entry : SectionHdrLayout)
672     addSecFlag(Entry, SecCommonFlags::SecFlagCompress);
673 }
674 
675 void SampleProfileWriterExtBinaryBase::setToCompressSection(SecType Type) {
676   addSectionFlag(Type, SecCommonFlags::SecFlagCompress);
677 }
678 
679 void SampleProfileWriterExtBinaryBase::allocSecHdrTable() {
680   support::endian::Writer Writer(*OutputStream, support::little);
681 
682   Writer.write(static_cast<uint64_t>(SectionHdrLayout.size()));
683   SecHdrTableOffset = OutputStream->tell();
684   for (uint32_t i = 0; i < SectionHdrLayout.size(); i++) {
685     Writer.write(static_cast<uint64_t>(-1));
686     Writer.write(static_cast<uint64_t>(-1));
687     Writer.write(static_cast<uint64_t>(-1));
688     Writer.write(static_cast<uint64_t>(-1));
689   }
690 }
691 
692 std::error_code SampleProfileWriterExtBinaryBase::writeSecHdrTable() {
693   auto &OFS = static_cast<raw_fd_ostream &>(*OutputStream);
694   uint64_t Saved = OutputStream->tell();
695 
696   // Set OutputStream to the location saved in SecHdrTableOffset.
697   if (OFS.seek(SecHdrTableOffset) == (uint64_t)-1)
698     return sampleprof_error::ostream_seek_unsupported;
699   support::endian::Writer Writer(*OutputStream, support::little);
700 
701   assert(SecHdrTable.size() == SectionHdrLayout.size() &&
702          "SecHdrTable entries doesn't match SectionHdrLayout");
703   SmallVector<uint32_t, 16> IndexMap(SecHdrTable.size(), -1);
704   for (uint32_t TableIdx = 0; TableIdx < SecHdrTable.size(); TableIdx++) {
705     IndexMap[SecHdrTable[TableIdx].LayoutIndex] = TableIdx;
706   }
707 
708   // Write the section header table in the order specified in
709   // SectionHdrLayout. SectionHdrLayout specifies the sections
710   // order in which profile reader expect to read, so the section
711   // header table should be written in the order in SectionHdrLayout.
712   // Note that the section order in SecHdrTable may be different
713   // from the order in SectionHdrLayout, for example, SecFuncOffsetTable
714   // needs to be computed after SecLBRProfile (the order in SecHdrTable),
715   // but it needs to be read before SecLBRProfile (the order in
716   // SectionHdrLayout). So we use IndexMap above to switch the order.
717   for (uint32_t LayoutIdx = 0; LayoutIdx < SectionHdrLayout.size();
718        LayoutIdx++) {
719     assert(IndexMap[LayoutIdx] < SecHdrTable.size() &&
720            "Incorrect LayoutIdx in SecHdrTable");
721     auto Entry = SecHdrTable[IndexMap[LayoutIdx]];
722     Writer.write(static_cast<uint64_t>(Entry.Type));
723     Writer.write(static_cast<uint64_t>(Entry.Flags));
724     Writer.write(static_cast<uint64_t>(Entry.Offset));
725     Writer.write(static_cast<uint64_t>(Entry.Size));
726   }
727 
728   // Reset OutputStream.
729   if (OFS.seek(Saved) == (uint64_t)-1)
730     return sampleprof_error::ostream_seek_unsupported;
731 
732   return sampleprof_error::success;
733 }
734 
735 std::error_code SampleProfileWriterExtBinaryBase::writeHeader(
736     const SampleProfileMap &ProfileMap) {
737   auto &OS = *OutputStream;
738   FileStart = OS.tell();
739   writeMagicIdent(Format);
740 
741   allocSecHdrTable();
742   return sampleprof_error::success;
743 }
744 
745 std::error_code SampleProfileWriterCompactBinary::writeHeader(
746     const SampleProfileMap &ProfileMap) {
747   support::endian::Writer Writer(*OutputStream, support::little);
748   if (auto EC = SampleProfileWriterBinary::writeHeader(ProfileMap))
749     return EC;
750 
751   // Reserve a slot for the offset of function offset table. The slot will
752   // be populated with the offset of FuncOffsetTable later.
753   TableOffset = OutputStream->tell();
754   Writer.write(static_cast<uint64_t>(-2));
755   return sampleprof_error::success;
756 }
757 
758 std::error_code SampleProfileWriterBinary::writeSummary() {
759   auto &OS = *OutputStream;
760   encodeULEB128(Summary->getTotalCount(), OS);
761   encodeULEB128(Summary->getMaxCount(), OS);
762   encodeULEB128(Summary->getMaxFunctionCount(), OS);
763   encodeULEB128(Summary->getNumCounts(), OS);
764   encodeULEB128(Summary->getNumFunctions(), OS);
765   const std::vector<ProfileSummaryEntry> &Entries =
766       Summary->getDetailedSummary();
767   encodeULEB128(Entries.size(), OS);
768   for (auto Entry : Entries) {
769     encodeULEB128(Entry.Cutoff, OS);
770     encodeULEB128(Entry.MinCount, OS);
771     encodeULEB128(Entry.NumCounts, OS);
772   }
773   return sampleprof_error::success;
774 }
775 std::error_code SampleProfileWriterBinary::writeBody(const FunctionSamples &S) {
776   auto &OS = *OutputStream;
777   if (std::error_code EC = writeContextIdx(S.getContext()))
778     return EC;
779 
780   encodeULEB128(S.getTotalSamples(), OS);
781 
782   // Emit all the body samples.
783   encodeULEB128(S.getBodySamples().size(), OS);
784   for (const auto &I : S.getBodySamples()) {
785     LineLocation Loc = I.first;
786     const SampleRecord &Sample = I.second;
787     encodeULEB128(Loc.LineOffset, OS);
788     encodeULEB128(Loc.Discriminator, OS);
789     encodeULEB128(Sample.getSamples(), OS);
790     encodeULEB128(Sample.getCallTargets().size(), OS);
791     for (const auto &J : Sample.getSortedCallTargets()) {
792       StringRef Callee = J.first;
793       uint64_t CalleeSamples = J.second;
794       if (std::error_code EC = writeNameIdx(Callee))
795         return EC;
796       encodeULEB128(CalleeSamples, OS);
797     }
798   }
799 
800   // Recursively emit all the callsite samples.
801   uint64_t NumCallsites = 0;
802   for (const auto &J : S.getCallsiteSamples())
803     NumCallsites += J.second.size();
804   encodeULEB128(NumCallsites, OS);
805   for (const auto &J : S.getCallsiteSamples())
806     for (const auto &FS : J.second) {
807       LineLocation Loc = J.first;
808       const FunctionSamples &CalleeSamples = FS.second;
809       encodeULEB128(Loc.LineOffset, OS);
810       encodeULEB128(Loc.Discriminator, OS);
811       if (std::error_code EC = writeBody(CalleeSamples))
812         return EC;
813     }
814 
815   return sampleprof_error::success;
816 }
817 
818 /// Write samples of a top-level function to a binary file.
819 ///
820 /// \returns true if the samples were written successfully, false otherwise.
821 std::error_code
822 SampleProfileWriterBinary::writeSample(const FunctionSamples &S) {
823   encodeULEB128(S.getHeadSamples(), *OutputStream);
824   return writeBody(S);
825 }
826 
827 std::error_code
828 SampleProfileWriterCompactBinary::writeSample(const FunctionSamples &S) {
829   uint64_t Offset = OutputStream->tell();
830   StringRef Name = S.getName();
831   FuncOffsetTable[Name] = Offset;
832   encodeULEB128(S.getHeadSamples(), *OutputStream);
833   return writeBody(S);
834 }
835 
836 /// Create a sample profile file writer based on the specified format.
837 ///
838 /// \param Filename The file to create.
839 ///
840 /// \param Format Encoding format for the profile file.
841 ///
842 /// \returns an error code indicating the status of the created writer.
843 ErrorOr<std::unique_ptr<SampleProfileWriter>>
844 SampleProfileWriter::create(StringRef Filename, SampleProfileFormat Format) {
845   std::error_code EC;
846   std::unique_ptr<raw_ostream> OS;
847   if (Format == SPF_Binary || Format == SPF_Ext_Binary ||
848       Format == SPF_Compact_Binary)
849     OS.reset(new raw_fd_ostream(Filename, EC, sys::fs::OF_None));
850   else
851     OS.reset(new raw_fd_ostream(Filename, EC, sys::fs::OF_TextWithCRLF));
852   if (EC)
853     return EC;
854 
855   return create(OS, Format);
856 }
857 
858 /// Create a sample profile stream writer based on the specified format.
859 ///
860 /// \param OS The output stream to store the profile data to.
861 ///
862 /// \param Format Encoding format for the profile file.
863 ///
864 /// \returns an error code indicating the status of the created writer.
865 ErrorOr<std::unique_ptr<SampleProfileWriter>>
866 SampleProfileWriter::create(std::unique_ptr<raw_ostream> &OS,
867                             SampleProfileFormat Format) {
868   std::error_code EC;
869   std::unique_ptr<SampleProfileWriter> Writer;
870 
871   // Currently only Text and Extended Binary format are supported for CSSPGO.
872   if ((FunctionSamples::ProfileIsCS || FunctionSamples::ProfileIsProbeBased) &&
873       (Format == SPF_Binary || Format == SPF_Compact_Binary))
874     return sampleprof_error::unsupported_writing_format;
875 
876   if (Format == SPF_Binary)
877     Writer.reset(new SampleProfileWriterRawBinary(OS));
878   else if (Format == SPF_Ext_Binary)
879     Writer.reset(new SampleProfileWriterExtBinary(OS));
880   else if (Format == SPF_Compact_Binary)
881     Writer.reset(new SampleProfileWriterCompactBinary(OS));
882   else if (Format == SPF_Text)
883     Writer.reset(new SampleProfileWriterText(OS));
884   else if (Format == SPF_GCC)
885     EC = sampleprof_error::unsupported_writing_format;
886   else
887     EC = sampleprof_error::unrecognized_format;
888 
889   if (EC)
890     return EC;
891 
892   Writer->Format = Format;
893   return std::move(Writer);
894 }
895 
896 void SampleProfileWriter::computeSummary(const SampleProfileMap &ProfileMap) {
897   SampleProfileSummaryBuilder Builder(ProfileSummaryBuilder::DefaultCutoffs);
898   Summary = Builder.computeSummaryForProfiles(ProfileMap);
899 }
900