Odin/core/image/png/helpers.odin

/*
	Copyright 2021 Jeroen van Rijn <nom@duclavier.com>.
	Made available under Odin's BSD-2 license.

	List of contributors:
		Jeroen van Rijn: Initial implementation.
		Ginger Bill:     Cosmetic changes.

	These are a few useful utility functions to work with PNG images.
*/
package png

import "core:image"
import "core:compress/zlib"
import coretime "core:time"
import "core:strings"
import "core:bytes"
import "core:mem"
import "core:runtime"

/*
	Cleanup of image-specific data.
	There are other helpers for cleanup of PNG-specific data.
	Those are named *_destroy, where * is the name of the helper.
*/

destroy :: proc(img: ^Image) {
	if img == nil {
		/*
			Nothing to do.
			Load must've returned with an error.
		*/
		return
	}

	bytes.buffer_destroy(&img.pixels)

	if v, ok := img.metadata.(^image.PNG_Info); ok {
		for chunk in v.chunks {
			delete(chunk.data)
		}
		delete(v.chunks)
		free(v)
	}
	free(img)
}

/*
	Chunk helpers
*/

gamma :: proc(c: image.PNG_Chunk) -> (res: f32, ok: bool) {
	if c.header.type != .gAMA || len(c.data) != size_of(gAMA) {
		return {}, false
	}
	gama := (^gAMA)(raw_data(c.data))^
	return f32(gama.gamma_100k) / 100_000.0, true
}

INCHES_PER_METER :: 1000.0 / 25.4

phys :: proc(c: image.PNG_Chunk) -> (res: pHYs, ok: bool) {
	if c.header.type != .pHYs || len(c.data) != size_of(pHYs) {
		return {}, false
	}

	return (^pHYs)(raw_data(c.data))^, true 
}

phys_to_dpi :: proc(p: pHYs) -> (x_dpi, y_dpi: f32) {
	return f32(p.ppu_x) / INCHES_PER_METER, f32(p.ppu_y) / INCHES_PER_METER
}

time :: proc(c: image.PNG_Chunk) -> (res: tIME, ok: bool) {
	if c.header.type != .tIME || len(c.data) != size_of(tIME) {
		return {}, false
	}

	return (^tIME)(raw_data(c.data))^, true
}

core_time :: proc(c: image.PNG_Chunk) -> (t: coretime.Time, ok: bool) {
	if t, png_ok := time(c); png_ok {
		return coretime.datetime_to_time(
			int(t.year), int(t.month),  int(t.day),
			int(t.hour), int(t.minute), int(t.second),
		)
	} else {
		return {}, false
	}
}

text :: proc(c: image.PNG_Chunk) -> (res: Text, ok: bool) {
	runtime.DEFAULT_TEMP_ALLOCATOR_TEMP_GUARD(ignore = context.temp_allocator == context.allocator)

	assert(len(c.data) == int(c.header.length))
	#partial switch c.header.type {
	case .tEXt:
		ok = true

		fields := bytes.split(c.data, sep=[]u8{0}, allocator=context.temp_allocator)
		if len(fields) == 2 {
			res.keyword = strings.clone(string(fields[0]))
			res.text    = strings.clone(string(fields[1]))
		} else {
			ok = false
		}
		return
	case .zTXt:
		ok = true

		fields := bytes.split_n(c.data, sep=[]u8{0}, n=3, allocator=context.temp_allocator)
		if len(fields) != 3 || len(fields[1]) != 0 {
			// Compression method must be 0=Deflate, which thanks to the split above turns
			// into an empty slice
			ok = false; return
		}

		// Set up ZLIB context and decompress text payload.
		buf: bytes.Buffer
		zlib_error := zlib.inflate_from_byte_array(fields[2], &buf)
		defer bytes.buffer_destroy(&buf)
		if zlib_error != nil {
			ok = false; return
		}

		res.keyword = strings.clone(string(fields[0]))
		res.text = strings.clone(bytes.buffer_to_string(&buf))
		return
	case .iTXt:
		ok = true

		s := string(c.data)
		null := strings.index_byte(s, 0)
		if null == -1 {
			ok = false; return
		}
		if len(c.data) < null + 4 {
			// At a minimum, including the \0 following the keyword, we require 5 more bytes.
			ok = false;	return
		}
		res.keyword = strings.clone(string(c.data[:null]))
		rest := c.data[null+1:]

		compression_flag := rest[:1][0]
		if compression_flag > 1 {
			ok = false; return
		}
		compression_method := rest[1:2][0]
		if compression_flag == 1 && compression_method > 0 {
			// Only Deflate is supported
			ok = false; return
		}
		rest = rest[2:]

		// We now expect an optional language keyword and translated keyword, both followed by a \0
		null = strings.index_byte(string(rest), 0)
		if null == -1 {
			ok = false; return
		}
		res.language = strings.clone(string(rest[:null]))
		rest = rest[null+1:]

		null = strings.index_byte(string(rest), 0)
		if null == -1 {
			ok = false; return
		}
		res.keyword_localized = strings.clone(string(rest[:null]))
		rest = rest[null+1:]
		if compression_flag == 0 {
			res.text = strings.clone(string(rest))
		} else {
			// Set up ZLIB context and decompress text payload.
			buf: bytes.Buffer
			zlib_error := zlib.inflate_from_byte_array(rest, &buf)
			defer bytes.buffer_destroy(&buf)
			if zlib_error != nil {

				ok = false; return
			}

			res.text = strings.clone(bytes.buffer_to_string(&buf))
		}
		return
	case:
		// PNG text helper called with an unrecognized chunk type.
		ok = false; return
	}
}

