From 0972f05b6c93dc950ea692a389c29d43077b86f9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Gerg=C5=91=20S=C3=A1lyi?= Date: Tue, 15 Apr 2025 23:02:38 +0200 Subject: [PATCH] Refactor image loading Start separating image loading logic from wl_shm buffer handling Add fast path decoding directly to wl_shm buffers if possible Add auto-vectorized RGB to BGRA swizzling --- src/image.rs | 444 +++++++++++++++++++++++++++---------------------- src/wayland.rs | 2 +- 2 files changed, 245 insertions(+), 201 deletions(-) diff --git a/src/image.rs b/src/image.rs index dff9625..c744906 100644 --- a/src/image.rs +++ b/src/image.rs @@ -1,15 +1,19 @@ +#![allow(clippy::too_many_arguments)] + use std::{ - fs::read_dir, + fs::{DirEntry, read_dir}, + io, path::Path, }; +use anyhow::{bail, Context}; use fast_image_resize::{ FilterType, PixelType, Resizer, ResizeAlg, ResizeOptions, images::Image, }; -use image::{ImageBuffer, ImageError, ImageReader, Rgb}; -use log::{debug, error}; -use smithay_client_toolkit::shm::slot::{Buffer, SlotPool}; +use image::{ColorType, DynamicImage, ImageBuffer, ImageDecoder, ImageReader}; +use log::{debug, error, warn}; +use smithay_client_toolkit::shm::slot::SlotPool; use smithay_client_toolkit::reexports::client::protocol::wl_shm; use crate::wayland::WorkspaceBackground; @@ -20,218 +24,258 @@ pub fn workspace_bgs_from_output_image_dir( format: wl_shm::Format, brightness: i32, contrast: f32, + width: u32, + height: u32, +) -> anyhow::Result> { + let mut buffers = Vec::new(); + let mut resizer = Resizer::new(); + let stride = match format { + wl_shm::Format::Xrgb8888 => width as usize * 4, + wl_shm::Format::Bgr888 => { + // Align buffer stride to both 4 and pixel format block size + // Not being aligned to 4 caused + // https://github.com/gergo-salyi/multibg-sway/issues/6 + (width as usize * 3).next_multiple_of(4) + }, + _ => unreachable!() + }; + let dir = read_dir(&dir_path).context("Failed to open directory")?; + for entry_result in dir { + match workspace_bg_from_file( + entry_result, + slot_pool, + format, + brightness, + contrast, + width, + height, + stride, + &mut resizer + ) { + Ok(Some(workspace_bg)) => buffers.push(workspace_bg), + Ok(None) => continue, + Err(e) => { + error!("Skipping a directory entry in {:?} \ + due to an error: {:#}", dir_path.as_ref(), e); + continue; + } + } + } + if buffers.is_empty() { + bail!("Found no suitable images in the directory") + } + Ok(buffers) +} + +fn workspace_bg_from_file( + dir_entry_result: io::Result, + slot_pool: &mut SlotPool, + format: wl_shm::Format, + brightness: i32, + contrast: f32, + width: u32, + height: u32, + stride: usize, + resizer: &mut Resizer, +) -> anyhow::Result> { + let entry = dir_entry_result.context("Failed to read direectory")?; + let path = entry.path(); + // Skip dirs + if path.is_dir() { return Ok(None) } + // Use the file stem as the name of the workspace for this wallpaper + let workspace_name = path.file_stem().unwrap() + .to_string_lossy().into_owned(); + let (buffer, canvas) = slot_pool.create_buffer( + width.try_into().unwrap(), + height.try_into().unwrap(), + stride.try_into().unwrap(), + format, + ).context("Failed to create Wayland shared memory buffer")?; + let color_transform = if brightness == 0 && contrast == 0.0 { + ColorTransform::None + } else { + ColorTransform::Legacy { brightness, contrast } + }; + load_wallpaper( + &path, + &mut canvas[..stride * height as usize], + width, + height, + stride, + format, + color_transform, + resizer + ).context("Failed to load wallpaper")?; + Ok(Some(WorkspaceBackground { workspace_name, buffer })) +} + +#[derive(Clone, Copy, PartialEq)] +pub enum ColorTransform { + // Levels { input_max: u8, input_min: u8, output_max: u8, output_min: u8 }, + Legacy { brightness: i32, contrast: f32 }, + None, +} + +fn load_wallpaper( + path: &Path, + dst: &mut [u8], surface_width: u32, surface_height: u32, -) - -> Result, String> -{ - let mut buffers = Vec::new(); - - let dir = read_dir(&dir_path) - .map_err(|e| format!("Failed to open directory: {}", e))?; - - for entry_result in dir { - - let entry = match entry_result { - Ok(entry) => entry, - Err(e) => { - error!( - "Skipping a directory entry in '{:?}' due to an error: {}", - dir_path.as_ref(), e - ); - continue; - } - }; - - let path = entry.path(); - - // Skip dirs - if path.is_dir() { continue } - - // Use the file stem as the name of the workspace for this wallpaper - let workspace_name = path.file_stem().unwrap() - .to_string_lossy().into_owned(); - - let raw_image = match ImageReader::open(&path) - .map_err(ImageError::IoError) - .and_then(|r| r.with_guessed_format() - .map_err(ImageError::IoError) - ) - .and_then(|r| r.decode()) - { - Ok(raw_image) => raw_image, - Err(e) => { - error!( - "Failed to open image '{:?}': {}", - path, e - ); - continue; - } - }; - - // It is possible to adjust the contrast and brightness here - let mut image = raw_image; + surface_stride: usize, + surface_format: wl_shm::Format, + color_transform: ColorTransform, + resizer: &mut Resizer, +) -> anyhow::Result<()> { + let reader = ImageReader::open(path) + .context("Failed to open image file")? + .with_guessed_format() + .context("Failed to read image file format")?; + let file_format = reader.format() + .context("Failed to determine image file format")?; + if !file_format.can_read() { + bail!("Unsupported image file format {file_format:?}") + } else if !file_format.reading_enabled() { + bail!("Application was compiled with support \ + for image file format {file_format:?} disabled") + } + let mut decoder = reader.into_decoder() + .context("Failed to initialize image decoder")?; + let (image_width, image_height) = decoder.dimensions(); + let image_size = decoder.total_bytes(); + let image_color_type = decoder.color_type(); + if image_width == 0 || image_height == 0 || image_size > isize::MAX as u64 { + bail!("Image has invalid dimensions {image_width}x{image_height}") + }; + let image_size = image_size as usize; + debug!("Image {image_width}x{image_height} {image_color_type:?}"); + if image_color_type.has_alpha() { + warn!("Image has alpha channel which will be ignored"); + } + if let Ok(Some(_)) = decoder.icc_profile() { + debug!("Image has an embedded ICC color profile \ + but ICC color profile handling is not yet implemented"); + } + let needs_resize = image_width != surface_width + || image_height != surface_height; + let surface_row_len = surface_width as usize * 3; + if !needs_resize + && image_color_type == ColorType::Rgb8 + && surface_format == wl_shm::Format::Bgr888 + && color_transform == ColorTransform::None + && surface_row_len == surface_stride + { + debug!("Decoding image directly to destination buffer"); + decoder.read_image(&mut dst[..image_size]) + .context("Failed to decode image")?; + return Ok(()); + } + let mut image = DynamicImage::from_decoder(decoder) + .context("Failed to decode image")?; + if let ColorTransform::Legacy { brightness, contrast } = color_transform { if contrast != 0.0 { image = image.adjust_contrast(contrast) } if brightness != 0 { image = image.brighten(brightness) } + } + let mut image = image.into_rgb8(); + if needs_resize { + debug!("Resizing image from {}x{} to {}x{}", + image_width, image_height, + surface_width, surface_height + ); + let src_image = Image::from_vec_u8( + image_width, + image_height, + image.into_raw(), + PixelType::U8x3, + ).unwrap(); + let mut dst_image = Image::new( + surface_width, + surface_height, + PixelType::U8x3, + ); + resizer.resize( + &src_image, + &mut dst_image, + &ResizeOptions::new() + .fit_into_destination(None) + .resize_alg(ResizeAlg::Convolution(FilterType::Lanczos3)) + ).context("Failed to resize image")?; + image = ImageBuffer::from_raw( + surface_width, + surface_height, + dst_image.into_vec() + ).unwrap(); + } + match surface_format { + wl_shm::Format::Bgr888 => { + if surface_row_len == surface_stride { + dst.copy_from_slice(&image); + } else { + copy_pad_stride( + &image, + dst, + surface_row_len, + surface_stride, + surface_height as usize, + ); + } + }, + wl_shm::Format::Xrgb8888 => { + swizzle_bgra_from_rgb(&image, dst); + }, + _ => unreachable!(), + } + Ok(()) +} - let mut image = image.into_rgb8(); - let image_width = image.width(); - let image_height = image.height(); +fn copy_pad_stride( + src: &[u8], + dst: &mut [u8], + src_stride: usize, + dst_stride: usize, + height: usize, +) { + for row in 0..height { + dst[row * dst_stride..][..src_stride].copy_from_slice( + &src[row * src_stride..][..src_stride] + ); + } +} - if image_width == 0 { - error!( - "Image '{}' has zero width, skipping", workspace_name - ); - continue; - }; - if image_height == 0 { - error!( - "Image '{}' has zero height, skipping", workspace_name - ); - continue; - }; - - if image_width != surface_width || image_height != surface_height - { - debug!("Resizing image '{}' from {}x{} to {}x{}", - workspace_name, - image_width, image_height, - surface_width, surface_height - ); - - let src_image = Image::from_vec_u8( - image_width, - image_height, - image.into_raw(), - PixelType::U8x3, - ).unwrap(); - - let mut dst_image = Image::new( - surface_width, - surface_height, - PixelType::U8x3, - ); - - let mut resizer = Resizer::new(); - resizer.resize( - &src_image, - &mut dst_image, - &ResizeOptions::new() - .fit_into_destination(None) - .resize_alg(ResizeAlg::Convolution(FilterType::Lanczos3)) - ).unwrap(); - - image = ImageBuffer::from_raw( - surface_width, - surface_height, - dst_image.into_vec() - ).unwrap(); +fn swizzle_bgra_from_rgb(src: &[u8], dst: &mut [u8]) { + let pixel_count = dst.len() / 4; + assert_eq!(src.len(), pixel_count * 3); + assert_eq!(dst.len(), pixel_count * 4); + unsafe { + #[cfg(target_arch = "x86_64")] + if is_x86_feature_detected!("avx2") { + bgra_from_rgb_avx2(src, dst, pixel_count); + return } - - let buffer = match format { - wl_shm::Format::Xrgb8888 => - buffer_xrgb8888_from_image(image, slot_pool), - wl_shm::Format::Bgr888 => - buffer_bgr888_from_image(image, slot_pool), - _ => unreachable!() - }; - - buffers.push(WorkspaceBackground { workspace_name, buffer }); - } - - if buffers.is_empty() { - Err("Found 0 suitable images in the directory".to_string()) - } - else { - Ok(buffers) + bgra_from_rgb(src, dst, pixel_count) } } -fn buffer_xrgb8888_from_image( - image: ImageBuffer, Vec>, - slot_pool: &mut SlotPool, -) - -> Buffer -{ - let (buffer, canvas) = slot_pool - .create_buffer( - image.width() as i32, - image.height() as i32, - image.width() as i32 * 4, - wl_shm::Format::Xrgb8888 - ) - .unwrap(); - - let canvas_len = image.len() / 3 * 4; - - let image_pixels = image.pixels(); - let canvas_pixels = canvas[..canvas_len].chunks_exact_mut(4); - - for (image_pixel, canvas_pixel) in image_pixels.zip(canvas_pixels) { - canvas_pixel[0] = image_pixel.0[2]; - canvas_pixel[1] = image_pixel.0[1]; - canvas_pixel[2] = image_pixel.0[0]; - } - - buffer +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "avx2")] +unsafe fn bgra_from_rgb_avx2(src: &[u8], dst: &mut [u8], pixel_count: usize) { + unsafe { bgra_from_rgb(src, dst, pixel_count) } } -fn buffer_bgr888_from_image( - image: ImageBuffer, Vec>, - slot_pool: &mut SlotPool, -) - -> Buffer -{ - // Align buffer stride to both 4 and pixel format block size - // Not being aligned to 4 caused - // https://github.com/gergo-salyi/multibg-sway/issues/6 - const BUFFER_STRIDE_ALIGNEMENT: u32 = 4 * 3; - - let width = image.width(); - let height = image.height(); - let image_stride = width * 3; - - let unaligned_bytes = image_stride % BUFFER_STRIDE_ALIGNEMENT; - - let buffer_stride = - if unaligned_bytes == 0 { - image_stride - } else { - let padding = BUFFER_STRIDE_ALIGNEMENT - unaligned_bytes; - image_stride + padding - }; - - let (buffer, canvas) = slot_pool - .create_buffer( - width.try_into().unwrap(), - height.try_into().unwrap(), - buffer_stride.try_into().unwrap(), - wl_shm::Format::Bgr888 - ) - .unwrap(); - - if unaligned_bytes == 0 { - canvas[..image.len()].copy_from_slice(&image); - } - else { - let height: usize = height.try_into().unwrap(); - let buffer_stride: usize = buffer_stride.try_into().unwrap(); - let image_stride: usize = image_stride.try_into().unwrap(); - - for row in 0..height { - let canvas_start = row * buffer_stride; - let image_start = row * image_stride; - let len = image_stride; - - canvas[canvas_start..(canvas_start + len)].copy_from_slice( - &image.as_raw()[image_start..(image_start + len)] - ); +unsafe fn bgra_from_rgb(src: &[u8], dst: &mut [u8], pixel_count: usize) { + unsafe { + let mut src = src.as_ptr(); + let mut dst = dst.as_mut_ptr(); + for _ in 0..pixel_count { + *dst.add(0) = *src.add(2); // B + *dst.add(1) = *src.add(1); // G + *dst.add(2) = *src.add(0); // R + *dst.add(3) = u8::MAX; // A + src = src.add(3); + dst = dst.add(4); } } - - buffer } diff --git a/src/wayland.rs b/src/wayland.rs index 1a7201c..afffdca 100644 --- a/src/wayland.rs +++ b/src/wayland.rs @@ -293,7 +293,7 @@ logical size: {}x{}, transform: {:?}", }, Err(e) => { error!( - "Failed to get wallpapers for new output '{}' form '{:?}': {}", + "Failed to get wallpapers for new output '{}' form '{:?}': {:#}", output_name, output_wallpaper_dir, e ); return;