heifer 0.3.6

//! Integration with the [`image`](https://docs.rs/image) crate (feature `image`).
//!
//! ```no_run
//! heifer::image_crate::register_image_decoder_hooks();
//! let img = image::open("photo.heic").unwrap(); // decoded by heifer
//! ```
//!
//! [`HeifDecoder`] can also be used directly with `image::DynamicImage::from_decoder`.
//!
//! [`HeifEncoder`] writes HEIC files from any image of the `image` crate:
//!
//! ```no_run
//! let img = image::open("photo.png").unwrap();
//! let file = std::fs::File::create("photo.heic").unwrap();
//! img.write_with_encoder(heifer::image_crate::HeifEncoder::new(file)).unwrap();
//! ```
//!
//! (`image::save` and `DynamicImage::save` only know the formats built into `image`.)

use std::io::{Read, Write};

use image::error::{
    DecodingError, EncodingError, ImageFormatHint, UnsupportedError, UnsupportedErrorKind,
};
use image::metadata::Orientation;
use image::{ColorType, ExtendedColorType, ImageDecoder, ImageEncoder, ImageError, ImageResult};

use crate::{EncodeOptions, Image, Metadata, Options};

/// File extensions handled by heifer.
const EXTENSIONS: [&str; 3] = ["heic", "heif", "hif"];

/// HEVC-based HEIF brands recognised from the file content (the `ftyp` major brand).
const BRANDS: [&[u8; 4]; 4] = [b"heic", b"heix", b"heim", b"heis"];

/// Registers heifer with the `image` crate, for the `.heic`, `.heif` and `.hif` extensions and
/// for content detection of HEVC-coded HEIF files. Returns `false` if another decoder was
/// already registered for one of these extensions.
///
/// After this call, `image::open`, `image::ImageReader` and `image::load_from_memory` decode
/// HEIC files with heifer.
pub fn register_image_decoder_hooks() -> bool {
    let mut all = true;
    for ext in EXTENSIONS {
        all &= image::hooks::register_decoding_hook(
            ext.into(),
            Box::new(|reader| Ok(Box::new(HeifDecoder::new(reader)?))),
        );
    }
    // Box size (4 bytes, any value) then `ftyp` and the major brand.
    const MASK: &[u8] = b"\x00\x00\x00\x00\xff\xff\xff\xff\xff\xff\xff\xff";
    static SIGNATURES: [[u8; 12]; 4] = {
        let mut s = [[0u8; 12]; 4];
        let mut i = 0;
        while i < 4 {
            let mut j = 0;
            while j < 4 {
                s[i][4 + j] = b"ftyp"[j];
                s[i][8 + j] = BRANDS[i][j];
                j += 1;
            }
            i += 1;
        }
        s
    };
    for signature in &SIGNATURES {
        image::hooks::register_format_detection_hook("heic".into(), signature, Some(MASK));
    }
    all
}

/// An [`image::ImageDecoder`] for HEIF/HEIC files.
///
/// The image is decoded when the decoder is created; crop, rotation and mirror transforms are
/// already applied, so [`ImageDecoder::orientation`] reports no remaining transform.
pub struct HeifDecoder {
    image: Image,
    metadata: Metadata,
}

impl HeifDecoder {
    /// Reads a whole HEIF file and decodes its primary image.
    pub fn new(mut reader: impl Read) -> ImageResult<Self> {
        let mut bytes = Vec::new();
        reader
            .read_to_end(&mut bytes)
            .map_err(ImageError::IoError)?;
        Self::from_bytes(&bytes, &Options::default())
    }

    /// Decodes a HEIF file held in memory, with heifer options.
    pub fn from_bytes(bytes: &[u8], options: &Options) -> ImageResult<Self> {
        let image = crate::decode_with_options(bytes, options).map_err(decoding_error)?;
        let metadata = crate::read_metadata(bytes).map_err(decoding_error)?;
        Ok(Self { image, metadata })
    }
}

fn decoding_error(e: crate::Error) -> ImageError {
    ImageError::Decoding(DecodingError::new(ImageFormatHint::Name("HEIF".into()), e))
}

impl ImageDecoder for HeifDecoder {
    fn dimensions(&self) -> (u32, u32) {
        (self.image.width, self.image.height)
    }

    fn color_type(&self) -> ColorType {
        match (self.image.bit_depth > 8, self.image.has_alpha) {
            (false, false) => ColorType::Rgb8,
            (false, true) => ColorType::Rgba8,
            (true, false) => ColorType::Rgb16,
            (true, true) => ColorType::Rgba16,
        }
    }

    fn icc_profile(&mut self) -> ImageResult<Option<Vec<u8>>> {
        Ok(self.metadata.icc_profile.clone())
    }

    fn exif_metadata(&mut self) -> ImageResult<Option<Vec<u8>>> {
        Ok(self.metadata.exif.clone())
    }

    fn xmp_metadata(&mut self) -> ImageResult<Option<Vec<u8>>> {
        Ok(self.metadata.xmp.clone())
    }

    /// HEIF transforms (`irot`/`imir`) are already applied, and the EXIF orientation is only
    /// informative in HEIF: nothing is left to do.
    fn orientation(&mut self) -> ImageResult<Orientation> {
        Ok(Orientation::NoTransforms)
    }

