diff --git a/apps/desktop-gpui/kernels/Cargo.toml b/apps/desktop-gpui/kernels/Cargo.toml index 7547d3518c..fbfcf8da89 100644 --- a/apps/desktop-gpui/kernels/Cargo.toml +++ b/apps/desktop-gpui/kernels/Cargo.toml @@ -26,3 +26,7 @@ harness = false [[bench]] name = "png_encode" harness = false + +[[bench]] +name = "frame_export" +harness = false diff --git a/apps/desktop-gpui/kernels/benches/frame_export.rs b/apps/desktop-gpui/kernels/benches/frame_export.rs new file mode 100644 index 0000000000..57740639dd --- /dev/null +++ b/apps/desktop-gpui/kernels/benches/frame_export.rs @@ -0,0 +1,133 @@ +use cap_gpui_kernels::{codec as candidate_codec, screenshot as candidate}; +use std::{hint::black_box, time::Instant}; + +fn opaque_rgba(pixels: usize) -> Vec { + let mut rgba = Vec::with_capacity(pixels * 4); + for i in 0..pixels { + rgba.extend_from_slice(&[ + (i & 0xFF) as u8, + ((i >> 8) & 0xFF) as u8, + (i.wrapping_mul(17) & 0xFF) as u8, + 255, + ]); + } + rgba +} + +fn mixed_rgba(pixels: usize) -> Vec { + let mut rgba = Vec::with_capacity(pixels * 4); + for i in 0..pixels { + rgba.extend_from_slice(&[ + (i & 0xFF) as u8, + ((i >> 8) & 0xFF) as u8, + (i.wrapping_mul(17) & 0xFF) as u8, + if i % 7 == 0 { (i % 200) as u8 } else { 255 }, + ]); + } + rgba +} + +fn fused_two_step(rgba: &[u8]) -> Vec { + baseline::rgba_to_rgb(&baseline::flatten_onto_white(rgba)) +} + +fn run_flatten(name: &str, rgba: &[u8]) { + assert_eq!( + baseline::flatten_onto_white(rgba), + candidate::flatten_onto_white(rgba), + "{name}" + ); + emit( + name, + baseline::flatten_onto_white, + candidate::flatten_onto_white, + rgba, + ); +} + +fn run_strip(name: &str, rgba: &[u8]) { + assert_eq!( + baseline::rgba_to_rgb(rgba), + candidate_codec::rgba_to_rgb(rgba), + "{name}" + ); + emit( + name, + baseline::rgba_to_rgb, + candidate_codec::rgba_to_rgb, + rgba, + ); +} + +fn run_fused(name: &str, rgba: &[u8]) { + assert_eq!( + fused_two_step(rgba), + candidate_codec::flatten_rgba_to_rgb_over_white(rgba), + "{name}" + ); + emit( + name, + fused_two_step, + candidate_codec::flatten_rgba_to_rgb_over_white, + rgba, + ); +} + +fn emit(name: &str, before: fn(T) -> Vec, after: fn(T) -> Vec, input: T) +where + T: Copy, +{ + for sample in 0..30 { + let (old, new) = if sample % 2 == 0 { + (timed(before, input), timed(after, input)) + } else { + let new = timed(after, input); + let old = timed(before, input); + (old, new) + }; + if sample >= 5 { + println!("{name},{},{old},{new}", sample - 5); + } + } +} + +fn timed(f: fn(T) -> Vec, input: T) -> u128 +where + T: Copy, +{ + let start = Instant::now(); + black_box(f(black_box(input))); + start.elapsed().as_nanos() +} + +fn main() { + let hd = opaque_rgba(1920 * 1080); + let uhd_opaque = opaque_rgba(3840 * 2160); + let uhd_mixed = mixed_rgba(3840 * 2160); + println!("workload,sample,before_ns,after_ns"); + run_flatten("flatten-opaque-1080p", &hd); + run_flatten("flatten-opaque-4k", &uhd_opaque); + run_flatten("flatten-mixed-4k", &uhd_mixed); + run_strip("rgba-strip-4k", &uhd_opaque); + run_fused("export-fused-4k", &uhd_mixed); +} + +mod baseline { + pub fn rgba_to_rgb(rgba: &[u8]) -> Vec { + rgba.chunks_exact(4) + .flat_map(|px| [px[0], px[1], px[2]]) + .collect() + } + + pub fn flatten_onto_white(rgba: &[u8]) -> Vec { + let mut rgba = rgba.to_vec(); + for pixel in rgba.chunks_exact_mut(4) { + let alpha = u32::from(pixel[3]); + for channel in &mut pixel[..3] { + *channel = ((u32::from(*channel) * alpha + 255 * (255 - alpha) + 127) / 255) as u8; + } + pixel[3] = 255; + } + rgba + } +} diff --git a/apps/desktop-gpui/kernels/src/codec.rs b/apps/desktop-gpui/kernels/src/codec.rs index b184447c33..334373f068 100644 --- a/apps/desktop-gpui/kernels/src/codec.rs +++ b/apps/desktop-gpui/kernels/src/codec.rs @@ -17,9 +17,35 @@ pub fn read_cursor_events( } pub fn rgba_to_rgb(rgba: &[u8]) -> Vec { - rgba.chunks_exact(4) - .flat_map(|px| [px[0], px[1], px[2]]) - .collect() + let len = rgba.len() / 4 * 3; + let mut rgb = vec![0u8; len]; + for (src, dst) in rgba.chunks_exact(4).zip(rgb.chunks_exact_mut(3)) { + dst[0] = src[0]; + dst[1] = src[1]; + dst[2] = src[2]; + } + rgb +} + +pub fn flatten_rgba_to_rgb_over_white(rgba: &[u8]) -> Vec { + let mut rgb = Vec::with_capacity(rgba.len() / 4 * 3); + for px in rgba.chunks_exact(4) { + let alpha = px[3]; + if alpha == 255 { + rgb.extend_from_slice(&[px[0], px[1], px[2]]); + } else if alpha == 0 { + rgb.extend_from_slice(&[255, 255, 255]); + } else { + let a = u32::from(alpha); + let inv = 255 - a; + rgb.extend_from_slice(&[ + ((u32::from(px[0]) * a + 255 * inv + 127) / 255) as u8, + ((u32::from(px[1]) * a + 255 * inv + 127) / 255) as u8, + ((u32::from(px[2]) * a + 255 * inv + 127) / 255) as u8, + ]); + } + } + rgb } pub fn encode_png_rgba(rgba: &[u8], width: u32, height: u32) -> image::ImageResult> { @@ -192,6 +218,27 @@ mod tests { ); } + #[test] + fn fused_flatten_to_rgb_matches_two_step() { + let rgba: Vec = (0..=255u8) + .flat_map(|alpha| { + (0..=255u8).flat_map(move |channel| { + [channel, 255 - channel, channel.wrapping_mul(17), alpha] + }) + }) + .collect(); + assert_eq!( + flatten_rgba_to_rgb_over_white(&rgba), + rgba_to_rgb(&crate::screenshot::flatten_onto_white(&rgba)) + ); + assert_eq!(flatten_rgba_to_rgb_over_white(&[]), Vec::::new()); + let trailing: [u8; 7] = [10, 20, 30, 40, 50, 60, 70]; + assert_eq!( + flatten_rgba_to_rgb_over_white(&trailing), + rgba_to_rgb(&crate::screenshot::flatten_onto_white(&trailing)) + ); + } + #[test] fn png_round_trips_rgba() { let rgba = [10, 20, 30, 40, 50, 60, 70, 255]; diff --git a/apps/desktop-gpui/kernels/src/screenshot.rs b/apps/desktop-gpui/kernels/src/screenshot.rs index 93f969b67f..e6000dcd10 100644 --- a/apps/desktop-gpui/kernels/src/screenshot.rs +++ b/apps/desktop-gpui/kernels/src/screenshot.rs @@ -19,9 +19,21 @@ pub fn has_translucent_pixel(rgba: &[u8]) -> bool { pub fn flatten_onto_white(rgba: &[u8]) -> Vec { let mut rgba = rgba.to_vec(); for pixel in rgba.chunks_exact_mut(4) { - let alpha = u32::from(pixel[3]); + let alpha = pixel[3]; + if alpha == 255 { + continue; + } + if alpha == 0 { + pixel[0] = 255; + pixel[1] = 255; + pixel[2] = 255; + pixel[3] = 255; + continue; + } + let a = u32::from(alpha); + let inv = 255 - a; for channel in &mut pixel[..3] { - *channel = ((u32::from(*channel) * alpha + 255 * (255 - alpha) + 127) / 255) as u8; + *channel = ((u32::from(*channel) * a + 255 * inv + 127) / 255) as u8; } pixel[3] = 255; } diff --git a/apps/desktop-gpui/src/screenshot_export.rs b/apps/desktop-gpui/src/screenshot_export.rs index 8e98a379c5..5533f18844 100644 --- a/apps/desktop-gpui/src/screenshot_export.rs +++ b/apps/desktop-gpui/src/screenshot_export.rs @@ -238,9 +238,8 @@ fn encode_png(out: &Composited) -> Result, String> { /// share flow only ever picks it when the output is fully opaque, and the /// flatten here is the same belt the browser's own encode wears. fn encode_jpeg(out: &Composited) -> Result, String> { - let flat = flatten_onto_white(out); - let rgb = codec::rgba_to_rgb(&flat.rgba); - codec::encode_jpeg_rgb(&rgb, flat.width, flat.height, 90) + let rgb = codec::flatten_rgba_to_rgb_over_white(&out.rgba); + codec::encode_jpeg_rgb(&rgb, out.width, out.height, 90) .map_err(|e| format!("Failed to encode screenshot export: {e}")) }