Implement .alpha_add_if_missing for JPEG

1 files changed, 83 insertions, 17 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
 						}
 					}
 				}