llvm/unittests/tools/llvm-profdata/OutputSizeLimitTest.cpp

   1 //===- llvm/unittests/tools/llvm-profdata/OutputSizeLimitTest.cpp ---------===//
   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 #include "llvm/ProfileData/SampleProfReader.h"
  10 #include "llvm/ProfileData/SampleProfWriter.h"
  11 #include "llvm/Support/FileSystem.h"
  12 #include "llvm/Support/VirtualFileSystem.h"
  13 #include "llvm/Testing/Support/Error.h"
  14 #include "gtest/gtest.h"
  15
  16 using namespace llvm;
  17 using llvm::unittest::TempFile;
  18
  19 std::string Input1 = R"(main:184019:0
  20  4: 534
  21  4.2: 534
  22  5: 1075
  23  5.1: 1075
  24  6: 2080
  25  7: 534
  26  9: 2064 _Z3bari:1471 _Z3fooi:631
  27  10: inline1:1000
  28   1: 1000
  29  10: inline2:2000
  30   1: 2000
  31 _Z3bari:20301:1437
  32  1: 1437
  33 _Z3fooi:7711:610
  34  1: 610)";
  35
  36 const char EmptyProfile[18] = "\xff\xe5\xd0\xb1\xf4\xc9\x94\xa8\x53\x67";
  37
  38 /// sys::fs and SampleProf mix Error and error_code, making an adapter class
  39 /// to keep code elegant.
  40 template <typename T> class ExpectedErrorOr : public Expected<T> {
  41 public:
  42   ExpectedErrorOr(T &&Obj) : Expected<T>(Obj) {}
  43
  44   ExpectedErrorOr(std::error_code EC) : Expected<T>(errorCodeToError(EC)) {}
  45
  46   ExpectedErrorOr(Error &&E) : Expected<T>(std::move(E)) {}
  47
  48   template <typename U>
  49   ExpectedErrorOr(ErrorOr<U> &&E)
  50       : Expected<T>(errorCodeToError(E.getError())) {}
  51
  52   template <typename U>
  53   ExpectedErrorOr(Expected<U> &&E) : Expected<T>(E.takeError()) {}
  54 };
  55
  56 #define DEF_VAR_RETURN_IF_ERROR(Var, Value)                                    \
  57   auto Var##OrErr = Value;                                                     \
  58   if (!Var##OrErr)                                                             \
  59     return Var##OrErr;                                                         \
  60   auto Var = std::move(Var##OrErr.get())
  61
  62 #define VAR_RETURN_IF_ERROR(Var, Value)                                        \
  63   Var##OrErr = Value;                                                          \
  64   if (!Var##OrErr)                                                             \
  65     return Var##OrErr;                                                         \
  66   Var = std::move(Var##OrErr.get())
  67
  68 #define RETURN_IF_ERROR(Value)                                                 \
  69   if (auto E = Value)                                                          \
  70   return std::move(E)
  71
  72 /// The main testing routine. After rewriting profiles with size limit, check
  73 /// the following:
  74 /// 1. The file size of the new profile is within the size limit.
  75 /// 2. The new profile is a subset of the old profile, and the content of every
  76 /// sample in the new profile is unchanged.
  77 /// Note that even though by default samples with fewest total count are dropped
  78 /// first, this is not a requirement. Samples can be dropped by any order.
  79 static ExpectedErrorOr<void *> RunTest(StringRef Input, size_t SizeLimit,
  80                                        SampleProfileFormat Format,
  81                                        bool Compress = false) {
  82   // Read Input profile.
  83   auto FS = vfs::getRealFileSystem();
  84   LLVMContext Context;
  85   auto InputBuffer = MemoryBuffer::getMemBuffer(Input);
  86   DEF_VAR_RETURN_IF_ERROR(
  87       Reader, SampleProfileReader::create(InputBuffer, Context, *FS));
  88   RETURN_IF_ERROR(Reader->read());
  89   SampleProfileMap OldProfiles = Reader->getProfiles();
  90
  91   // Rewrite it to a temp file with size limit.
  92   TempFile Temp("profile", "afdo", "", true);
  93   bool isEmpty = false;
  94   {
  95     DEF_VAR_RETURN_IF_ERROR(Writer,
  96                             SampleProfileWriter::create(Temp.path(), Format));
  97     if (Compress)
  98       Writer->setToCompressAllSections();
  99     std::error_code EC = Writer->writeWithSizeLimit(OldProfiles, SizeLimit);
 100     // too_large means no sample could be written because SizeLimit is too
 101     // small. Otherwise any other error code indicates unexpected failure.
 102     if (EC == sampleprof_error::too_large)
 103       isEmpty = true;
 104     else if (EC)
 105       return EC;
 106   }
 107
 108   // Read the temp file to get new profiles. Use the default empty profile if
 109   // temp file was not written because size limit is too small.
 110   SampleProfileMap NewProfiles;
 111   InputBuffer = MemoryBuffer::getMemBuffer(StringRef(EmptyProfile, 17));
 112   DEF_VAR_RETURN_IF_ERROR(
 113       NewReader, SampleProfileReader::create(InputBuffer, Context, *FS));
 114   if (!isEmpty) {
 115     VAR_RETURN_IF_ERROR(NewReader, SampleProfileReader::create(
 116                                        Temp.path().str(), Context, *FS));
 117     RETURN_IF_ERROR(NewReader->read());
 118     NewProfiles = NewReader->getProfiles();
 119   }
 120
 121   // Check temp file is actually within size limit.
 122   uint64_t FileSize;
 123   RETURN_IF_ERROR(sys::fs::file_size(Temp.path(), FileSize));
 124   EXPECT_LE(FileSize, SizeLimit);
 125
 126   // For every sample in the new profile, confirm it is in the old profile and
 127   // unchanged.
 128   for (auto Sample : NewProfiles) {
 129     auto FindResult = OldProfiles.find(Sample.second.getContext());
 130     EXPECT_NE(FindResult, OldProfiles.end());
 131     if (FindResult != OldProfiles.end()) {
 132       EXPECT_EQ(Sample.second.getHeadSamples(),
 133                 FindResult->second.getHeadSamples());
 134       EXPECT_EQ(Sample.second, FindResult->second);
 135     }
 136   }
 137   return nullptr;
 138 }
 139
 140 TEST(TestOutputSizeLimit, TestOutputSizeLimitExtBinary) {
 141   for (size_t OutputSizeLimit : {490, 489, 488, 475, 474, 459, 400})
 142     ASSERT_THAT_EXPECTED(
 143         RunTest(Input1, OutputSizeLimit, llvm::sampleprof::SPF_Ext_Binary),
 144         Succeeded());
 145 }
 146
 147 TEST(TestOutputSizeLimit, TestOutputSizeLimitBinary) {
 148   for (size_t OutputSizeLimit : {250, 249, 248, 237, 236, 223, 200})
 149     ASSERT_THAT_EXPECTED(
 150         RunTest(Input1, OutputSizeLimit, llvm::sampleprof::SPF_Binary),
 151         Succeeded());
 152 }
 153
 154 TEST(TestOutputSizeLimit, TestOutputSizeLimitText) {
 155   for (size_t OutputSizeLimit :
 156        {229, 228, 227, 213, 212, 211, 189, 188, 187, 186, 150})
 157     ASSERT_THAT_EXPECTED(
 158         RunTest(Input1, OutputSizeLimit, llvm::sampleprof::SPF_Text),
 159         Succeeded());
 160 }
 161
 162 #if LLVM_ENABLE_ZLIB
 163 TEST(TestOutputSizeLimit, TestOutputSizeLimitExtBinaryCompressed) {
 164   for (size_t OutputSizeLimit :
 165        {507, 506, 505, 494, 493, 492, 483, 482, 481, 480})
 166     ASSERT_THAT_EXPECTED(RunTest(Input1, OutputSizeLimit,
 167                                  llvm::sampleprof::SPF_Ext_Binary, true),
 168                          Succeeded());
 169 }
 170 #endif