    fn set_limits(&mut self, limits: image::Limits) -> ImageResult<()> {
        limits.check_support(&image::LimitSupport::default())?;
        limits.check_dimensions(self.image.width, self.image.height)?;
        Ok(())
    }

    fn read_image(self, buf: &mut [u8]) -> ImageResult<()> {
        let channels = if self.image.has_alpha { 4 } else { 3 };
        let max = u32::from(self.image.max_value());
        let pixels = self.image.data.chunks_exact(4).map(|p| &p[..channels]);
        if self.image.bit_depth > 8 {
            // 16-bit native endian, scaled to the full 0..=65535 range.
            for (out, v) in buf.chunks_exact_mut(2).zip(pixels.flatten()) {
                let v16 = ((u32::from(*v) * 65535 + max / 2) / max) as u16;
                out.copy_from_slice(&v16.to_ne_bytes());
            }
        } else {
            for (out, v) in buf.iter_mut().zip(pixels.flatten()) {
                *out = *v as u8;
            }
        }
        Ok(())
    }

    fn read_image_boxed(self: Box<Self>, buf: &mut [u8]) -> ImageResult<()> {
        (*self).read_image(buf)
    }
}

/// Encodes images as HEIC files, for `DynamicImage::write_with_encoder` and the other users of
/// [`ImageEncoder`].
///
/// Accepts 8- and 16-bit grey, grey + alpha, RGB and RGBA images (16-bit images are coded in
/// 10 bits). The ICC profile and EXIF metadata given by `image` are stored in the file; the EXIF
/// orientation is set to 1, since the pixels are stored as given.
#[derive(Debug)]
pub struct HeifEncoder<W: Write> {
    writer: W,
    options: EncodeOptions,
    metadata: Metadata,
}

impl<W: Write> HeifEncoder<W> {
    /// An encoder with the default options (quality 80, all cores).
    pub fn new(writer: W) -> Self {
        Self::with_options(writer, EncodeOptions::default())
    }

    /// An encoder with the given quality and threads.
    pub fn with_options(writer: W, options: EncodeOptions) -> Self {
        Self {
            writer,
            options,
            metadata: Metadata::default(),
        }
    }
}

fn hint() -> ImageFormatHint {
    ImageFormatHint::Name("HEIF".into())
}

impl<W: Write> ImageEncoder for HeifEncoder<W> {
    fn write_image(
        mut self,
        buf: &[u8],
        width: u32,
        height: u32,
        color_type: ExtendedColorType,
    ) -> ImageResult<()> {
        use ExtendedColorType as C;
        let (channels, wide) = match color_type {
            C::L8 => (1, false),
            C::La8 => (2, false),
            C::Rgb8 => (3, false),
            C::Rgba8 => (4, false),
            C::L16 => (1, true),
            C::La16 => (2, true),
            C::Rgb16 => (3, true),
            C::Rgba16 => (4, true),
            _ => {
                return Err(ImageError::Unsupported(
                    UnsupportedError::from_format_and_kind(
                        hint(),
                        UnsupportedErrorKind::Color(color_type),
                    ),
                ));
            }
        };
        let bytes = if wide { 2 } else { 1 };
        let expected = u64::from(width) * u64::from(height) * channels * bytes;
        if buf.len() as u64 != expected {
            return Err(ImageError::Parameter(
                image::error::ParameterError::from_kind(
                    image::error::ParameterErrorKind::DimensionMismatch,
                ),
            ));
        }
        let channels = channels as usize;
        // 16-bit images are coded in 10 bits (no banding), 8-bit ones in 8 bits.
        let (bit_depth, top) = if wide { (10, 1023u16) } else { (8, 255) };
        let sample = |i: usize| -> u16 {
            if wide {
                let v = u16::from_ne_bytes([buf[2 * i], buf[2 * i + 1]]);
                ((u32::from(v) * 1023 + 32767) / 65535) as u16
            } else {
                u16::from(buf[i])
            }
        };
        let pixel = |i: usize| -> [u16; 4] {
            let p = i * channels;
            match channels {
                1 => {
                    let v = sample(p);
                    [v, v, v, top]
                }
                2 => {
                    let v = sample(p);
                    [v, v, v, sample(p + 1)]
                }
                3 => [sample(p), sample(p + 1), sample(p + 2), top],
                _ => [sample(p), sample(p + 1), sample(p + 2), sample(p + 3)],
            }
        };
        let n = width as usize * height as usize;
        let has_alpha = channels % 2 == 0 && (0..n).any(|i| pixel(i)[3] != top);
        let heic = crate::encode::encode_pixels(
            width,
            height,
            bit_depth,
            has_alpha,
            Some(&self.metadata),
            &self.options,
            pixel,
        )
        .map_err(|e| ImageError::Encoding(EncodingError::new(hint(), e)))?;
        self.writer.write_all(&heic).map_err(ImageError::IoError)
    }

    fn set_icc_profile(&mut self, icc_profile: Vec<u8>) -> Result<(), UnsupportedError> {
        self.metadata.icc_profile = Some(icc_profile);
        Ok(())
    }

    fn set_exif_metadata(&mut self, exif: Vec<u8>) -> Result<(), UnsupportedError> {
        // Some callers keep the `Exif\0\0` prefix of JPEG APP1 segments.
        let tiff = exif.strip_prefix(b"Exif\0\0").unwrap_or(&exif);
        self.metadata.exif = Some(tiff.to_vec());
        Ok(())
    }
}