[SampleProfileLoader] Fix integer overflow in generateMDProfMetadata (#90217)
[llvm-project.git] / llvm / lib / ObjectYAML / GOFFEmitter.cpp
blob345904407e1d24d6e571c28f84d5a1fd66d4094d
1 //===- yaml2goff - Convert YAML to a GOFF object file ---------------------===//
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 /// \file
10 /// The GOFF component of yaml2obj.
11 ///
12 //===----------------------------------------------------------------------===//
14 #include "llvm/ADT/IndexedMap.h"
15 #include "llvm/ObjectYAML/ObjectYAML.h"
16 #include "llvm/ObjectYAML/yaml2obj.h"
17 #include "llvm/Support/ConvertEBCDIC.h"
18 #include "llvm/Support/Endian.h"
19 #include "llvm/Support/raw_ostream.h"
21 using namespace llvm;
23 namespace {
25 // Common flag values on records.
26 enum {
27 // Flag: This record is continued.
28 Rec_Continued = 1,
30 // Flag: This record is a continuation.
31 Rec_Continuation = 1 << (8 - 6 - 1),
34 template <typename ValueType> struct BinaryBeImpl {
35 ValueType Value;
36 BinaryBeImpl(ValueType V) : Value(V) {}
39 template <typename ValueType>
40 raw_ostream &operator<<(raw_ostream &OS, const BinaryBeImpl<ValueType> &BBE) {
41 char Buffer[sizeof(BBE.Value)];
42 support::endian::write<ValueType, llvm::endianness::big, support::unaligned>(
43 Buffer, BBE.Value);
44 OS.write(Buffer, sizeof(BBE.Value));
45 return OS;
48 template <typename ValueType> BinaryBeImpl<ValueType> binaryBe(ValueType V) {
49 return BinaryBeImpl<ValueType>(V);
52 struct ZerosImpl {
53 size_t NumBytes;
56 raw_ostream &operator<<(raw_ostream &OS, const ZerosImpl &Z) {
57 OS.write_zeros(Z.NumBytes);
58 return OS;
61 ZerosImpl zeros(const size_t NumBytes) { return ZerosImpl{NumBytes}; }
63 // The GOFFOstream is responsible to write the data into the fixed physical
64 // records of the format. A user of this class announces the start of a new
65 // logical record and the size of its payload. While writing the payload, the
66 // physical records are created for the data. Possible fill bytes at the end of
67 // a physical record are written automatically.
68 class GOFFOstream : public raw_ostream {
69 public:
70 explicit GOFFOstream(raw_ostream &OS)
71 : OS(OS), LogicalRecords(0), RemainingSize(0), NewLogicalRecord(false) {
72 SetBufferSize(GOFF::PayloadLength);
75 ~GOFFOstream() { finalize(); }
77 void makeNewRecord(GOFF::RecordType Type, size_t Size) {
78 fillRecord();
79 CurrentType = Type;
80 RemainingSize = Size;
81 if (size_t Gap = (RemainingSize % GOFF::PayloadLength))
82 RemainingSize += GOFF::PayloadLength - Gap;
83 NewLogicalRecord = true;
84 ++LogicalRecords;
87 void finalize() { fillRecord(); }
89 uint32_t logicalRecords() { return LogicalRecords; }
91 private:
92 // The underlying raw_ostream.
93 raw_ostream &OS;
95 // The number of logical records emitted so far.
96 uint32_t LogicalRecords;
98 // The remaining size of this logical record, including fill bytes.
99 size_t RemainingSize;
101 // The type of the current (logical) record.
102 GOFF::RecordType CurrentType;
104 // Signals start of new record.
105 bool NewLogicalRecord;
107 // Return the number of bytes left to write until next physical record.
108 // Please note that we maintain the total number of bytes left, not the
109 // written size.
110 size_t bytesToNextPhysicalRecord() {
111 size_t Bytes = RemainingSize % GOFF::PayloadLength;
112 return Bytes ? Bytes : GOFF::PayloadLength;
115 // Write the record prefix of a physical record, using the current record
116 // type.
117 static void writeRecordPrefix(raw_ostream &OS, GOFF::RecordType Type,
118 size_t RemainingSize,
119 uint8_t Flags = Rec_Continuation) {
120 uint8_t TypeAndFlags = Flags | (Type << 4);
121 if (RemainingSize > GOFF::RecordLength)
122 TypeAndFlags |= Rec_Continued;
123 OS << binaryBe(static_cast<unsigned char>(GOFF::PTVPrefix))
124 << binaryBe(static_cast<unsigned char>(TypeAndFlags))
125 << binaryBe(static_cast<unsigned char>(0));
128 // Fill the last physical record of a logical record with zero bytes.
129 void fillRecord() {
130 assert((GetNumBytesInBuffer() <= RemainingSize) &&
131 "More bytes in buffer than expected");
132 size_t Remains = RemainingSize - GetNumBytesInBuffer();
133 if (Remains) {
134 assert((Remains < GOFF::RecordLength) &&
135 "Attempting to fill more than one physical record");
136 raw_ostream::write_zeros(Remains);
138 flush();
139 assert(RemainingSize == 0 && "Not fully flushed");
140 assert(GetNumBytesInBuffer() == 0 && "Buffer not fully empty");
143 // See raw_ostream::write_impl.
144 void write_impl(const char *Ptr, size_t Size) override {
145 assert((RemainingSize >= Size) && "Attempt to write too much data");
146 assert(RemainingSize && "Logical record overflow");
147 if (!(RemainingSize % GOFF::PayloadLength)) {
148 writeRecordPrefix(OS, CurrentType, RemainingSize,
149 NewLogicalRecord ? 0 : Rec_Continuation);
150 NewLogicalRecord = false;
152 assert(!NewLogicalRecord &&
153 "New logical record not on physical record boundary");
155 size_t Idx = 0;
156 while (Size > 0) {
157 size_t BytesToWrite = bytesToNextPhysicalRecord();
158 if (BytesToWrite > Size)
159 BytesToWrite = Size;
160 OS.write(Ptr + Idx, BytesToWrite);
161 Idx += BytesToWrite;
162 Size -= BytesToWrite;
163 RemainingSize -= BytesToWrite;
164 if (Size) {
165 writeRecordPrefix(OS, CurrentType, RemainingSize);
170 // Return the current position within the stream, not counting the bytes
171 // currently in the buffer.
172 uint64_t current_pos() const override { return OS.tell(); }
175 class GOFFState {
176 void writeHeader(GOFFYAML::FileHeader &FileHdr);
177 void writeEnd();
179 void reportError(const Twine &Msg) {
180 ErrHandler(Msg);
181 HasError = true;
184 GOFFState(raw_ostream &OS, GOFFYAML::Object &Doc,
185 yaml::ErrorHandler ErrHandler)
186 : GW(OS), Doc(Doc), ErrHandler(ErrHandler), HasError(false) {}
188 ~GOFFState() { GW.finalize(); }
190 bool writeObject();
192 public:
193 static bool writeGOFF(raw_ostream &OS, GOFFYAML::Object &Doc,
194 yaml::ErrorHandler ErrHandler);
196 private:
197 GOFFOstream GW;
198 GOFFYAML::Object &Doc;
199 yaml::ErrorHandler ErrHandler;
200 bool HasError;
203 void GOFFState::writeHeader(GOFFYAML::FileHeader &FileHdr) {
204 SmallString<16> CCSIDName;
205 if (std::error_code EC =
206 ConverterEBCDIC::convertToEBCDIC(FileHdr.CharacterSetName, CCSIDName))
207 reportError("Conversion error on " + FileHdr.CharacterSetName);
208 if (CCSIDName.size() > 16) {
209 reportError("CharacterSetName too long");
210 CCSIDName.resize(16);
212 SmallString<16> LangProd;
213 if (std::error_code EC = ConverterEBCDIC::convertToEBCDIC(
214 FileHdr.LanguageProductIdentifier, LangProd))
215 reportError("Conversion error on " + FileHdr.LanguageProductIdentifier);
216 if (LangProd.size() > 16) {
217 reportError("LanguageProductIdentifier too long");
218 LangProd.resize(16);
221 GW.makeNewRecord(GOFF::RT_HDR, GOFF::PayloadLength);
222 GW << binaryBe(FileHdr.TargetEnvironment) // TargetEnvironment
223 << binaryBe(FileHdr.TargetOperatingSystem) // TargetOperatingSystem
224 << zeros(2) // Reserved
225 << binaryBe(FileHdr.CCSID) // CCSID
226 << CCSIDName // CharacterSetName
227 << zeros(16 - CCSIDName.size()) // Fill bytes
228 << LangProd // LanguageProductIdentifier
229 << zeros(16 - LangProd.size()) // Fill bytes
230 << binaryBe(FileHdr.ArchitectureLevel); // ArchitectureLevel
231 // The module propties are optional. Figure out if we need to write them.
232 uint16_t ModPropLen = 0;
233 if (FileHdr.TargetSoftwareEnvironment)
234 ModPropLen = 3;
235 else if (FileHdr.InternalCCSID)
236 ModPropLen = 2;
237 if (ModPropLen) {
238 GW << binaryBe(ModPropLen) << zeros(6);
239 if (ModPropLen >= 2)
240 GW << binaryBe(FileHdr.InternalCCSID ? *FileHdr.InternalCCSID : 0);
241 if (ModPropLen >= 3)
242 GW << binaryBe(FileHdr.TargetSoftwareEnvironment
243 ? *FileHdr.TargetSoftwareEnvironment
244 : 0);
248 void GOFFState::writeEnd() {
249 GW.makeNewRecord(GOFF::RT_END, GOFF::PayloadLength);
250 GW << binaryBe(uint8_t(0)) // No entry point
251 << binaryBe(uint8_t(0)) // No AMODE
252 << zeros(3) // Reserved
253 << binaryBe(GW.logicalRecords());
254 // No entry point yet. Automatically fill remaining space with zero bytes.
255 GW.finalize();
258 bool GOFFState::writeObject() {
259 writeHeader(Doc.Header);
260 if (HasError)
261 return false;
262 writeEnd();
263 return true;
266 bool GOFFState::writeGOFF(raw_ostream &OS, GOFFYAML::Object &Doc,
267 yaml::ErrorHandler ErrHandler) {
268 GOFFState State(OS, Doc, ErrHandler);
269 return State.writeObject();
271 } // namespace
273 namespace llvm {
274 namespace yaml {
276 bool yaml2goff(llvm::GOFFYAML::Object &Doc, raw_ostream &Out,
277 ErrorHandler ErrHandler) {
278 return GOFFState::writeGOFF(Out, Doc, ErrHandler);
281 } // namespace yaml
282 } // namespace llvm