text_destroy :: proc(text: Text) {
	delete(text.keyword)
	delete(text.keyword_localized)
	delete(text.language)
	delete(text.text)
}

iccp :: proc(c: image.PNG_Chunk) -> (res: iCCP, ok: bool) {
	runtime.DEFAULT_TEMP_ALLOCATOR_TEMP_GUARD(ignore = context.temp_allocator == context.allocator)

	fields := bytes.split_n(c.data, sep=[]u8{0}, n=3, allocator=context.temp_allocator)

	if len(fields[0]) < 1 || len(fields[0]) > 79 {
		// Invalid profile name
		return
	}

	if len(fields[1]) != 0 {
		// Compression method should be a zero, which the split turned into an empty slice.
		return
	}

	// Set up ZLIB context and decompress iCCP payload
	buf: bytes.Buffer
	zlib_error := zlib.inflate_from_byte_array(fields[2], &buf)
	if zlib_error != nil {
		bytes.buffer_destroy(&buf)
		return
	}

	res.name = strings.clone(string(fields[0]))
	res.profile = bytes.buffer_to_bytes(&buf)
	ok = true
	return
}

iccp_destroy :: proc(i: iCCP) {
	delete(i.name)

	delete(i.profile)

}

srgb :: proc(c: image.PNG_Chunk) -> (res: sRGB, ok: bool) {
	if c.header.type != .sRGB || len(c.data) != size_of(sRGB_Rendering_Intent) {
		return {}, false
	}

	res.intent = sRGB_Rendering_Intent(c.data[0])
	if res.intent > max(sRGB_Rendering_Intent) {
		ok = false; return
	}
	return res, true
}

plte :: proc(c: image.PNG_Chunk) -> (res: PLTE, ok: bool) {
	if c.header.type != .PLTE || c.header.length % 3 != 0 || c.header.length > 768 {
		return {}, false
	}

	plte := mem.slice_data_cast([]image.RGB_Pixel, c.data[:])
	for color, i in plte {
		res.entries[i] = color
	}
	res.used = u16(len(plte))
	return res, true
}

splt :: proc(c: image.PNG_Chunk) -> (res: sPLT, ok: bool) {
	if c.header.type != .sPLT {
		return
	}
	runtime.DEFAULT_TEMP_ALLOCATOR_TEMP_GUARD(ignore = context.temp_allocator == context.allocator)

	fields := bytes.split_n(c.data, sep=[]u8{0}, n=2, allocator=context.temp_allocator)
	if len(fields) != 2 {
		return
	}

	res.depth = fields[1][0]
	if res.depth != 8 && res.depth != 16 {
		return
	}

	data := fields[1][1:]
	count: int

	if res.depth == 8 {
		if len(data) % 6 != 0 {
			return
		}
		count = len(data) / 6
		if count > 256 {
			return
		}

		res.entries = mem.slice_data_cast([][4]u8, data)
	} else { // res.depth == 16
		if len(data) % 10 != 0 {
			return
		}
		count = len(data) / 10
		if count > 256 {
			return
		}

		res.entries = mem.slice_data_cast([][4]u16, data)
	}

	res.name = strings.clone(string(fields[0]))
	res.used = u16(count)
	ok = true
	return
}

splt_destroy :: proc(s: sPLT) {
	delete(s.name)
}

sbit :: proc(c: image.PNG_Chunk) -> (res: [4]u8, ok: bool) {
	/*
		Returns [4]u8 with the significant bits in each channel.
		A channel will contain zero if not applicable to the PNG color type.
	*/

	if len(c.data) < 1 || len(c.data) > 4 {
		ok = false; return
	}
	ok = true

	for i := 0; i < len(c.data); i += 1 {
		res[i] = c.data[i]
	}
	return

}

hist :: proc(c: image.PNG_Chunk) -> (res: hIST, ok: bool) {
	if c.header.type != .hIST {
		return {}, false
	}
	if c.header.length & 1 == 1 || c.header.length > 512 {
		// The entries are u16be, so the length must be even.
		// At most 256 entries must be present
		return {}, false
	}

	ok = true
	data := mem.slice_data_cast([]u16be, c.data)
	i := 0
	for len(data) > 0 {
		// HIST entries are u16be, we unpack them to machine format
		res.entries[i] = u16(data[0])
		i += 1; data = data[1:]
	}
	res.used = u16(i)
	return
}

