Implement .alpha_add_if_missing for JPEG

2 files changed, 96 insertions, 30 deletions

diff --git a/core/image/jpeg/jpeg.odin b/core/image/jpeg/jpeg.odin
index 58c6dff50..391b4316c 100644
--- a/core/image/jpeg/jpeg.odin
+++ b/core/image/jpeg/jpeg.odin
@@ -927,9 +927,9 @@ load_from_context :: proc(ctx: ^$C, options := Options{}, allocator := context.a
 										cbcr_pixel_column := k / luma_h_sampling_factor + 4 * h
 										cbcr_pixel := cbcr_pixel_row * BLOCK_SIZE + cbcr_pixel_column
 
-										r := cast(i16)math.clamp(cast(f32)y_blk[.Y][i] + 1.402 * cast(f32)cbcr_blk[.Cr][cbcr_pixel] + 128, 0, 255)
-										g := cast(i16)math.clamp(cast(f32)y_blk[.Y][i] - 0.344 * cast(f32)cbcr_blk[.Cb][cbcr_pixel] - 0.714 * cast(f32)cbcr_blk[.Cr][cbcr_pixel] + 128, 0, 255)
-										b := cast(i16)math.clamp(cast(f32)y_blk[.Y][i] + 1.772 * cast(f32)cbcr_blk[.Cb][cbcr_pixel] + 128, 0, 255)
+										r := cast(i16)clamp(cast(f32)y_blk[.Y][i] + 1.402 * cast(f32)cbcr_blk[.Cr][cbcr_pixel] + 128, 0, 255)
+										g := cast(i16)clamp(cast(f32)y_blk[.Y][i] - 0.344 * cast(f32)cbcr_blk[.Cb][cbcr_pixel] - 0.714 * cast(f32)cbcr_blk[.Cr][cbcr_pixel] + 128, 0, 255)
+										b := cast(i16)clamp(cast(f32)y_blk[.Y][i] + 1.772 * cast(f32)cbcr_blk[.Cb][cbcr_pixel] + 128, 0, 255)
 
 										y_blk[.Y][i]  = r
 										y_blk[.Cb][i] = g
@@ -941,28 +941,94 @@ load_from_context :: proc(ctx: ^$C, options := Options{}, allocator := context.a
 					}
 				}
 
+				if .alpha_add_if_missing in options {
+					img.channels += 1
+				}
+
 				if resize(&img.pixels.buf, img.width * img.height * img.channels) != nil {
 					return img, .Unable_To_Allocate_Or_Resize
 				}
 
