rawshift-image 0.1.1

//! Safe wrapper around libjxl's C encode API for JPEG XL encoding.
//!
//! All `unsafe`/FFI interaction with libjxl is confined to this module (the
//! `src/codecs` safety boundary — see `PRINCIPLES.md`). Callers receive a plain
//! `Vec<u8>` codestream and `Result<_, String>`; `src/formats/encode.rs` maps the
//! `String` into [`EncodeError::Jxl`](crate::error::EncodeError::Jxl).
//!
//! The raw bindings are generated by `build.rs` (bindgen over libjxl's BSD-3
//! `jxl/encode.h`) into `$OUT_DIR/jxl_bindings.rs` and `include!`d below — we do
//! not vendor anyone else's bindings, keeping the crate free of GPL code.

use std::os::raw::c_void;
use std::ptr;

/// Generated libjxl C-API bindings (see `build.rs`). bindgen output, so the
/// usual non-idiomatic-naming and clippy lints are silenced wholesale.
#[allow(
    non_upper_case_globals,
    non_camel_case_types,
    non_snake_case,
    dead_code,
    improper_ctypes,
    clippy::all,
    clippy::pedantic
)]
mod ffi {
    include!(concat!(env!("OUT_DIR"), "/jxl_bindings.rs"));
}

/// Fully-resolved libjxl encoder settings.
///
/// This is the plain-data contract between [`crate::formats::encode`] and the
/// FFI layer: every field is already the integer/float libjxl expects. The
/// integer toggles use libjxl's own sentinel `-1` for "leave at the encoder
/// default"; such values are simply not applied.
#[derive(Debug, Clone)]
pub struct JxlEncodeParams {
    /// Butteraugli distance (`0.0` = lossless, up to `25.0`). Ignored when
    /// `quality` is `Some` or `lossless` is set.
    pub distance: f32,
    /// JPEG-style quality `0..=100`; when `Some`, converted to a distance by
    /// libjxl and used instead of `distance`.
    pub quality: Option<f32>,
    /// Force mathematically lossless encoding.
    pub lossless: bool,
    /// `JXL_ENC_FRAME_SETTING_EFFORT`, `1..=10`.
    pub effort: i64,
    /// `JXL_ENC_FRAME_SETTING_BROTLI_EFFORT`, `-1` (default) or `0..=11`.
    pub brotli_effort: i64,
    /// `JXL_ENC_FRAME_SETTING_DECODING_SPEED`, `0..=4`.
    pub decoding_speed: i64,
    /// Emit a responsive/progressive stream (sets `RESPONSIVE` + `PROGRESSIVE_DC`).
    pub progressive: bool,
    /// `JXL_ENC_FRAME_SETTING_MODULAR`, `-1`/`0`/`1`.
    pub modular: i64,
    /// `JXL_ENC_FRAME_SETTING_COLOR_TRANSFORM`, `-1`/`0`(XYB)/`1`(None)/`2`(YCbCr).
    pub color_transform: i64,
    /// `JXL_ENC_FRAME_SETTING_EPF`, `-1..=3`.
    pub epf: i64,
    /// `JXL_ENC_FRAME_SETTING_GABORISH`, `-1`/`0`/`1`.
    pub gaborish: i64,
    /// `JXL_ENC_FRAME_SETTING_NOISE`, `-1`/`0`/`1`.
    pub noise: i64,
    /// `JXL_ENC_FRAME_SETTING_DOTS`, `-1`/`0`/`1`.
    pub dots: i64,
    /// `JXL_ENC_FRAME_SETTING_PATCHES`, `-1`/`0`/`1`.
    pub patches: i64,
    /// `JXL_ENC_FRAME_SETTING_PHOTON_NOISE` (ISO; `0.0` = off).
    pub photon_noise_iso: f32,
    /// `JXL_ENC_FRAME_SETTING_RESAMPLING`, `-1`/`1`/`2`/`4`/`8`.
    pub resampling: i64,
    /// Force the BMFF container (`JxlEncoderUseContainer`).
    pub use_container: bool,
    /// `JxlEncoderSetCodestreamLevel`: `-1` (auto), `5`, or `10`.
    pub codestream_level: i32,
    /// Raw integer frame settings passthrough: `(JxlEncoderFrameSettingId, value)`.
    pub extra_int_options: Vec<(i32, i64)>,
    /// Raw float frame settings passthrough: `(JxlEncoderFrameSettingId, value)`.
    pub extra_float_options: Vec<(i32, f32)>,
}

