[ProfileData] Add unit test infrastructure for sample profile reader/writer
[oota-llvm.git] / include / llvm / ProfileData / SampleProfWriter.h
1 //===- SampleProfWriter.h - Write LLVM sample profile data ----------------===//
2 //
3 //                      The LLVM Compiler Infrastructure
4 //
5 // This file is distributed under the University of Illinois Open Source
6 // License. See LICENSE.TXT for details.
7 //
8 //===----------------------------------------------------------------------===//
9 //
10 // This file contains definitions needed for writing sample profiles.
11 //
12 //===----------------------------------------------------------------------===//
13 #ifndef LLVM_PROFILEDATA_SAMPLEPROFWRITER_H
14 #define LLVM_PROFILEDATA_SAMPLEPROFWRITER_H
15
16 #include "llvm/ADT/MapVector.h"
17 #include "llvm/ADT/StringRef.h"
18 #include "llvm/ProfileData/SampleProf.h"
19 #include "llvm/Support/ErrorOr.h"
20 #include "llvm/Support/FileSystem.h"
21 #include "llvm/Support/raw_ostream.h"
22
23 namespace llvm {
24
25 namespace sampleprof {
26
27 enum SampleProfileFormat { SPF_None = 0, SPF_Text, SPF_Binary, SPF_GCC };
28
29 /// \brief Sample-based profile writer. Base class.
30 class SampleProfileWriter {
31 public:
32   virtual ~SampleProfileWriter() {}
33
34   /// Write sample profiles in \p S for function \p FName.
35   ///
36   /// \returns status code of the file update operation.
37   virtual std::error_code write(StringRef FName, const FunctionSamples &S) = 0;
38
39   /// Write all the sample profiles in the given map of samples.
40   ///
41   /// \returns status code of the file update operation.
42   std::error_code write(const StringMap<FunctionSamples> &ProfileMap) {
43     if (std::error_code EC = writeHeader(ProfileMap))
44       return EC;
45
46     for (const auto &I : ProfileMap) {
47       StringRef FName = I.first();
48       const FunctionSamples &Profile = I.second;
49       if (std::error_code EC = write(FName, Profile))
50         return EC;
51     }
52     return sampleprof_error::success;
53   }
54
55   raw_ostream &getOutputStream() { return *OutputStream; }
56
57   /// Profile writer factory.
58   ///
59   /// Create a new file writer based on the value of \p Format.
60   static ErrorOr<std::unique_ptr<SampleProfileWriter>>
61   create(StringRef Filename, SampleProfileFormat Format);
62
63   /// Create a new stream writer based on the value of \p Format.
64   /// For testing.
65   static ErrorOr<std::unique_ptr<SampleProfileWriter>>
66   create(std::unique_ptr<raw_ostream> &OS, SampleProfileFormat Format);
67
68 protected:
69   SampleProfileWriter(std::unique_ptr<raw_ostream> &OS)
70       : OutputStream(std::move(OS)) {}
71
72   /// \brief Write a file header for the profile file.
73   virtual std::error_code
74   writeHeader(const StringMap<FunctionSamples> &ProfileMap) = 0;
75
76   /// \brief Output stream where to emit the profile to.
77   std::unique_ptr<raw_ostream> OutputStream;
78 };
79
80 /// \brief Sample-based profile writer (text format).
81 class SampleProfileWriterText : public SampleProfileWriter {
82 public:
83   std::error_code write(StringRef FName, const FunctionSamples &S) override;
84
85 protected:
86   SampleProfileWriterText(std::unique_ptr<raw_ostream> &OS)
87       : SampleProfileWriter(OS), Indent(0) {}
88
89   std::error_code
90   writeHeader(const StringMap<FunctionSamples> &ProfileMap) override {
91     return sampleprof_error::success;
92   }
93
94 private:
95   /// Indent level to use when writing.
96   ///
97   /// This is used when printing inlined callees.
98   unsigned Indent;
99
100   friend ErrorOr<std::unique_ptr<SampleProfileWriter>>
101   SampleProfileWriter::create(std::unique_ptr<raw_ostream> &OS,
102                               SampleProfileFormat Format);
103 };
104
105 /// \brief Sample-based profile writer (binary format).
106 class SampleProfileWriterBinary : public SampleProfileWriter {
107 public:
108   std::error_code write(StringRef F, const FunctionSamples &S) override;
109
110 protected:
111   SampleProfileWriterBinary(std::unique_ptr<raw_ostream> &OS)
112       : SampleProfileWriter(OS), NameTable() {}
113
114   std::error_code
115   writeHeader(const StringMap<FunctionSamples> &ProfileMap) override;
116   std::error_code writeNameIdx(StringRef FName);
117   std::error_code writeBody(StringRef FName, const FunctionSamples &S);
118
119 private:
120   void addName(StringRef FName);
121   void addNames(const FunctionSamples &S);
122
123   MapVector<StringRef, uint32_t> NameTable;
124
125   friend ErrorOr<std::unique_ptr<SampleProfileWriter>>
126   SampleProfileWriter::create(std::unique_ptr<raw_ostream> &OS,
127                               SampleProfileFormat Format);
128 };
129
130 } // End namespace sampleprof
131
132 } // End namespace llvm
133
134 #endif // LLVM_PROFILEDATA_SAMPLEPROFWRITER_H