blurhash-update 0.1.0

#![doc = include_str!("../README.md")]
#![deny(missing_docs)]

mod base83;
mod srgb_lookup;

use std::f32::consts::PI;

use srgb_lookup::srgb_to_linear;

const BYTES_PER_PIXEL: usize = 4;

/// How many components should be used in blurhash creation
///
/// More components will increase the definition of the blurhash, but also increase processing
/// time.
pub struct Components {
    /// How many components to process in the x direction
    pub x: u32,

    /// How many components to process in the y direction
    pub y: u32,
}

/// Bounds for the input image
///
/// These are required since they can not be inferred from the RGBA values
#[derive(Clone, Copy, Debug)]
pub struct ImageBounds {
    /// The input image's width
    pub width: u32,

    /// The input image's height
    pub height: u32,
}

struct ComponentState {
    x: u32,
    y: u32,
    basis: f32,
}

/// Error raised when too many components are requested
#[derive(Debug)]
pub enum ConfigurationError {
    /// Component values are not within the required range.
    InvalidComponentCount,

    /// Skip value must not be zero
    ZeroSkip,
}

/// Encoder type used to produce blurhashes
pub struct Encoder {
    index: usize,
    skip: u32,
    components: Components,
    factors: Box<[(ComponentState, [f32; BYTES_PER_PIXEL])]>,
    bounds: ImageBounds,
}

/// A simple "encode this image please" function
///
/// The input image must be in the sRGB colorspace and be formatted as 8bit RGBA
pub fn encode(
    components: Components,
    bounds: ImageBounds,
    rgba8_image: &[u8],
) -> Result<String, ConfigurationError> {
    let mut encoder = Encoder::new(components, bounds, 1)?;
    encoder.update(rgba8_image);
    Ok(encoder.finalize())
}

/// A simple "encode this image please" function that automatically selects component and skip
/// values
///
/// The input image must be in the sRGB colorspace and be formatted as 8bit RGBA
pub fn auto_encode(bounds: ImageBounds, rgba8_image: &[u8]) -> String {
    let mut encoder = Encoder::auto(bounds);
    encoder.update(rgba8_image);
    encoder.finalize()
}

// determine closest component ratio to input bounds
fn calculate_components(ImageBounds { width, height }: ImageBounds) -> Components {
    let mut out = Components { x: 0, y: 0 };

    let (out_longer, out_shorter, in_longer, in_shorter) = if width > height {
        (&mut out.x, &mut out.y, width as f32, height as f32)
    } else {
        (&mut out.y, &mut out.x, height as f32, width as f32)
    };

    struct State {
        similarity: f32,
        ratio: (u32, u32),
    }

    let ratios = [(3, 3), (4, 3), (5, 3), (6, 3), (5, 2), (6, 2), (7, 2)];

    let in_ratio = in_longer / in_shorter;

    let State { ratio, .. } = ratios.into_iter().fold(
        State {
            similarity: f32::MAX,
            ratio: (0, 0),
        },
        |state, (ratio_longer, ratio_shorter)| {
            let ratio = ratio_longer as f32 / ratio_shorter as f32;
            let diff = (ratio - in_ratio).abs();

            if diff < state.similarity {
                State {
                    similarity: diff,
                    ratio: (ratio_longer, ratio_shorter),
                }
            } else {
                state
            }
        },
    );

    *out_longer = ratio.0;
    *out_shorter = ratio.1;

    out
}

// target 256ish total pixels to process
fn calculate_skip(ImageBounds { width, height }: ImageBounds) -> u32 {
    let target_1d = f32::sqrt((width * height / 512) as f32).floor() as u32;

    let mut base = 1;

    loop {
        if base * 2 < target_1d {
            base *= 2;
        } else {
            break base;
        }
    }
}

impl Encoder {
    /// Create an encoder that automatically picks Compoent and Skip values
    ///
    /// This is a best-effort configuration
    pub fn auto(bounds: ImageBounds) -> Self {
        Self::new(calculate_components(bounds), bounds, calculate_skip(bounds))
            .expect("Generated bounds are always valid")
    }

