Add --stamp flag

src/main.rs

@@ -1,23 +1,34 @@
-//! Image format is a buffer of `SCREEN_W * SCREEN_H + 1` bytes.
-//! The 4 least-significant bits of each byte is one pixel. The 4 most-significant bits are unused.
-//! Images are stored in landscape format, i.e. the first byte is the top-left pixel, and the
-//! second-to-last byte is the bottom-right pixel. The final byte is a null-terminator, and must be 0.
 use clap::Parser;
-use eyre::{Context, bail, eyre};
-use image::ImageReader;
+use eyre::{Context, ContextCompat, bail, ensure, eyre};
+use image::{
+    DynamicImage, GrayAlphaImage, ImageReader,
+    imageops::{flip_vertical_in_place, overlay},
+};
 use nix::ioctl_write_ptr;
-use std::{fs::File, os::fd::AsRawFd, path::PathBuf};
+use std::{
+    fs::File,
+    os::fd::AsRawFd,
+    path::{Path, PathBuf},
+};
 
 #[derive(Parser)]
 struct Opt {
     /// Path to the image-blob.
     pub path: PathBuf,
+
+    /// Stamp an image onto the final image. "<x>:<y>:<path>"
+    #[clap(long)]
+    pub stamp: Vec<String>,
 }
 
 const SCREEN_W: usize = 1872;
 const SCREEN_H: usize = 1404;
 
-const FILE_SIZE: usize = SCREEN_W * SCREEN_H + 1;
+/// Image format is a buffer of `SCREEN_W * SCREEN_H + 1` bytes.
+/// The 4 least-significant bits of each byte is one pixel. The 4 most-significant bits are unused.
+/// Images are stored in landscape format, i.e. the first byte is the top-left pixel, and the
+/// second-to-last byte is the bottom-right pixel. The final byte is a null-terminator, and must be 0.
+const IOCTL_BUFFER_SIZE: usize = SCREEN_W * SCREEN_H + 1;
 
 // const SPI_IOC_MAGIC: u8 = b'k'; // Defined in linux/spi/spidev.h
 // const SPI_IOC_TYPE_MESSAGE: u8 = 0;
@@ -41,11 +52,22 @@ ioctl_write_ptr!(
     DrmRockchipEbcOffScreen
 );
 
-fn main() -> eyre::Result<()> {
-    let opt = Opt::parse();
-    color_eyre::install()?;
+/// Convert [DynamicImage] into a [GrayAlphaImage].
+/// Note that we use [GrayAlphaImage] instead of `GrayImage` because we need the alpha channel to
+/// be able to add stamps.
+fn to_grayscale(path: &Path, image: DynamicImage) -> GrayAlphaImage {
+    if let DynamicImage::ImageLumaA8(image) = image {
+        return image;
+    }
 
-    let mut image = ImageReader::open(&opt.path)?.decode()?;
+    eprintln!("Image {path:?} is not 16-bit grayscale+alpha. Converting...");
+    image.to_luma_alpha8()
+}
+
+fn load_wallpaper(path: &Path) -> eyre::Result<GrayAlphaImage> {
+    eprintln!("Opening image at {path:?}");
+
+    let mut image = ImageReader::open(path)?.decode()?;
 
     let image_dimensions = (image.width() as usize, image.height() as usize);
 
@@ -55,15 +77,20 @@ fn main() -> eyre::Result<()> {
         _ => bail!("Image must be {SCREEN_W}x{SCREEN_H}"),
     }
 
-    let image = image.to_luma8();
-    let mut pixels = image.into_vec();
-    for pixel in &mut pixels {
-        *pixel >>= 4; // convert 8-bit colorspace to 4-bits.
+    Ok(to_grayscale(path, image))
 }
-    pixels.push(0u8); // add a null-terminator
 
-    // sanity check buffer length
-    assert!(pixels.len() == FILE_SIZE);
+fn main() -> eyre::Result<()> {
+    let opt = Opt::parse();
+    color_eyre::install()?;
+
+    let mut image = load_wallpaper(&opt.path)?;
+
+    for stamp in &opt.stamp {
+        apply_stamp(stamp, &mut image)?;
+    }
+
+    let pixels = image_to_offimage_buffer(image)?;
 
     let set_off_screen_data = DrmRockchipEbcOffScreen {
         info1: 0, // TODO: what is this?
@@ -80,3 +107,74 @@ fn main() -> eyre::Result<()> {
 
     Ok(())
 }
+
+fn image_to_offimage_buffer(
+    mut image: GrayAlphaImage,
+) -> eyre::Result<Box<[u8; IOCTL_BUFFER_SIZE]>> {
+    let mut buf: Box<[u8; IOCTL_BUFFER_SIZE]> = vec![0u8; IOCTL_BUFFER_SIZE]
+        .into_boxed_slice()
+        .try_into()
+        .expect("Buffer is the correct size");
+
+    // y is reversed in the ioctl buffer.
+    flip_vertical_in_place(&mut image);
+
+    let image = image.into_vec();
+
+    ensure!(
+        image.len() / 2 == SCREEN_W * SCREEN_H,
+        "Invalid size of image backing buffer. Bug?",
+    );
+
+    // Iterate over pixels and copy them into `buf`.
+    for (i, pixel) in image.chunks_exact(2).enumerate() {
+        let &[gray, alpha] = pixel else {
+            unreachable!()
+        };
+
+        // ignore the alpha channel.
+        let _ = alpha;
+
+        // Convert 8-bit colorspace to 4-bits by truncating the 4 least significant bits.
+        // Yes, this means that the most-significant 4 bits of each byte are unused.
+        // Not idea why they designed it this way. Maybe it's faster?
+        buf[i] = gray >> 4;
+    }
+
+    // Sanity check buffer length
+    assert_eq!(buf.len(), IOCTL_BUFFER_SIZE);
+
+    // Sanity-check that last byte is null.
+    assert_eq!(buf.last(), Some(&0u8));
+
+    Ok(buf)
+}
+
+fn parse_stamp(stamp: &str) -> Option<(i64, i64, &Path)> {
+    let (x, stamp) = stamp.split_once(':')?;
+    let (y, path) = stamp.split_once(':')?;
+
+    let x = x.parse().ok()?;
+    let y = y.parse().ok()?;
+    let path = Path::new(path);
+    Some((x, y, path))
+}
+
+fn apply_stamp(stamp: &str, onto: &mut GrayAlphaImage) -> eyre::Result<()> {
+    let (x, y, path) =
+        parse_stamp(stamp).wrap_err_with(|| eyre!("Invalid stamp format: {stamp:?}"))?;
+
+    let image = load_stamp(path)?;
+
+    overlay(onto, &image, x, y);
+
+    Ok(())
+}
+
+fn load_stamp(path: &Path) -> eyre::Result<GrayAlphaImage> {
+    eprintln!("Opening image at {path:?}");
+
+    let image = ImageReader::open(path)?.decode()?;
+
+    Ok(to_grayscale(path, image))
+}