DarkRoom/core/dr-decode/src/lib.rs

//! RAW decoding for DarkRoom.
//!
//! Four separate entry points rather than one `decode`, because callers differ
//! sharply in what they need (ARCH §3.2):
//!
//! - **Culling** wants [`embedded_preview`] and nothing else — a ~200 KB read
//!   against a 34 MB file.
//! - **The grid** wants [`metadata`].
//! - **Develop and export** need [`decode`], the only path that touches sensor
//!   data.
//!
//! Fusing them would force a full decode where a header read suffices, which
//! is exactly why Lightroom stalls ~2 s per image during culling.

mod error;
mod preview;

pub use error::DecodeError;
pub use preview::{
    decode_jpeg, extract_embedded_preview, extract_preview, Preview, PreviewSize,
    PREVIEW_PROBE_BYTES,
};

use dr_types::Format;

/// Capture metadata read from a file header.
#[derive(Debug, Clone, Default, PartialEq)]
pub struct Metadata {
    pub make: Option<String>,
    pub model: Option<String>,
    pub lens: Option<String>,
    /// Exposure time in seconds.
    pub shutter: Option<f32>,
    pub aperture: Option<f32>,
    pub iso: Option<u32>,
    pub focal_length: Option<f32>,
    /// Full sensor dimensions, before crop.
    pub width: Option<u32>,
    pub height: Option<u32>,
}

/// Decoded sensor data, before demosaic.
///
/// Deliberately *not* RGB: demosaic is a GPU pipeline stage (ARCH §5.2), so
/// this carries CFA-pattern samples plus what the shader needs to interpret
/// them.
#[derive(Debug, Clone)]
pub struct RawImage {
    pub width: u32,
    pub height: u32,
    /// One sample per photosite, in sensor order.
    pub data: Vec<u16>,
    pub cfa_pattern: CfaPattern,
    pub black_level: [u16; 4],
    pub white_level: u16,
    /// As-shot white balance, as per-channel multipliers.
    pub wb_coeffs: [f32; 4],
    /// Camera-to-XYZ colour matrix (FR-DEV-3e).
    pub color_matrix: Option<[f32; 9]>,
}

/// TRACES: FR-RAW-5
/// The colour filter array layout.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CfaPattern {
    Rggb,
    Bggr,
    Grbg,
    Gbrg,
    /// Fujifilm's 6×6 pattern. Needs a different demosaic entirely
    /// (FR-RAW-5), at roughly 2× the cost of Bayer.
    XTrans,
    Unknown,
}

impl CfaPattern {
    /// Whether this needs the X-Trans demosaic path rather than Bayer.
    pub fn is_xtrans(self) -> bool {
        matches!(self, CfaPattern::XTrans)
    }
}

/// TRACES: FR-RAW-1 | M-9
/// Identify a format from a file header.
///
/// Content-based, not extension-based: an extension is a hint, and a
/// mismatched one should not produce a confusing decode failure downstream.
pub fn probe(header: &[u8]) -> Option<Format> {
    if header.len() < 16 {
        return None;
    }

    // JPEG: SOI marker.
    if header.starts_with(&[0xFF, 0xD8, 0xFF]) {
        return Some(Format::Jpeg);
    }

    // Fujifilm RAF carries an ASCII signature.
    if header.starts_with(b"FUJIFILMCCD-RAW") {
        return Some(Format::Raf);
    }

    // CR3 is ISO-BMFF: a `ftyp` box with a Canon brand.
    if header.len() >= 12 && &header[4..8] == b"ftyp" && &header[8..11] == b"crx" {
        return Some(Format::Cr3);
    }

    // The TIFF-derived formats share a byte-order mark plus magic. CR2 adds
    // its own marker at offset 8; the rest are indistinguishable from the
    // header alone and need the extension to disambiguate.
    let le = header.starts_with(&[0x49, 0x49, 0x2A, 0x00]);
    let be = header.starts_with(&[0x4D, 0x4D, 0x00, 0x2A]);
    if le || be {
        if header.len() >= 11 && &header[8..10] == b"CR" {
            return Some(Format::Cr2);
        }
        // Ambiguous between NEF, ARW, DNG, ORF, RW2 — caller falls back to
        // the extension.
        return None;
    }

    None
}

/// TRACES: FR-CAT-5 | M-12
/// Read capture metadata without decoding sensor data.
pub fn metadata(bytes: &[u8]) -> Result<Metadata, DecodeError> {
    use rawler::rawsource::RawSource;

    let source = RawSource::new_from_slice(bytes);
    let decoder =
        rawler::get_decoder(&source).map_err(|e| DecodeError::Unsupported(e.to_string()))?;
    let md = decoder
        .raw_metadata(&source, &Default::default())
        .map_err(|e| DecodeError::Metadata(e.to_string()))?;

    let exif = &md.exif;
    Ok(Metadata {
        make: Some(md.make.clone()).filter(|s| !s.is_empty()),
        model: Some(md.model.clone()).filter(|s| !s.is_empty()),
        lens: exif.lens_model.clone(),
        shutter: exif.exposure_time.map(|r| r.n as f32 / r.d.max(1) as f32),
        aperture: exif.fnumber.map(|r| r.n as f32 / r.d.max(1) as f32),
        iso: exif.iso_speed_ratings.map(|v| v as u32),
        focal_length: exif.focal_length.map(|r| r.n as f32 / r.d.max(1) as f32),
        width: None,
        height: None,
    })
}