    /// Create a new Encoder to produce a blurhash
    ///
    /// The provided component x and y values must be between 1 and 9 inclusive.
    ///
    /// The `skip` value indicates how many pixels can be skipped when proccessing the image. this
    /// value will be squared to produce the final skip value. When set to 1, no pixels will be
    /// skipped, when set to 2, one in four pixels will be processed. when set to 3, one in 9
    /// pixels will be processed, etc. This improves performance at the cost of losing accuracy.
    ///
    /// Errors if too many components are requested
    pub fn new(
        Components { x, y }: Components,
        bounds: ImageBounds,
        skip: u32,
    ) -> Result<Self, ConfigurationError> {
        if !(1..=9).contains(&x) || !(1..=9).contains(&y) {
            return Err(ConfigurationError::InvalidComponentCount);
        }

        if skip == 0 {
            return Err(ConfigurationError::ZeroSkip);
        }

        Ok(Self {
            index: 0,
            skip,
            components: Components { x, y },
            factors: Box::from(
                (0..y)
                    .flat_map(|y| {
                        (0..x).map(move |x| (ComponentState { x, y, basis: 0. }, [0., 0., 0., 0.]))
                    })
                    .collect::<Vec<_>>(),
            ),
            bounds,
        })
    }

    /// Update the encoder with bytes from an image
    ///
    /// The input image must be in the sRGB colorspace and be formatted as 8bit RGBA
    /// The input doesn't need to contain whole pixels, the encoder is capable of handling partial
    /// pixels
    pub fn update(&mut self, rgba8_image: &[u8]) {
        if self.skip == 1 {
            self.update_noskip(rgba8_image)
        } else {
            self.update_skip(rgba8_image)
        }
    }

    fn update_skip(&mut self, rgba8_image: &[u8]) {
        let basis_scale_x = PI / self.bounds.width as f32;
        let basis_scale_y = PI / self.bounds.height as f32;

        let mut current_index = self.index;

        loop {
            let (px_x, px_y) = self.next_px(current_index);

            let scale_x = px_x as f32 * basis_scale_x;
            let scale_y = px_y as f32 * basis_scale_y;

            let next_index = (px_y * self.bounds.width + px_x) as usize * BYTES_PER_PIXEL;

            let skip_rgb = current_index.saturating_sub(next_index);
            let index_into = next_index.saturating_sub(self.index);

            if index_into >= rgba8_image.len() {
                break;
            }

            assert!(skip_rgb < BYTES_PER_PIXEL, "{skip_rgb}");

            for (ComponentState { x, y, basis }, rgb) in self.factors.iter_mut() {
                *basis = f32::cos(*x as f32 * scale_x) * f32::cos(*y as f32 * scale_y);

                let slot_iter = rgb.iter_mut().skip(skip_rgb);
                let value_iter = rgba8_image[index_into..]
                    .iter()
                    .take(BYTES_PER_PIXEL)
                    .map(|byte| *basis * srgb_to_linear(*byte));

                for (val, slot) in value_iter.zip(slot_iter) {
                    *slot += val;
                }
            }

            current_index = next_index + BYTES_PER_PIXEL;
        }

        self.index += rgba8_image.len();
    }

    fn next_px(&self, index: usize) -> (u32, u32) {
        let pixel = (index / BYTES_PER_PIXEL) as u32;
        let pixel_x = pixel % self.bounds.width;
        let pixel_y = pixel / self.bounds.width;

        let y_offset = pixel_y % self.skip;

        if y_offset == 0 {
            let x_offset = pixel_x % self.skip;

            if x_offset == 0 {
                (pixel_x, pixel_y)
            } else {
                let next_px_x = pixel_x + self.skip - x_offset;

                if next_px_x >= self.bounds.width {
                    (0, pixel_y + self.skip)
                } else {
                    (next_px_x, pixel_y)
                }
            }
        } else {
            (0, pixel_y + self.skip - y_offset)
        }
    }