-				out := mem.slice_data_cast([]image.RGB_Pixel, img.pixels.buf[:])
-				for y in 0..<img.height {
-					mcu_row := y / BLOCK_SIZE
-					pixel_row := y % BLOCK_SIZE
-					for x in 0..<img.width {
-						mcu_col := x / BLOCK_SIZE
-						pixel_col := x % BLOCK_SIZE
-						mcu_idx := mcu_row * block_width + mcu_col
-						pixel_idx := pixel_row * BLOCK_SIZE + pixel_col
-
-						if img.channels == 3 {
-							out[y * img.width + x] = {
+				switch img.channels {
+				case 1:
+					out_idx := 0
+					for y in 0..<img.height {
+						mcu_row   := y / BLOCK_SIZE
+						pixel_row := y % BLOCK_SIZE
+						for x in 0..<img.width {
+							mcu_col   := x / BLOCK_SIZE
+							pixel_col := x % BLOCK_SIZE
+							mcu_idx   := mcu_row   * block_width + mcu_col
+							pixel_idx := pixel_row * BLOCK_SIZE  + pixel_col
+
+							img.pixels.buf[out_idx] = cast(byte)blocks[mcu_idx][.Y][pixel_idx]
+							out_idx += 1
+						}
+					}
+
+				case 2:
+					out     := mem.slice_data_cast([]image.GA_Pixel, img.pixels.buf[:])
+					out_idx := 0
+					for y in 0..<img.height {
+						mcu_row   := y / BLOCK_SIZE
+						pixel_row := y % BLOCK_SIZE
+
+						for x in 0..<img.width {
+							mcu_col   := x / BLOCK_SIZE
+							pixel_col := x % BLOCK_SIZE
+							mcu_idx   := mcu_row   * block_width + mcu_col
+							pixel_idx := pixel_row * BLOCK_SIZE  + pixel_col
+
+							out[out_idx] = {
+								cast(byte)blocks[mcu_idx][.Y][pixel_idx],
+								255, // Alpha
+							}
+							out_idx += 1
+						}
+					}
+
+				case 3:
+					out     := mem.slice_data_cast([]image.RGB_Pixel, img.pixels.buf[:])
+					out_idx := 0
+					for y in 0..<img.height {
+						mcu_row   := y / BLOCK_SIZE
+						pixel_row := y % BLOCK_SIZE
+
+						for x in 0..<img.width {
+							mcu_col   := x / BLOCK_SIZE
+							pixel_col := x % BLOCK_SIZE
+							mcu_idx   := mcu_row   * block_width + mcu_col
+							pixel_idx := pixel_row * BLOCK_SIZE  + pixel_col
+
+							out[out_idx] = {
 								cast(byte)blocks[mcu_idx][.Y][pixel_idx],
 								cast(byte)blocks[mcu_idx][.Cb][pixel_idx],
 								cast(byte)blocks[mcu_idx][.Cr][pixel_idx],
 							}
-						} else {
-							img.pixels.buf[y * img.width + x] = cast(byte)blocks[mcu_idx][.Y][pixel_idx]
+							out_idx += 1
+						}
+					}
+
+				case 4:
+					out     := mem.slice_data_cast([]image.RGBA_Pixel, img.pixels.buf[:])
+					out_idx := 0
+					for y in 0..<img.height {
+						mcu_row   := y / BLOCK_SIZE
+						pixel_row := y % BLOCK_SIZE
+
+						for x in 0..<img.width {
+							mcu_col   := x / BLOCK_SIZE
+							pixel_col := x % BLOCK_SIZE
+							mcu_idx   := mcu_row   * block_width + mcu_col
+							pixel_idx := pixel_row * BLOCK_SIZE  + pixel_col
+
+							out[out_idx] = {
+								cast(byte)blocks[mcu_idx][.Y][pixel_idx],
+								cast(byte)blocks[mcu_idx][.Cb][pixel_idx],
+								cast(byte)blocks[mcu_idx][.Cr][pixel_idx],
+								255, // Alpha
+							}
+							out_idx += 1
 						}
 					}
 				}
diff --git a/tests/core/image/test_core_image.odin b/tests/core/image/test_core_image.odin
index 82ff2704a..3c7efc8a5 100644
--- a/tests/core/image/test_core_image.odin
+++ b/tests/core/image/test_core_image.odin
@@ -2367,7 +2367,7 @@ Basic_JPG_Tests := []Test{
 	{
 		"emblem-1024", {
 			{Default,   nil, {1024, 1024, 3, 8}, 0x_46a29e0f},
-			// {Alpha_Add, nil, {1024, 1024, 4, 8}, 0x_46a29e0f},
+			{Alpha_Add, nil, {1024, 1024, 4, 8}, 0x_cae2d532},
 		},
 	},
 	{
@@ -2394,21 +2394,21 @@ run_jpg_suite :: proc(t: ^testing.T, suite: []Test) {
 			passed := (test.expected_error == nil && err == nil) || (test.expected_error == err)
 			testing.expectf(t, passed, "%q failed to load with error %v.", file.file, err)
 
-			if err == nil { // No point in running the other tests if it didn't load.
-				pixels := bytes.buffer_to_bytes(&img.pixels)
+			// No point in running the other tests if it didn't load.
+			(err == nil) or_continue
 
-				dims   := Dims{img.width, img.height, img.channels, img.depth}
-				testing.expectf(t, test.dims == dims, "%v has %v, expected: %v.", file.file, dims, test.dims)
+			pixels := bytes.buffer_to_bytes(&img.pixels)
+			dims   := Dims{img.width, img.height, img.channels, img.depth}
+			testing.expectf(t, test.dims == dims, "%v has %v, expected: %v.", file.file, dims, test.dims)
 
-				img_hash := hash.crc32(pixels)
-				testing.expectf(t, test.hash == img_hash, "%v test #1's hash is %08x, expected %08x with %v.", file.file, img_hash, test.hash, test.options)
+			img_hash := hash.crc32(pixels)
+			testing.expectf(t, test.hash == img_hash, "%v test #1's hash is %08x, expected %08x with %v.", file.file, img_hash, test.hash, test.options)
 
-				// Optionally save to BMP file to check file loaded properly during development
-				when false  {
-					test_bmp := strings.concatenate({TEST_SUITE_PATH_JPG, "/", file.file, ".bmp"}, context.temp_allocator)
-					save_err := bmp.save(test_bmp, img)
-					testing.expectf(t, save_err == nil, "expected saving to BMP not to raise error, got %v", save_err)
-				}
+			// Optionally save to QOI file to check file loaded properly during development
+			when false {
+				test_qoi := strings.concatenate({TEST_SUITE_PATH_JPG, "/", file.file, ".qoi"}, context.temp_allocator)
+				save_err := qoi.save(test_qoi, img)
+				testing.expectf(t, save_err == nil, "expected saving to QOI not to raise error, got %v", save_err)
 			}
 		}
 	}