/// TRACES: FR-RAW-3 | FR-EXP-9
/// Fully decode sensor data.
///
/// The expensive path — reads the whole file and unpacks every photosite.
/// Only develop and export should call it; culling and the grid must not
/// (FR-CULL-1).
pub fn decode(bytes: &[u8]) -> Result<RawImage, DecodeError> {
    use rawler::rawsource::RawSource;

    let source = RawSource::new_from_slice(bytes);
    let decoder =
        rawler::get_decoder(&source).map_err(|e| DecodeError::Unsupported(e.to_string()))?;
    let image = decoder
        .raw_image(&source, &Default::default(), false)
        .map_err(|e| DecodeError::Decode(e.to_string()))?;

    let data = match image.data {
        rawler::RawImageData::Integer(v) => v,
        rawler::RawImageData::Float(v) => {
            // Float sensor data is rare; normalise to the u16 the pipeline
            // expects rather than carrying two representations.
            v.iter()
                .map(|&f| (f * 65535.0).clamp(0.0, 65535.0) as u16)
                .collect()
        }
    };

    let cfa = cfa_from_rawler(&image.camera.cfa, image.camera.model.as_str());
    // Black levels are rationals; the pipeline wants plain u16 samples.
    let bl = &image.blacklevel.levels;
    let level_at = |i: usize| -> u16 {
        bl.get(i)
            .map(|r| (r.n as f32 / r.d.max(1) as f32).round() as u16)
            .unwrap_or(0)
    };
    let black_level = [level_at(0), level_at(1), level_at(2), level_at(3)];

    Ok(RawImage {
        width: image.width as u32,
        height: image.height as u32,
        data,
        cfa_pattern: cfa,
        black_level,
        white_level: image
            .whitelevel
            .0
            .first()
            .map(|v| *v as u16)
            .unwrap_or(u16::MAX),
        wb_coeffs: image.wb_coeffs,
        color_matrix: None,
    })
}

fn cfa_from_rawler(cfa: &rawler::CFA, model: &str) -> CfaPattern {
    // rawler exposes the pattern as a string; X-Trans is 6x6 rather than 2x2.
    let name = cfa.name.to_ascii_uppercase();
    if name.len() > 4 || model.contains("X-") {
        return CfaPattern::XTrans;
    }
    match name.as_str() {
        "RGGB" => CfaPattern::Rggb,
        "BGGR" => CfaPattern::Bggr,
        "GRBG" => CfaPattern::Grbg,
        "GBRG" => CfaPattern::Gbrg,
        _ => CfaPattern::Unknown,
    }
}

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

    #[test]
    fn probe_identifies_jpeg() {
        let mut h = vec![0xFF, 0xD8, 0xFF, 0xE0];
        h.extend_from_slice(&[0u8; 16]);
        assert_eq!(probe(&h), Some(Format::Jpeg));
    }

    #[test]
    fn probe_identifies_cr2_by_its_marker() {
        // Little-endian TIFF, then CR2's own magic at offset 8.
        let mut h = vec![0x49, 0x49, 0x2A, 0x00, 0x10, 0, 0, 0];
        h.extend_from_slice(b"CR\x02\x00");
        h.extend_from_slice(&[0u8; 8]);
        assert_eq!(probe(&h), Some(Format::Cr2));
    }

    #[test]
    fn probe_identifies_raf_by_signature() {
        let mut h = b"FUJIFILMCCD-RAW ".to_vec();
        h.extend_from_slice(&[0u8; 16]);
        assert_eq!(probe(&h), Some(Format::Raf));
    }

    #[test]
    fn probe_returns_none_for_ambiguous_tiff() {
        // NEF, ARW, DNG and ORF share this header; the extension has to
        // disambiguate, and claiming a format here would be a lie.
        let mut h = vec![0x49, 0x49, 0x2A, 0x00];
        h.extend_from_slice(&[0u8; 20]);
        assert_eq!(probe(&h), None);
    }

    #[test]
    fn probe_rejects_short_input() {
        assert_eq!(probe(&[0xFF, 0xD8]), None);
    }

    #[test]
    fn xtrans_is_distinguishable() {
        assert!(CfaPattern::XTrans.is_xtrans());
        assert!(!CfaPattern::Rggb.is_xtrans());
    }
}