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

@@ -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 image::{ColorType, DynamicImage, ImageBuffer, ImageDecoder, ImageReader};
-use log::{debug, error};
+use log::{debug, error, warn};
-use smithay_client_toolkit::shm::slot::{Buffer, SlotPool};
+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,
-    surface_width: u32,
+    width: u32,
-    surface_height: u32,
+    height: u32,
-)
+) -> anyhow::Result<Vec<WorkspaceBackground>> {
-    -> Result<Vec<WorkspaceBackground>, String>
-{
     let mut buffers = Vec::new();
-
+    let mut resizer = Resizer::new();
-    let dir = read_dir(&dir_path)
+    let stride = match format {
-        .map_err(|e| format!("Failed to open directory: {}", e))?;
+        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(
-        let entry = match entry_result {
+            entry_result,
-            Ok(entry) => entry,
+            slot_pool,
+            format,
+            brightness,
+            contrast,
+            width,
+            height,
+            stride,
+            &mut resizer
+        ) {
+            Ok(Some(workspace_bg)) => buffers.push(workspace_bg),
+            Ok(None) => continue,
             Err(e) => {
-                error!(
+                error!("Skipping a directory entry in {:?} \
-                    "Skipping a directory entry in '{:?}' due to an error: {}",
+                    due to an error: {:#}", dir_path.as_ref(), e);
-                    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<DirEntry>,
+    slot_pool: &mut SlotPool,
+    format: wl_shm::Format,
+    brightness: i32,
+    contrast: f32,
+    width: u32,
+    height: u32,
+    stride: usize,
+    resizer: &mut Resizer,
+) -> anyhow::Result<Option<WorkspaceBackground>> {
+    let entry = dir_entry_result.context("Failed to read direectory")?;
     let path = entry.path();
-
     // Skip dirs
-        if path.is_dir() { continue }
+    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(
-        let raw_image = match ImageReader::open(&path)
+        width.try_into().unwrap(),
-            .map_err(ImageError::IoError)
+        height.try_into().unwrap(),
-            .and_then(|r| r.with_guessed_format()
+        stride.try_into().unwrap(),
-                .map_err(ImageError::IoError)
+        format,
-            )
+    ).context("Failed to create Wayland shared memory buffer")?;
-            .and_then(|r| r.decode())
+    let color_transform = if brightness == 0 && contrast == 0.0 {
-        {
+        ColorTransform::None
-            Ok(raw_image) => raw_image,
+    } else {
-            Err(e) => {
+        ColorTransform::Legacy { brightness, contrast }
-                error!(
-                    "Failed to open image '{:?}': {}",
-                    path, e
-                );
-                continue;
-            }
     };
+    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 }))
+}
 
-        // It is possible to adjust the contrast and brightness here
+#[derive(Clone, Copy, PartialEq)]
-        let mut image = raw_image;
+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,
+    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();
-        let image_width = image.width();
+    if needs_resize {
-        let image_height = image.height();
+        debug!("Resizing image from {}x{} to {}x{}",
-
-        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();
+        ).context("Failed to resize image")?;
-
         image = ImageBuffer::from_raw(
             surface_width,
             surface_height,
             dst_image.into_vec()
         ).unwrap();
     }
-
+    match surface_format {
-        let buffer = match format {
+        wl_shm::Format::Bgr888 => {
-            wl_shm::Format::Xrgb8888 =>
+            if surface_row_len == surface_stride {
-                buffer_xrgb8888_from_image(image, slot_pool),
+                dst.copy_from_slice(&image);
-            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)
-    }
-}
-
-fn buffer_xrgb8888_from_image(
-    image: ImageBuffer<Rgb<u8>, Vec<u8>>,
-    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
-}
-
-fn buffer_bgr888_from_image(
-    image: ImageBuffer<Rgb<u8>, Vec<u8>>,
-    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;
+                copy_pad_stride(
-        image_stride + padding
+                    &image,
-    };
+                    dst,
-
+                    surface_row_len,
-    let (buffer, canvas) = slot_pool
+                    surface_stride,
-        .create_buffer(
+                    surface_height as usize,
-            width.try_into().unwrap(),
+                );
-            height.try_into().unwrap(),
+            }
-            buffer_stride.try_into().unwrap(),
+        },
-            wl_shm::Format::Bgr888
+        wl_shm::Format::Xrgb8888 => {
-        )
+            swizzle_bgra_from_rgb(&image, dst);
-        .unwrap();
+        },
-
+        _ => unreachable!(),
-    if unaligned_bytes == 0 {
+    }
-        canvas[..image.len()].copy_from_slice(&image);
+    Ok(())
 }
-    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();
 
+fn copy_pad_stride(
+    src: &[u8],
+    dst: &mut [u8],
+    src_stride: usize,
+    dst_stride: usize,
+    height: usize,
+) {
     for row in 0..height {
-            let canvas_start = row * buffer_stride;
+        dst[row * dst_stride..][..src_stride].copy_from_slice(
-            let image_start = row * image_stride;
+            &src[row * src_stride..][..src_stride]
-            let len = image_stride;
-
-            canvas[canvas_start..(canvas_start + len)].copy_from_slice(
-                &image.as_raw()[image_start..(image_start + len)]
         );
     }
 }
 
-    buffer
+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
+        }
+        bgra_from_rgb(src, dst, pixel_count)
+    }
+}
+
+#[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) }
+}
+
+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);
+        }
+    }
 }

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;