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