    fn update_noskip(&mut self, rgba8_image: &[u8]) {
        // get offset in terms of already-processed bytes
        let offset = self.index % BYTES_PER_PIXEL;
        // get offset in terms of remaining bytes on head of rgba8_image
        let offset = (BYTES_PER_PIXEL - offset) % BYTES_PER_PIXEL;

        let basis_scale_x = PI / self.bounds.width as f32;
        let basis_scale_y = PI / self.bounds.height as f32;

        for (ComponentState { basis, .. }, [_, g, b, _]) in self.factors.iter_mut() {
            for (val, slot) in rgba8_image[..offset]
                .iter()
                .map(|byte| *basis * srgb_to_linear(*byte))
                .zip(
                    [b, g][..offset.saturating_sub(BYTES_PER_PIXEL - 2)]
                        .iter_mut()
                        .rev(),
                )
            {
                **slot += val;
            }
        }

        let pixels = ((self.index + offset) / BYTES_PER_PIXEL) as u32;

        let mut chunks = rgba8_image[offset..].chunks_exact(BYTES_PER_PIXEL);

        for (i, chunk) in (&mut chunks).enumerate() {
            let px = pixels + i as u32;
            let px_x = px % self.bounds.width;
            let px_y = px / self.bounds.width;

            let scale_x = px_x as f32 * basis_scale_x;
            let scale_y = px_y as f32 * basis_scale_y;

            for (ComponentState { x, y, .. }, rgb) in self.factors.iter_mut() {
                let basis = f32::cos(*x as f32 * scale_x) * f32::cos(*y as f32 * scale_y);

                assert_eq!(chunk.len(), rgb.len());
                for (val, slot) in chunk
                    .iter()
                    .map(|byte| basis * srgb_to_linear(*byte))
                    .zip(rgb)
                {
                    *slot += val;
                }
            }
        }

        if !chunks.remainder().is_empty() {
            let px = pixels + (rgba8_image[offset..].len() / BYTES_PER_PIXEL) as u32;
            let px_x = px % self.bounds.width;
            let px_y = px / self.bounds.width;

            let scale_x = px_x as f32 * basis_scale_x;
            let scale_y = px_y as f32 * basis_scale_y;

            for (ComponentState { x, y, basis }, rgb) in self.factors.iter_mut() {
                *basis = f32::cos(*x as f32 * scale_x) * f32::cos(*y as f32 * scale_y);

                for (val, slot) in chunks
                    .remainder()
                    .iter()
                    .map(|byte| *basis * srgb_to_linear(*byte))
                    .zip(rgb)
                {
                    *slot += val;
                }
            }
        }

        self.index += rgba8_image.len();
    }

    /// Produce a blurhash from the provided encoder
    pub fn finalize(mut self) -> String {
        for (ComponentState { x, y, .. }, rgb) in self.factors.iter_mut() {
            let normalisation = if *x == 0 && *y == 0 { 1. } else { 2. };

            let scale = self.skip.pow(2) as f32 * normalisation
                / (self.bounds.width * self.bounds.height) as f32;

            for slot in rgb {
                *slot *= scale;
            }
        }

        let mut blurhash = String::with_capacity(30);

        let (_, dc) = self.factors[0];
        let ac = &self.factors[1..];

        let size_flag = self.components.x - 1 + (self.components.y - 1) * 9;
        base83::encode(size_flag, 1, &mut blurhash);

        let maximum = ac.iter().fold(0.0_f32, |maximum, (_, [r, g, b, _])| {
            maximum.max(r.abs()).max(g.abs()).max(b.abs())
        });

        let quantized_maximum = (maximum * 166. - 0.5).floor().max(0.) as u32;

        base83::encode(quantized_maximum, 1, &mut blurhash);

        let maximum_value = (quantized_maximum + 1) as f32 / 166.;

        base83::encode(encode_dc(dc), 4, &mut blurhash);

        for (_, rgb) in ac {
            base83::encode(encode_ac(*rgb, maximum_value), 2, &mut blurhash);
        }

        blurhash
    }
}

fn encode_dc(rgb: [f32; BYTES_PER_PIXEL]) -> u32 {
    let [r, g, b, _] = rgb.map(linear_to_srgb);

    (r << 16) + (g << 8) + b
}

fn encode_ac(rgb: [f32; BYTES_PER_PIXEL], maximum_value: f32) -> u32 {
    let [r, g, b, _] = rgb.map(|c| encode_ac_digit(c, maximum_value));

    r * 19 * 19 + g * 19 + b
}

fn encode_ac_digit(d: f32, maximum_value: f32) -> u32 {
    ((sign_pow(d / maximum_value, 0.5) * 9. + 9.5) as i32).clamp(0, 18) as u32
}

fn linear_to_srgb(value: f32) -> u32 {
    let v = f32::max(0., f32::min(1., value));
    if v <= 0.003_130_8 {
        (v * 12.92 * 255. + 0.5).round() as u32
    } else {
        ((1.055 * f32::powf(v, 1. / 2.4) - 0.055) * 255. + 0.5).round() as u32
    }
}

fn sign(n: f32) -> f32 {
    if n < 0. {
        -1.
    } else {
        1.
    }
}

fn sign_pow(val: f32, exp: f32) -> f32 {
    sign(val) * val.abs().powf(exp)
}