chrm :: proc(c: image.PNG_Chunk) -> (res: cHRM, ok: bool) {
	ok = true
	if c.header.length != size_of(cHRM_Raw) {
		return {}, false
	}
	chrm := (^cHRM_Raw)(raw_data(c.data))^

	res.w.x = f32(chrm.w.x) / 100_000.0
	res.w.y = f32(chrm.w.y) / 100_000.0
	res.r.x = f32(chrm.r.x) / 100_000.0
	res.r.y = f32(chrm.r.y) / 100_000.0
	res.g.x = f32(chrm.g.x) / 100_000.0
	res.g.y = f32(chrm.g.y) / 100_000.0
	res.b.x = f32(chrm.b.x) / 100_000.0
	res.b.y = f32(chrm.b.y) / 100_000.0
	return
}

exif :: proc(c: image.PNG_Chunk) -> (res: Exif, ok: bool) {

	ok = true

	if len(c.data) < 4 {
		ok = false; return
	}

	if c.data[0] == 'M' && c.data[1] == 'M' {
		res.byte_order = .big_endian
		if c.data[2] != 0 || c.data[3] != 42 {
			ok = false; return
		}
	} else if c.data[0] == 'I' && c.data[1] == 'I' {
		res.byte_order = .little_endian
		if c.data[2] != 42 || c.data[3] != 0 {
			ok = false; return
		}
	} else {
		ok = false; return
	}

	res.data = c.data
	return
}

/*
	General helper functions
*/

compute_buffer_size :: image.compute_buffer_size

/*
	PNG save helpers
*/

when false {

	make_chunk :: proc(c: any, t: Chunk_Type) -> (res: Chunk) {

		data: []u8
		if v, ok := c.([]u8); ok {
			data = v
		} else {
			data = mem.any_to_bytes(c)
		}

		res.header.length = u32be(len(data))
		res.header.type   = t
		res.data   = data

		// CRC the type
		crc    := hash.crc32(mem.any_to_bytes(res.header.type))
		// Extend the CRC with the data
		res.crc = u32be(hash.crc32(data, crc))
		return
	}

	write_chunk :: proc(fd: os.Handle, chunk: Chunk) {
		c := chunk
		// Write length + type
		os.write_ptr(fd, &c.header, 8)
		// Write data
		os.write_ptr(fd, mem.raw_data(c.data), int(c.header.length))
		// Write CRC32
		os.write_ptr(fd, &c.crc, 4)
	}

	write_image_as_png :: proc(filename: string, image: Image) -> (err: Error) {
		profiler.timed_proc()
		using image
		using os
		flags: int = O_WRONLY|O_CREATE|O_TRUNC

		if len(image.pixels) == 0 || len(image.pixels) < image.width * image.height * int(image.channels) {
			return .Invalid_Image_Dimensions
		}

		mode: int = 0
		when ODIN_OS == .Linux || ODIN_OS == .Darwin {
			// NOTE(justasd): 644 (owner read, write; group read; others read)
			mode = S_IRUSR | S_IWUSR | S_IRGRP | S_IROTH
		}

		fd, fderr := open(filename, flags, mode)
		if fderr != 0 {
			return .Cannot_Open_File
		}
		defer close(fd)

		magic := Signature

		write_ptr(fd, &magic, 8)

		ihdr := IHDR{
			width              = u32be(width),
			height             = u32be(height),
			bit_depth          = depth,
			compression_method = 0,
			filter_method      = 0,
			interlace_method   = .None,
		}

		switch channels {
		case 1: ihdr.color_type = Color_Type{}
		case 2: ihdr.color_type = Color_Type{.Alpha}
		case 3: ihdr.color_type = Color_Type{.Color}
		case 4: ihdr.color_type = Color_Type{.Color, .Alpha}
		case:// Unhandled
			return .Unknown_Color_Type
		}
		h := make_chunk(ihdr, .IHDR)
		write_chunk(fd, h)

		bytes_needed := width * height * int(channels) + height
		filter_bytes := mem.make_dynamic_array_len_cap([dynamic]u8, bytes_needed, bytes_needed, context.allocator)
		defer delete(filter_bytes)

		i := 0; j := 0
		// Add a filter byte 0 per pixel row
		for y := 0; y < height; y += 1 {
			filter_bytes[j] = 0; j += 1
			for x := 0; x < width; x += 1 {
				for z := 0; z < channels; z += 1 {
					filter_bytes[j+z] = image.pixels[i+z]
				}
				i += channels; j += channels
			}
		}
		assert(j == bytes_needed)

		a: []u8 = filter_bytes[:]

		out_buf: ^[dynamic]u8
		defer free(out_buf)

		ctx := zlib.ZLIB_Context{
			in_buf  = &a,
			out_buf = out_buf,
		}
		err = zlib.write_zlib_stream_from_memory(&ctx)

		b: []u8
		if err == nil {
			b = ctx.out_buf[:]
		} else {
			return err
		}

		idat := make_chunk(b, .IDAT)

		write_chunk(fd, idat)

		iend := make_chunk([]u8{}, .IEND)
		write_chunk(fd, iend)

		return nil
	}
}