impl Default for JxlEncodeParams {
    fn default() -> Self {
        Self {
            distance: 1.0,
            quality: None,
            lossless: false,
            effort: 7,
            brotli_effort: -1,
            decoding_speed: 0,
            progressive: false,
            modular: -1,
            color_transform: -1,
            epf: -1,
            gaborish: -1,
            noise: -1,
            dots: -1,
            patches: -1,
            photon_noise_iso: 0.0,
            resampling: -1,
            use_container: false,
            codestream_level: -1,
            extra_int_options: Vec::new(),
            extra_float_options: Vec::new(),
        }
    }
}

/// Owns a `JxlEncoder` and frees it on drop (including on the early-return paths).
struct Encoder(*mut ffi::JxlEncoder);
impl Drop for Encoder {
    fn drop(&mut self) {
        unsafe { ffi::JxlEncoderDestroy(self.0) }
    }
}

/// Owns a resizable parallel runner and frees it on drop.
struct Runner(*mut c_void);
impl Drop for Runner {
    fn drop(&mut self) {
        unsafe { ffi::JxlResizableParallelRunnerDestroy(self.0) }
    }
}

/// Encode interleaved RGB samples to a JPEG XL codestream.
///
/// `samples` is packed RGB with no alpha: `width * height * 3` bytes for
/// `bits_per_sample == 8`, or `width * height * 3 * 2` native-endian bytes for
/// `bits_per_sample == 16`. The image is encoded as sRGB (matching how this crate
/// tags decoded pixels); EXIF/ICC/XMP boxes are appended by the caller.
pub fn encode(
    samples: &[u8],
    width: u32,
    height: u32,
    bits_per_sample: u32,
    params: &JxlEncodeParams,
) -> Result<Vec<u8>, String> {
    let bytes_per_sample = if bits_per_sample == 16 { 2 } else { 1 };
    let expected = (width as usize) * (height as usize) * 3 * bytes_per_sample;
    if samples.len() != expected {
        return Err(format!(
            "pixel buffer length mismatch: expected {expected}, got {}",
            samples.len()
        ));
    }

    unsafe {
        let enc = Encoder(ffi::JxlEncoderCreate(ptr::null()));
        if enc.0.is_null() {
            return Err("JxlEncoderCreate returned null".into());
        }

        // Multithread the encode, sized to the image. The runner must outlive the
        // encoder use below, so it is dropped after `enc` (declared after it).
        let runner = Runner(ffi::JxlResizableParallelRunnerCreate(ptr::null()));
        if !runner.0.is_null() {
            let threads =
                ffi::JxlResizableParallelRunnerSuggestThreads(u64::from(width), u64::from(height));
            ffi::JxlResizableParallelRunnerSetThreads(runner.0, threads as usize);
            if ffi::JxlEncoderSetParallelRunner(
                enc.0,
                Some(ffi::JxlResizableParallelRunner),
                runner.0,
            ) != ffi::JXL_ENC_SUCCESS
            {
                return Err("JxlEncoderSetParallelRunner failed".into());
            }
        }

        if params.codestream_level >= 0
            && ffi::JxlEncoderSetCodestreamLevel(enc.0, params.codestream_level)
                != ffi::JXL_ENC_SUCCESS
        {
            return Err("JxlEncoderSetCodestreamLevel failed".into());
        }
        if params.use_container && ffi::JxlEncoderUseContainer(enc.0, 1) != ffi::JXL_ENC_SUCCESS {
            return Err("JxlEncoderUseContainer failed".into());
        }

        // Basic info: dimensions, depth, RGB, no alpha. `uses_original_profile` is
        // required for bit-exact lossless and harmless (slightly larger) otherwise.
        let mut info: ffi::JxlBasicInfo = std::mem::zeroed();
        ffi::JxlEncoderInitBasicInfo(&mut info);
        info.xsize = width;
        info.ysize = height;
        info.bits_per_sample = bits_per_sample;
        info.exponent_bits_per_sample = 0;
        info.num_color_channels = 3;
        info.num_extra_channels = 0;
        info.alpha_bits = 0;
        info.uses_original_profile =
            i32::from(params.lossless || params.distance == 0.0 || params.quality == Some(100.0));
        if ffi::JxlEncoderSetBasicInfo(enc.0, &info) != ffi::JXL_ENC_SUCCESS {
            return Err("JxlEncoderSetBasicInfo failed (unsupported dimensions or depth)".into());
        }

        // Color: sRGB, matching the crate's display-referred convention.
        let mut color: ffi::JxlColorEncoding = std::mem::zeroed();
        ffi::JxlColorEncodingSetToSRGB(&mut color, 0);
        if ffi::JxlEncoderSetColorEncoding(enc.0, &color) != ffi::JXL_ENC_SUCCESS {
            return Err("JxlEncoderSetColorEncoding failed".into());
        }

        let fs = ffi::JxlEncoderFrameSettingsCreate(enc.0, ptr::null());
        if fs.is_null() {
            return Err("JxlEncoderFrameSettingsCreate returned null".into());
        }

        // Integer frame settings. These are macros (not closures) so each FFI
        // call expands directly into this fn's `unsafe` scope — a closure body
        // would form its own scope and need its own `unsafe`. `-1` is libjxl's
        // "leave at default" sentinel, which `set_int_opt!` skips.
        macro_rules! set_int {
            ($id:expr, $v:expr) => {{
                let value: i64 = $v;
                if ffi::JxlEncoderFrameSettingsSetOption(fs, $id, value) != ffi::JXL_ENC_SUCCESS {
                    return Err(format!(
                        "JxlEncoderFrameSettingsSetOption(id={}, v={value}) failed",
                        $id as i64
                    ));
                }
            }};
        }
        macro_rules! set_int_opt {
            ($id:expr, $v:expr) => {{
                if $v != -1 {
                    set_int!($id, $v);
                }
            }};
        }

        set_int!(ffi::JXL_ENC_FRAME_SETTING_EFFORT, params.effort);
        set_int!(
            ffi::JXL_ENC_FRAME_SETTING_DECODING_SPEED,
            params.decoding_speed
        );
        set_int_opt!(
            ffi::JXL_ENC_FRAME_SETTING_BROTLI_EFFORT,
            params.brotli_effort
        );
        set_int_opt!(ffi::JXL_ENC_FRAME_SETTING_MODULAR, params.modular);
        set_int_opt!(
            ffi::JXL_ENC_FRAME_SETTING_COLOR_TRANSFORM,
            params.color_transform
        );
        set_int_opt!(ffi::JXL_ENC_FRAME_SETTING_EPF, params.epf);
        set_int_opt!(ffi::JXL_ENC_FRAME_SETTING_GABORISH, params.gaborish);
        set_int_opt!(ffi::JXL_ENC_FRAME_SETTING_NOISE, params.noise);
        set_int_opt!(ffi::JXL_ENC_FRAME_SETTING_DOTS, params.dots);
        set_int_opt!(ffi::JXL_ENC_FRAME_SETTING_PATCHES, params.patches);
        set_int_opt!(ffi::JXL_ENC_FRAME_SETTING_RESAMPLING, params.resampling);
        if params.progressive {
            set_int!(ffi::JXL_ENC_FRAME_SETTING_RESPONSIVE, 1);
            set_int!(ffi::JXL_ENC_FRAME_SETTING_PROGRESSIVE_DC, 1);
        }
        if params.photon_noise_iso > 0.0
            && ffi::JxlEncoderFrameSettingsSetFloatOption(
                fs,
                ffi::JXL_ENC_FRAME_SETTING_PHOTON_NOISE,
                params.photon_noise_iso,
            ) != ffi::JXL_ENC_SUCCESS
        {
            return Err("JxlEncoderFrameSettingsSetFloatOption(PHOTON_NOISE) failed".into());
        }

        // Raw escape hatches for any frame setting not surfaced above.
        for &(id, v) in &params.extra_int_options {
            set_int!(id as ffi::JxlEncoderFrameSettingId, v);
        }
        for &(id, v) in &params.extra_float_options {
            if ffi::JxlEncoderFrameSettingsSetFloatOption(
                fs,
                id as ffi::JxlEncoderFrameSettingId,
                v,
            ) != ffi::JXL_ENC_SUCCESS
            {
                return Err(format!(
                    "JxlEncoderFrameSettingsSetFloatOption(id={id}) failed"
                ));
            }
        }

        // Quality/lossless. `SetFrameLossless` overrides distance-related options.
        if params.lossless {
            if ffi::JxlEncoderSetFrameLossless(fs, 1) != ffi::JXL_ENC_SUCCESS {
                return Err("JxlEncoderSetFrameLossless failed".into());
            }
        } else {
            let distance = match params.quality {
                Some(q) => ffi::JxlEncoderDistanceFromQuality(q),
                None => params.distance,
            };
            if ffi::JxlEncoderSetFrameDistance(fs, distance) != ffi::JXL_ENC_SUCCESS {
                return Err("JxlEncoderSetFrameDistance failed (distance out of range)".into());
            }
        }

        // Feed the single frame.
        let pixel_format = ffi::JxlPixelFormat {
            num_channels: 3,
            data_type: if bits_per_sample == 16 {
                ffi::JXL_TYPE_UINT16
            } else {
                ffi::JXL_TYPE_UINT8
            },
            endianness: ffi::JXL_NATIVE_ENDIAN,
            align: 0,
        };
        if ffi::JxlEncoderAddImageFrame(
            fs,
            &pixel_format,
            samples.as_ptr().cast::<c_void>(),
            samples.len(),
        ) != ffi::JXL_ENC_SUCCESS
        {
            return Err("JxlEncoderAddImageFrame failed".into());
        }
        ffi::JxlEncoderCloseInput(enc.0);

        // Pump the encoder until it reports success.
        let mut out: Vec<u8> = Vec::new();
        let mut chunk = vec![0u8; 1 << 16];
        loop {
            let mut next_out = chunk.as_mut_ptr();
            let mut avail_out = chunk.len();
            let status = ffi::JxlEncoderProcessOutput(enc.0, &mut next_out, &mut avail_out);
            let produced = chunk.len() - avail_out;
            out.extend_from_slice(&chunk[..produced]);
            match status {
                ffi::JXL_ENC_SUCCESS => break,
                ffi::JXL_ENC_NEED_MORE_OUTPUT => continue,
                _ => return Err("JxlEncoderProcessOutput failed".into()),
            }
        }

        Ok(out)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    /// 4x4 RGB test pattern at the requested depth (8 or 16 bits), as bytes.
    fn pattern(bits: u32) -> Vec<u8> {
        let px16: [[u16; 3]; 16] = std::array::from_fn(|i| {
            let v = (i as u16) * 4096;
            [v, v.wrapping_add(20000), v.wrapping_add(40000)]
        });
        let mut out = Vec::new();
        for px in px16 {
            for s in px {
                if bits == 16 {
                    out.extend_from_slice(&s.to_ne_bytes());
                } else {
                    out.push((s >> 8) as u8);
                }
            }
        }
        out
    }

    /// Decode a JXL codestream back to interleaved 16-bit RGB via jxl-oxide.
    #[cfg(feature = "jxl-decode")]
    fn decode(data: &[u8]) -> (u32, u32, Vec<u16>) {
        use jxl_oxide::JxlImage;
        let image = JxlImage::builder()
            .read(std::io::Cursor::new(data))
            .expect("decode header");
        let (w, h) = (image.width(), image.height());
        let render = image.render_frame(0).expect("render");
        let mut stream = render.stream_no_alpha();
        let ch = stream.channels() as usize;
        let mut buf = vec![0u16; (w * h) as usize * ch];
        stream.write_to_buffer(&mut buf);
        (w, h, buf)
    }

    #[test]
    #[cfg_attr(miri, ignore)] // FFI into libjxl cannot run under Miri.
    fn lossy_8bit_roundtrips() {
        let data = encode(&pattern(8), 4, 4, 8, &JxlEncodeParams::default()).expect("encode");
        assert!(!data.is_empty());
        // Bare codestream (FF 0A) or boxed container.
        assert!(data.starts_with(&[0xFF, 0x0A]) || data.len() > 12);
        #[cfg(feature = "jxl-decode")]
        {
            let (w, h, px) = decode(&data);
            assert_eq!((w, h), (4, 4));
            assert_eq!(px.len(), 4 * 4 * 3);
        }
    }

    #[test]
    #[cfg_attr(miri, ignore)]
    fn lossless_16bit_is_exact() {
        let src = pattern(16);
        let params = JxlEncodeParams {
            distance: 0.0,
            lossless: true,
            ..JxlEncodeParams::default()
        };
        let data = encode(&src, 4, 4, 16, &params).expect("encode");
        #[cfg(feature = "jxl-decode")]
        {
            let (w, h, px) = decode(&data);
            assert_eq!((w, h), (4, 4));
            // Reconstruct the source u16 samples (native-endian) and compare exactly.
            let want: Vec<u16> = src
                .chunks_exact(2)
                .map(|b| u16::from_ne_bytes([b[0], b[1]]))
                .collect();
            assert_eq!(px, want, "lossless 16-bit round-trip must be exact");
        }
    }

    #[test]
    #[cfg_attr(miri, ignore)]
    fn toggles_and_escape_hatch_apply() {
        let params = JxlEncodeParams {
            distance: 2.0,
            effort: 3,
            modular: 1, // force modular
            progressive: true,
            // JXL_ENC_FRAME_SETTING_KEEP_INVISIBLE (a setting not surfaced as a field).
            extra_int_options: vec![(ffi::JXL_ENC_FRAME_SETTING_KEEP_INVISIBLE as i32, 0)],
            ..JxlEncodeParams::default()
        };
        let data = encode(&pattern(8), 4, 4, 8, &params).expect("encode with toggles");
        assert!(!data.is_empty());
    }

    #[test]
    #[cfg_attr(miri, ignore)]
    fn rejects_wrong_buffer_length() {
        let err = encode(&[0u8; 10], 4, 4, 8, &JxlEncodeParams::default());
        assert!(err.is_err());
    }
}