impl std::fmt::Display for ConfigurationError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::InvalidComponentCount => write!(f, "Components out of bounds"),
            Self::ZeroSkip => write!(f, "Skip value cannot be zero"),
        }
    }
}

impl std::error::Error for ConfigurationError {}

#[cfg(test)]
mod tests {
    use image::{EncodableLayout, GenericImageView};

    #[test]
    fn contrived() {
        let input = [
            0, 60, 120, 0, 0, 60, 120, 0, 0, 60, 120, 0, 0, 60, 120, 0, 0, 60, 120, 0, 0, 60, 120,
            0, 0, 60, 120, 0, 0, 60, 120, 0, 120, 60, 0, 0, 120, 60, 0, 0, 120, 60, 0, 0, 120, 60,
            0, 0, 120, 60, 0, 0, 120, 60, 0, 0, 120, 60, 0, 0, 120, 60, 0, 0,
        ];
        let width = 4;
        let height = 4;

        let hash = super::encode(
            crate::Components { x: 4, y: 3 },
            crate::ImageBounds { width, height },
            &input,
        )
        .unwrap();

        assert_eq!(hash, "LQ9~?d$,fQ$,G1S%fQS%A{SPfQSP");
    }

    #[test]
    fn one_component() {
        let inputs = [
            ("data/19dd1c444d1c7939.png", "00AQtR"),
            ("data/f73d2ee39133d871.jpg", "00E{R{"),
            ("data/shenzi.png", "0039[D"),
        ];

        for (input, output) in inputs {
            let img = image::open(input).unwrap();
            let (width, height) = img.dimensions();

            let hash = super::encode(
                crate::Components { x: 1, y: 1 },
                crate::ImageBounds { width, height },
                img.to_rgba8().as_bytes(),
            )
            .unwrap();

            assert_eq!(hash, output, "wrong output for {input}");
        }
    }

    #[test]
    fn auto() {
        let inputs = [
            ("data/19dd1c444d1c7939.png", (5, 3), 32),
            ("data/f73d2ee39133d871.jpg", (4, 3), 32),
            ("data/shenzi.png", (3, 3), 16),
        ];

        for (input, expected_components, expected_skip) in inputs {
            let img = image::open(input).unwrap();
            let (width, height) = img.dimensions();

            let components = super::calculate_components(crate::ImageBounds { width, height });
            let skip = super::calculate_skip(crate::ImageBounds { width, height });

            assert_eq!(
                (components.x, components.y),
                expected_components,
                "wrong ratio for {input}"
            );
            assert_eq!(skip, expected_skip, "wrong skip for {input}");
        }
    }

    #[test]
    fn matches_blurhash() {
        let inputs = [
            ("data/19dd1c444d1c7939.png", "L3AQtR2FSz6NrsOCW:ODR*,EE};h"),
            ("data/f73d2ee39133d871.jpg", "LJE{R{Z}V?N#0JR*Rit7^htTfkaI"),
            ("data/shenzi.png", "L239[DQ.91t,rJX9Qns+8zt5.PR6"),
        ];

        for (input, output) in inputs {
            let img = image::open(input).unwrap();
            let (width, height) = img.dimensions();

            let hash = super::encode(
                crate::Components { x: 4, y: 3 },
                crate::ImageBounds { width, height },
                img.to_rgba8().as_bytes(),
            )
            .unwrap();

            assert_eq!(hash, output, "wrong output for {input}");
        }
    }

    #[test]
    fn matches_self_when_split() {
        let inputs = [
            "data/19dd1c444d1c7939.png",
            "data/f73d2ee39133d871.jpg",
            "data/shenzi.png",
        ];

        for input in inputs {
            let img = image::open(input).unwrap();
            let (width, height) = img.dimensions();
            let rgba8_img = img.to_rgba8();
            let bytes = rgba8_img.as_bytes();

            let b1 = super::encode(
                crate::Components { x: 4, y: 3 },
                crate::ImageBounds { width, height },
                bytes,
            )
            .unwrap();

            for chunk_count in 2..20 {
                let mut encoder = super::Encoder::new(
                    crate::Components { x: 4, y: 3 },
                    crate::ImageBounds { width, height },
                    1,
                )
                .unwrap();

                let chunk_size = bytes.len() / chunk_count;

                for chunk in bytes.chunks(chunk_size) {
                    encoder.update(chunk);
                }

                let b2 = encoder.finalize();

                assert_eq!(b1, b2, "wrong hash for {input} with {chunk_count} chunks");
            }
        }
    }
}