1#include "image_loader.h"
2
3#include <base/dbg.h>
4#include <base/io.h>
5#include <base/log.h>
6#include <base/mem.h>
7
8#include <png.h>
9
10#include <csetjmp>
11#include <cstdlib>
12
13bool CByteBufferReader::Read(void *pData, size_t Size)
14{
15 if(m_Error)
16 return false;
17
18 if(m_ReadOffset + Size <= m_Size)
19 {
20 mem_copy(dest: pData, source: &m_pData[m_ReadOffset], size: Size);
21 m_ReadOffset += Size;
22 return true;
23 }
24 else
25 {
26 m_Error = true;
27 return false;
28 }
29}
30
31void CByteBufferWriter::Write(const void *pData, size_t Size)
32{
33 if(!Size)
34 return;
35
36 const size_t WriteOffset = m_vBuffer.size();
37 m_vBuffer.resize(new_size: WriteOffset + Size);
38 mem_copy(dest: &m_vBuffer[WriteOffset], source: pData, size: Size);
39}
40
41class CUserErrorStruct
42{
43public:
44 CByteBufferReader *m_pReader;
45 const char *m_pContextName;
46 std::jmp_buf m_JmpBuf;
47};
48
49[[noreturn]] static void PngErrorCallback(png_structp pPngStruct, png_const_charp pErrorMessage)
50{
51 CUserErrorStruct *pUserStruct = static_cast<CUserErrorStruct *>(png_get_error_ptr(png_ptr: pPngStruct));
52 log_error("png", "error for file \"%s\": %s", pUserStruct->m_pContextName, pErrorMessage);
53 std::longjmp(env: pUserStruct->m_JmpBuf, val: 1);
54}
55
56static void PngWarningCallback(png_structp pPngStruct, png_const_charp pWarningMessage)
57{
58 CUserErrorStruct *pUserStruct = static_cast<CUserErrorStruct *>(png_get_error_ptr(png_ptr: pPngStruct));
59 log_warn("png", "warning for file \"%s\": %s", pUserStruct->m_pContextName, pWarningMessage);
60}
61
62static void PngReadDataCallback(png_structp pPngStruct, png_bytep pOutBytes, png_size_t ByteCountToRead)
63{
64 CByteBufferReader *pReader = static_cast<CByteBufferReader *>(png_get_io_ptr(png_ptr: pPngStruct));
65 if(!pReader->Read(pData: pOutBytes, Size: ByteCountToRead))
66 {
67 png_error(png_ptr: pPngStruct, error_message: "Could not read all bytes, file was too small");
68 }
69}
70
71static CImageInfo::EImageFormat ImageFormatFromChannelCount(int ColorChannelCount)
72{
73 switch(ColorChannelCount)
74 {
75 case 1:
76 return CImageInfo::FORMAT_R;
77 case 2:
78 return CImageInfo::FORMAT_RA;
79 case 3:
80 return CImageInfo::FORMAT_RGB;
81 case 4:
82 return CImageInfo::FORMAT_RGBA;
83 default:
84 dbg_assert_failed("ColorChannelCount invalid");
85 }
86}
87
88static int PngliteIncompatibility(png_structp pPngStruct, png_infop pPngInfo)
89{
90 int Result = 0;
91
92 const int ColorType = png_get_color_type(png_ptr: pPngStruct, info_ptr: pPngInfo);
93 switch(ColorType)
94 {
95 case PNG_COLOR_TYPE_GRAY:
96 case PNG_COLOR_TYPE_RGB:
97 case PNG_COLOR_TYPE_RGB_ALPHA:
98 case PNG_COLOR_TYPE_GRAY_ALPHA:
99 break;
100 default:
101 log_debug("png", "color type %d unsupported by pnglite", ColorType);
102 Result |= CImageLoader::PNGLITE_COLOR_TYPE;
103 }
104
105 const int BitDepth = png_get_bit_depth(png_ptr: pPngStruct, info_ptr: pPngInfo);
106 switch(BitDepth)
107 {
108 case 8:
109 case 16:
110 break;
111 default:
112 log_debug("png", "bit depth %d unsupported by pnglite", BitDepth);
113 Result |= CImageLoader::PNGLITE_BIT_DEPTH;
114 }
115
116 const int InterlaceType = png_get_interlace_type(png_ptr: pPngStruct, info_ptr: pPngInfo);
117 if(InterlaceType != PNG_INTERLACE_NONE)
118 {
119 log_debug("png", "interlace type %d unsupported by pnglite", InterlaceType);
120 Result |= CImageLoader::PNGLITE_INTERLACE_TYPE;
121 }
122
123 if(png_get_compression_type(png_ptr: pPngStruct, info_ptr: pPngInfo) != PNG_COMPRESSION_TYPE_BASE)
124 {
125 log_debug("png", "non-default compression type unsupported by pnglite");
126 Result |= CImageLoader::PNGLITE_COMPRESSION_TYPE;
127 }
128
129 if(png_get_filter_type(png_ptr: pPngStruct, info_ptr: pPngInfo) != PNG_FILTER_TYPE_BASE)
130 {
131 log_debug("png", "non-default filter type unsupported by pnglite");
132 Result |= CImageLoader::PNGLITE_FILTER_TYPE;
133 }
134
135 return Result;
136}
137
138bool CImageLoader::LoadPng(CByteBufferReader &Reader, const char *pContextName, CImageInfo &Image, int &PngliteIncompatible)
139{
140 CUserErrorStruct UserErrorStruct = {.m_pReader: &Reader, .m_pContextName: pContextName, .m_JmpBuf: {}};
141
142 if(setjmp(UserErrorStruct.m_JmpBuf))
143 {
144 return false;
145 }
146
147 png_structp pPngStruct = png_create_read_struct(PNG_LIBPNG_VER_STRING, error_ptr: &UserErrorStruct, error_fn: PngErrorCallback, warn_fn: PngWarningCallback);
148 if(pPngStruct == nullptr)
149 {
150 log_error("png", "libpng internal failure: png_create_read_struct failed.");
151 return false;
152 }
153
154 png_infop pPngInfo = nullptr;
155 png_bytepp pRowPointers = nullptr;
156 int Height = 0; // ensure this is not undefined for the Cleanup function
157 const auto &&Cleanup = [&]() {
158 if(pRowPointers != nullptr)
159 {
160 for(int y = 0; y < Height; ++y)
161 {
162 delete[] pRowPointers[y];
163 }
164 }
165 delete[] pRowPointers;
166 if(pPngInfo != nullptr)
167 {
168 png_destroy_info_struct(png_ptr: pPngStruct, info_ptr_ptr: &pPngInfo);
169 }
170 png_destroy_read_struct(png_ptr_ptr: &pPngStruct, info_ptr_ptr: nullptr, end_info_ptr_ptr: nullptr);
171 };
172 if(setjmp(UserErrorStruct.m_JmpBuf))
173 {
174 Cleanup();
175 return false;
176 }
177
178 pPngInfo = png_create_info_struct(png_ptr: pPngStruct);
179 if(pPngInfo == nullptr)
180 {
181 Cleanup();
182 log_error("png", "libpng internal failure: png_create_info_struct failed.");
183 return false;
184 }
185
186 png_byte aSignature[8];
187 if(!Reader.Read(pData: aSignature, Size: sizeof(aSignature)) || png_sig_cmp(sig: aSignature, start: 0, num_to_check: sizeof(aSignature)) != 0)
188 {
189 Cleanup();
190 log_error("png", "file is not a valid PNG file (signature mismatch).");
191 return false;
192 }
193
194 png_set_read_fn(png_ptr: pPngStruct, io_ptr: (png_bytep)&Reader, read_data_fn: PngReadDataCallback);
195 png_set_sig_bytes(png_ptr: pPngStruct, num_bytes: sizeof(aSignature));
196
197 png_read_info(png_ptr: pPngStruct, info_ptr: pPngInfo);
198
199 if(Reader.Error())
200 {
201 // error already logged
202 Cleanup();
203 return false;
204 }
205
206 const int Width = png_get_image_width(png_ptr: pPngStruct, info_ptr: pPngInfo);
207 Height = png_get_image_height(png_ptr: pPngStruct, info_ptr: pPngInfo);
208 const png_byte BitDepth = png_get_bit_depth(png_ptr: pPngStruct, info_ptr: pPngInfo);
209 const int ColorType = png_get_color_type(png_ptr: pPngStruct, info_ptr: pPngInfo);
210
211 if(Width == 0 || Height == 0)
212 {
213 log_error("png", "image has width (%d) or height (%d) of 0.", Width, Height);
214 Cleanup();
215 return false;
216 }
217
218 if(BitDepth == 16)
219 {
220 png_set_strip_16(png_ptr: pPngStruct);
221 }
222 else if(BitDepth > 8 || BitDepth == 0)
223 {
224 log_error("png", "bit depth %d not supported.", BitDepth);
225 Cleanup();
226 return false;
227 }
228
229 if(ColorType == PNG_COLOR_TYPE_PALETTE)
230 {
231 png_set_palette_to_rgb(png_ptr: pPngStruct);
232 }
233
234 if(ColorType == PNG_COLOR_TYPE_GRAY && BitDepth < 8)
235 {
236 png_set_expand_gray_1_2_4_to_8(png_ptr: pPngStruct);
237 }
238
239 if(png_get_valid(png_ptr: pPngStruct, info_ptr: pPngInfo, PNG_INFO_tRNS))
240 {
241 png_set_tRNS_to_alpha(png_ptr: pPngStruct);
242 }
243
244 png_read_update_info(png_ptr: pPngStruct, info_ptr: pPngInfo);
245
246 const int ColorChannelCount = png_get_channels(png_ptr: pPngStruct, info_ptr: pPngInfo);
247 const int BytesInRow = png_get_rowbytes(png_ptr: pPngStruct, info_ptr: pPngInfo);
248 dbg_assert(BytesInRow == Width * ColorChannelCount, "bytes in row incorrect.");
249
250 pRowPointers = new png_bytep[Height];
251 for(int y = 0; y < Height; ++y)
252 {
253 pRowPointers[y] = new png_byte[BytesInRow];
254 }
255
256 png_read_image(png_ptr: pPngStruct, image: pRowPointers);
257
258 if(!Reader.Error())
259 {
260 Image.m_Width = Width;
261 Image.m_Height = Height;
262 Image.m_Format = ImageFormatFromChannelCount(ColorChannelCount);
263 Image.Allocate();
264 for(int y = 0; y < Height; ++y)
265 {
266 mem_copy(dest: &Image.m_pData[y * BytesInRow], source: pRowPointers[y], size: BytesInRow);
267 }
268 PngliteIncompatible = PngliteIncompatibility(pPngStruct, pPngInfo);
269 }
270
271 Cleanup();
272
273 return !Reader.Error();
274}
275
276bool CImageLoader::LoadPng(IOHANDLE File, const char *pFilename, CImageInfo &Image, int &PngliteIncompatible)
277{
278 if(!File)
279 {
280 log_error("png", "failed to open file for reading. filename='%s'", pFilename);
281 return false;
282 }
283
284 void *pFileData;
285 unsigned FileDataSize;
286 const bool ReadSuccess = io_read_all(io: File, result: &pFileData, result_len: &FileDataSize);
287 io_close(io: File);
288 if(!ReadSuccess)
289 {
290 log_error("png", "failed to read file. filename='%s'", pFilename);
291 return false;
292 }
293
294 CByteBufferReader ImageReader(static_cast<const uint8_t *>(pFileData), FileDataSize);
295
296 const bool LoadResult = CImageLoader::LoadPng(Reader&: ImageReader, pContextName: pFilename, Image, PngliteIncompatible);
297 free(ptr: pFileData);
298 if(!LoadResult)
299 {
300 log_error("png", "failed to load image from file. filename='%s'", pFilename);
301 return false;
302 }
303
304 if(Image.m_Format != CImageInfo::FORMAT_RGB && Image.m_Format != CImageInfo::FORMAT_RGBA)
305 {
306 log_error("png", "image has unsupported format. filename='%s' format='%s'", pFilename, Image.FormatName());
307 Image.Free();
308 return false;
309 }
310
311 return true;
312}
313
314static void PngWriteDataCallback(png_structp pPngStruct, png_bytep pOutBytes, png_size_t ByteCountToWrite)
315{
316 CByteBufferWriter *pWriter = static_cast<CByteBufferWriter *>(png_get_io_ptr(png_ptr: pPngStruct));
317 pWriter->Write(pData: pOutBytes, Size: ByteCountToWrite);
318}
319
320static void PngOutputFlushCallback(png_structp pPngStruct)
321{
322 // no need to flush memory buffer
323}
324
325static int PngColorTypeFromFormat(CImageInfo::EImageFormat Format)
326{
327 switch(Format)
328 {
329 case CImageInfo::FORMAT_R:
330 return PNG_COLOR_TYPE_GRAY;
331 case CImageInfo::FORMAT_RA:
332 return PNG_COLOR_TYPE_GRAY_ALPHA;
333 case CImageInfo::FORMAT_RGB:
334 return PNG_COLOR_TYPE_RGB;
335 case CImageInfo::FORMAT_RGBA:
336 return PNG_COLOR_TYPE_RGBA;
337 default:
338 dbg_assert_failed("Format invalid");
339 }
340}
341
342bool CImageLoader::SavePng(CByteBufferWriter &Writer, const CImageInfo &Image)
343{
344 png_structp pPngStruct = png_create_write_struct(PNG_LIBPNG_VER_STRING, error_ptr: nullptr, error_fn: nullptr, warn_fn: nullptr);
345 if(pPngStruct == nullptr)
346 {
347 log_error("png", "libpng internal failure: png_create_write_struct failed.");
348 return false;
349 }
350
351 png_infop pPngInfo = png_create_info_struct(png_ptr: pPngStruct);
352 if(pPngInfo == nullptr)
353 {
354 png_destroy_read_struct(png_ptr_ptr: &pPngStruct, info_ptr_ptr: nullptr, end_info_ptr_ptr: nullptr);
355 log_error("png", "libpng internal failure: png_create_info_struct failed.");
356 return false;
357 }
358
359 png_set_write_fn(png_ptr: pPngStruct, io_ptr: (png_bytep)&Writer, write_data_fn: PngWriteDataCallback, output_flush_fn: PngOutputFlushCallback);
360
361 png_set_IHDR(png_ptr: pPngStruct, info_ptr: pPngInfo, width: Image.m_Width, height: Image.m_Height, bit_depth: 8, color_type: PngColorTypeFromFormat(Format: Image.m_Format), PNG_INTERLACE_NONE, PNG_COMPRESSION_TYPE_BASE, PNG_FILTER_TYPE_BASE);
362 png_write_info(png_ptr: pPngStruct, info_ptr: pPngInfo);
363
364 png_bytepp pRowPointers = new png_bytep[Image.m_Height];
365 const int WidthBytes = Image.m_Width * Image.PixelSize();
366 ptrdiff_t BufferOffset = 0;
367 for(size_t y = 0; y < Image.m_Height; ++y)
368 {
369 pRowPointers[y] = new png_byte[WidthBytes];
370 mem_copy(dest: pRowPointers[y], source: Image.m_pData + BufferOffset, size: WidthBytes);
371 BufferOffset += (ptrdiff_t)WidthBytes;
372 }
373 png_write_image(png_ptr: pPngStruct, image: pRowPointers);
374 png_write_end(png_ptr: pPngStruct, info_ptr: pPngInfo);
375
376 for(size_t y = 0; y < Image.m_Height; ++y)
377 {
378 delete[] pRowPointers[y];
379 }
380 delete[] pRowPointers;
381
382 png_destroy_info_struct(png_ptr: pPngStruct, info_ptr_ptr: &pPngInfo);
383 png_destroy_write_struct(png_ptr_ptr: &pPngStruct, info_ptr_ptr: nullptr);
384
385 return true;
386}
387
388bool CImageLoader::SavePng(IOHANDLE File, const char *pFilename, const CImageInfo &Image)
389{
390 if(!File)
391 {
392 log_error("png", "failed to open file for writing. filename='%s'", pFilename);
393 return false;
394 }
395
396 CByteBufferWriter Writer;
397 if(!CImageLoader::SavePng(Writer, Image))
398 {
399 // error already logged
400 io_close(io: File);
401 return false;
402 }
403
404 const bool WriteSuccess = io_write(io: File, buffer: Writer.Data(), size: Writer.Size()) == Writer.Size();
405 if(!WriteSuccess)
406 {
407 log_error("png", "failed to write PNG data to file. filename='%s'", pFilename);
408 }
409 io_close(io: File);
410 return WriteSuccess;
411}
412