Index: include/llvm/ProfileData/SampleProfReader.h =================================================================== --- include/llvm/ProfileData/SampleProfReader.h +++ include/llvm/ProfileData/SampleProfReader.h @@ -267,6 +267,10 @@ static ErrorOr> create(StringRef Filename, LLVMContext &C); + /// \brief Create a sample profile reader from the supplied memory buffer. + static ErrorOr> + create(std::unique_ptr &B, LLVMContext &C); + protected: /// \brief Map every function to its associated profile. /// Index: include/llvm/ProfileData/SampleProfWriter.h =================================================================== --- include/llvm/ProfileData/SampleProfWriter.h +++ include/llvm/ProfileData/SampleProfWriter.h @@ -29,9 +29,6 @@ /// \brief Sample-based profile writer. Base class. class SampleProfileWriter { public: - SampleProfileWriter(StringRef Filename, std::error_code &EC, - sys::fs::OpenFlags Flags) - : OS(Filename, EC, Flags) {} virtual ~SampleProfileWriter() {} /// Write sample profiles in \p S for function \p FName. @@ -55,30 +52,40 @@ return sampleprof_error::success; } + raw_ostream &getOutputStream() { return *OutputStream; } + /// Profile writer factory. /// - /// Create a new writer based on the value of \p Format. + /// Create a new file writer based on the value of \p Format. static ErrorOr> create(StringRef Filename, SampleProfileFormat Format); + /// Create a new stream writer based on the value of \p Format. + /// For testing. + static ErrorOr> + create(std::unique_ptr &OS, SampleProfileFormat Format); + protected: + SampleProfileWriter(std::unique_ptr &OS) + : OutputStream(std::move(OS)) {} + /// \brief Write a file header for the profile file. virtual std::error_code writeHeader(const StringMap &ProfileMap) = 0; /// \brief Output stream where to emit the profile to. - raw_fd_ostream OS; + std::unique_ptr OutputStream; }; /// \brief Sample-based profile writer (text format). class SampleProfileWriterText : public SampleProfileWriter { public: - SampleProfileWriterText(StringRef F, std::error_code &EC) - : SampleProfileWriter(F, EC, sys::fs::F_Text), Indent(0) {} - std::error_code write(StringRef FName, const FunctionSamples &S) override; protected: + SampleProfileWriterText(std::unique_ptr &OS) + : SampleProfileWriter(OS), Indent(0) {} + std::error_code writeHeader(const StringMap &ProfileMap) override { return sampleprof_error::success; @@ -89,17 +96,21 @@ /// /// This is used when printing inlined callees. unsigned Indent; + + friend ErrorOr> + SampleProfileWriter::create(std::unique_ptr &OS, + SampleProfileFormat Format); }; /// \brief Sample-based profile writer (binary format). class SampleProfileWriterBinary : public SampleProfileWriter { public: - SampleProfileWriterBinary(StringRef F, std::error_code &EC) - : SampleProfileWriter(F, EC, sys::fs::F_None), NameTable() {} - std::error_code write(StringRef F, const FunctionSamples &S) override; protected: + SampleProfileWriterBinary(std::unique_ptr &OS) + : SampleProfileWriter(OS), NameTable() {} + std::error_code writeHeader(const StringMap &ProfileMap) override; std::error_code writeNameIdx(StringRef FName); @@ -110,6 +121,10 @@ void addNames(const FunctionSamples &S); MapVector NameTable; + + friend ErrorOr> + SampleProfileWriter::create(std::unique_ptr &OS, + SampleProfileFormat Format); }; } // End namespace sampleprof Index: lib/ProfileData/SampleProfReader.cpp =================================================================== --- lib/ProfileData/SampleProfReader.cpp +++ lib/ProfileData/SampleProfReader.cpp @@ -693,15 +693,27 @@ auto BufferOrError = setupMemoryBuffer(Filename); if (std::error_code EC = BufferOrError.getError()) return EC; + return create(BufferOrError.get(), C); +} - auto Buffer = std::move(BufferOrError.get()); +/// \brief Create a sample profile reader based on the format of the input data. +/// +/// \param B The memory buffer to create the reader from (assumes ownership). +/// +/// \param Reader The reader to instantiate according to \p Filename's format. +/// +/// \param C The LLVM context to use to emit diagnostics. +/// +/// \returns an error code indicating the status of the created reader. +ErrorOr> +SampleProfileReader::create(std::unique_ptr &B, LLVMContext &C) { std::unique_ptr Reader; - if (SampleProfileReaderBinary::hasFormat(*Buffer)) - Reader.reset(new SampleProfileReaderBinary(std::move(Buffer), C)); - else if (SampleProfileReaderGCC::hasFormat(*Buffer)) - Reader.reset(new SampleProfileReaderGCC(std::move(Buffer), C)); - else if (SampleProfileReaderText::hasFormat(*Buffer)) - Reader.reset(new SampleProfileReaderText(std::move(Buffer), C)); + if (SampleProfileReaderBinary::hasFormat(*B)) + Reader.reset(new SampleProfileReaderBinary(std::move(B), C)); + else if (SampleProfileReaderGCC::hasFormat(*B)) + Reader.reset(new SampleProfileReaderGCC(std::move(B), C)); + else if (SampleProfileReaderText::hasFormat(*B)) + Reader.reset(new SampleProfileReaderText(std::move(B), C)); else return sampleprof_error::unrecognized_format; Index: lib/ProfileData/SampleProfWriter.cpp =================================================================== --- lib/ProfileData/SampleProfWriter.cpp +++ lib/ProfileData/SampleProfWriter.cpp @@ -39,6 +39,8 @@ /// it needs to be parsed by the SampleProfileReaderText class. std::error_code SampleProfileWriterText::write(StringRef FName, const FunctionSamples &S) { + auto &OS = *OutputStream; + OS << FName << ":" << S.getTotalSamples(); if (Indent == 0) OS << ":" << S.getHeadSamples(); @@ -84,7 +86,7 @@ const auto &ret = NameTable.find(FName); if (ret == NameTable.end()) return sampleprof_error::truncated_name_table; - encodeULEB128(ret->second, OS); + encodeULEB128(ret->second, *OutputStream); return sampleprof_error::success; } @@ -112,6 +114,8 @@ std::error_code SampleProfileWriterBinary::writeHeader( const StringMap &ProfileMap) { + auto &OS = *OutputStream; + // Write file magic identifier. encodeULEB128(SPMagic(), OS); encodeULEB128(SPVersion(), OS); @@ -134,6 +138,8 @@ std::error_code SampleProfileWriterBinary::writeBody(StringRef FName, const FunctionSamples &S) { + auto &OS = *OutputStream; + if (std::error_code EC = writeNameIdx(FName)) return EC; @@ -176,11 +182,11 @@ /// \returns true if the samples were written successfully, false otherwise. std::error_code SampleProfileWriterBinary::write(StringRef FName, const FunctionSamples &S) { - encodeULEB128(S.getHeadSamples(), OS); + encodeULEB128(S.getHeadSamples(), *OutputStream); return writeBody(FName, S); } -/// \brief Create a sample profile writer based on the specified format. +/// \brief Create a sample profile file writer based on the specified format. /// /// \param Filename The file to create. /// @@ -192,12 +198,36 @@ ErrorOr> SampleProfileWriter::create(StringRef Filename, SampleProfileFormat Format) { std::error_code EC; + std::unique_ptr OS; + if (Format == SPF_Binary) + OS.reset(new raw_fd_ostream(Filename, EC, sys::fs::F_None)); + else + OS.reset(new raw_fd_ostream(Filename, EC, sys::fs::F_Text)); + if (EC) + return EC; + + return create(OS, Format); +} + +/// \brief Create a sample profile stream writer based on the specified format. +/// +/// \param OS The output stream to store the profile data to. +/// +/// \param Writer The writer to instantiate according to the specified format. +/// +/// \param Format Encoding format for the profile file. +/// +/// \returns an error code indicating the status of the created writer. +ErrorOr> +SampleProfileWriter::create(std::unique_ptr &OS, + SampleProfileFormat Format) { + std::error_code EC; std::unique_ptr Writer; if (Format == SPF_Binary) - Writer.reset(new SampleProfileWriterBinary(Filename, EC)); + Writer.reset(new SampleProfileWriterBinary(OS)); else if (Format == SPF_Text) - Writer.reset(new SampleProfileWriterText(Filename, EC)); + Writer.reset(new SampleProfileWriterText(OS)); else if (Format == SPF_GCC) EC = sampleprof_error::unsupported_writing_format; else Index: unittests/ProfileData/CMakeLists.txt =================================================================== --- unittests/ProfileData/CMakeLists.txt +++ unittests/ProfileData/CMakeLists.txt @@ -7,4 +7,5 @@ add_llvm_unittest(ProfileDataTests CoverageMappingTest.cpp InstrProfTest.cpp + SampleProfTest.cpp ) Index: unittests/ProfileData/SampleProfTest.cpp =================================================================== --- /dev/null +++ unittests/ProfileData/SampleProfTest.cpp @@ -0,0 +1,102 @@ +//===- unittest/ProfileData/SampleProfTest.cpp -------------------*- C++ +//-*-===// +// +// The LLVM Compiler Infrastructure +// +// This file is distributed under the University of Illinois Open Source +// License. See LICENSE.TXT for details. +// +//===----------------------------------------------------------------------===// + +#include "llvm/ProfileData/SampleProfReader.h" +#include "llvm/ProfileData/SampleProfWriter.h" +#include "gtest/gtest.h" + +#include + +using namespace llvm; +using namespace sampleprof; + +static ::testing::AssertionResult NoError(std::error_code EC) { + if (!EC) + return ::testing::AssertionSuccess(); + return ::testing::AssertionFailure() << "error " << EC.value() << ": " + << EC.message(); +} + +namespace { + +struct SampleProfTest : ::testing::Test { + std::string Data; + std::unique_ptr OS; + std::unique_ptr Writer; + std::unique_ptr Reader; + + SampleProfTest() + : Data(), OS(new raw_string_ostream(Data)), Writer(), Reader() {} + + void createWriter(SampleProfileFormat Format) { + auto WriterOrErr = SampleProfileWriter::create(OS, Format); + ASSERT_TRUE(NoError(WriterOrErr.getError())); + Writer = std::move(WriterOrErr.get()); + } + + void readProfile(std::unique_ptr &Profile) { + auto ReaderOrErr = SampleProfileReader::create(Profile, getGlobalContext()); + ASSERT_TRUE(NoError(ReaderOrErr.getError())); + Reader = std::move(ReaderOrErr.get()); + } + + void testRoundTrip(SampleProfileFormat Format) { + createWriter(Format); + + StringRef FooName("_Z3fooi"); + FunctionSamples FooSamples; + FooSamples.addTotalSamples(7711); + FooSamples.addHeadSamples(610); + FooSamples.addBodySamples(1, 0, 610); + + StringRef BarName("_Z3bari"); + FunctionSamples BarSamples; + BarSamples.addTotalSamples(20301); + BarSamples.addHeadSamples(1437); + BarSamples.addBodySamples(1, 0, 1437); + + StringMap Profiles; + Profiles[FooName] = std::move(FooSamples); + Profiles[BarName] = std::move(BarSamples); + + std::error_code EC; + EC = Writer->write(Profiles); + ASSERT_TRUE(NoError(EC)); + + Writer->getOutputStream().flush(); + + auto Profile = MemoryBuffer::getMemBufferCopy(Data); + readProfile(Profile); + + EC = Reader->read(); + ASSERT_TRUE(NoError(EC)); + + StringMap &ReadProfiles = Reader->getProfiles(); + ASSERT_EQ(2u, ReadProfiles.size()); + + FunctionSamples &ReadFooSamples = ReadProfiles[FooName]; + ASSERT_EQ(7711u, ReadFooSamples.getTotalSamples()); + ASSERT_EQ(610u, ReadFooSamples.getHeadSamples()); + + FunctionSamples &ReadBarSamples = ReadProfiles[BarName]; + ASSERT_EQ(20301u, ReadBarSamples.getTotalSamples()); + ASSERT_EQ(1437u, ReadBarSamples.getHeadSamples()); + } +}; + +TEST_F(SampleProfTest, roundtrip_text_profile) { + testRoundTrip(SampleProfileFormat::SPF_Text); +} + +TEST_F(SampleProfTest, roundtrip_binary_profile) { + testRoundTrip(SampleProfileFormat::SPF_Binary); +} + +} // end anonymous namespace