forge-ibl-bake 0.1.0

// SPDX-License-Identifier: MIT OR Apache-2.0
// forge-ibl-bake · deveraux.dev
//! Radiance RGBE (`.hdr`) parser → linear f32 RGB pixel buffer.
//!
//! Hand-rolled: supports the new-style RLE scan-line encoding used by every
//! modern HDRI (Poly Haven, HDRI Haven, etc.). Inputs: raw file bytes.
//! Outputs: `HdrImage` with row-major `pixels: Vec<[f32;3]>` in linear light.
//! DESIGN_TARGET: old-style RLE (pre-1990 files) is not supported — returns Err.

use std::path::Path;

use crate::error::BakeError;

/// Linear-light equirectangular image decoded from a Radiance `.hdr` file.
#[derive(Debug)]
pub struct HdrImage {
    pub width: u32,
    pub height: u32,
    /// Row-major, left-to-right top-to-bottom, linear RGB (not sRGB).
    pub pixels: Vec<[f32; 3]>,
}

/// Load a `.hdr` file from `path` into linear f32 RGB.
pub fn load_hdr_file(path: &Path) -> Result<HdrImage, BakeError> {
    let data = std::fs::read(path)?;
    load_hdr(&data)
}

/// Parse Radiance RGBE bytes (already in memory) into linear f32 RGB.
pub fn load_hdr(data: &[u8]) -> Result<HdrImage, BakeError> {
    check_magic(data)?;
    let (body_offset, width, height) = parse_header(data)?;
    let pixels = decode_pixels(&data[body_offset..], width, height)?;
    Ok(HdrImage { width, height, pixels })
}

fn check_magic(data: &[u8]) -> Result<(), BakeError> {
    if !data.starts_with(b"#?RADIANCE") && !data.starts_with(b"#?RGBE") {
        return Err(BakeError::InvalidHdr(
            "missing #?RADIANCE / #?RGBE magic".into(),
        ));
    }
    Ok(())
}

/// Returns (offset_to_first_pixel_byte, width, height).
fn parse_header(data: &[u8]) -> Result<(usize, u32, u32), BakeError> {
    let mut i = 0usize;
    let n = data.len();

    loop {
        // Find end of line.
        let line_start = i;
        while i < n && data[i] != b'\n' {
            i += 1;
        }
        let raw_line = &data[line_start..i];
        if i < n {
            i += 1; // consume '\n'
        } else {
            return Err(BakeError::InvalidHdr("truncated header".into()));
        }

        // Strip optional '\r'.
        let line = raw_line.strip_suffix(b"\r").unwrap_or(raw_line);

        if line.is_empty() {
            // Blank line → next line is the resolution string.
            let res_start = i;
            while i < n && data[i] != b'\n' {
                i += 1;
            }
            let res_raw = &data[res_start..i];
            if i < n {
                i += 1;
            }
            let res_line = std::str::from_utf8(res_raw.strip_suffix(b"\r").unwrap_or(res_raw))
                .map_err(|_| BakeError::InvalidHdr("resolution line not UTF-8".into()))?;
            let (w, h) = parse_size_line(res_line)?;
            return Ok((i, w, h));
        }

        // Validate FORMAT if present; reject XYZE.
        if let Some(rest) = line.strip_prefix(b"FORMAT=") {
            let fmt = std::str::from_utf8(rest).unwrap_or("").trim();
            if fmt != "32-bit_rle_rgbe" {
                return Err(BakeError::InvalidHdr(format!(
                    "unsupported format: {fmt}"
                )));
            }
        }
    }
}

/// Parse "-Y <h> +X <w>" (standard Radiance orientation). Returns (width, height).
/// Any other orientation (e.g. "+Y", "+X -Y") is rejected to prevent silent flips/transposes.
fn parse_size_line(line: &str) -> Result<(u32, u32), BakeError> {
    let parts: Vec<&str> = line.split_whitespace().collect();
    if parts.len() < 4 {
        return Err(BakeError::InvalidHdr(format!(
            "bad resolution line: '{line}'"
        )));
    }
    if parts[0] != "-Y" || parts[2] != "+X" {
        return Err(BakeError::InvalidHdr(format!(
            "unsupported orientation '{} {}'; only '-Y +X' (standard Radiance) is supported",
            parts[0], parts[2]
        )));
    }
    let height: u32 = parts[1]
        .parse()
        .map_err(|_| BakeError::InvalidHdr("bad height".into()))?;
    let width: u32 = parts[3]
        .parse()
        .map_err(|_| BakeError::InvalidHdr("bad width".into()))?;
    Ok((width, height))
}

fn decode_pixels(
    data: &[u8],
    width: u32,
    height: u32,
) -> Result<Vec<[f32; 3]>, BakeError> {
    let w = width as usize;
    let h = height as usize;
    let mut pixels = vec![[0.0f32; 3]; w * h];
    let mut pos = 0usize;
    // Allocated once; cleared per scanline to avoid per-row heap churn.
    let mut ch_data: [Vec<u8>; 4] = [
        Vec::with_capacity(w),
        Vec::with_capacity(w),
        Vec::with_capacity(w),
        Vec::with_capacity(w),
    ];

    for y in 0..h {
        if pos + 4 > data.len() {
            return Err(BakeError::InvalidHdr(format!(
                "unexpected end at scan line {y}"
            )));
        }

        // New-style RLE: first two bytes are 0x02 0x02, then big-endian width.
        if data[pos] == 2 && data[pos + 1] == 2 {
            let scan_w =
                ((data[pos + 2] as usize) << 8) | (data[pos + 3] as usize);
            if scan_w != w {
                return Err(BakeError::InvalidHdr(format!(
                    "scan width {scan_w} ≠ image width {w} at y={y}"
                )));
            }
            pos += 4;

            // Decode 4 channels (R, G, B, E) independently with RLE.
            for ch in ch_data.iter_mut() { ch.clear(); }
            for ch in 0..4 {
                let out = &mut ch_data[ch];
                while out.len() < w {
                    if pos >= data.len() {
                        return Err(BakeError::InvalidHdr(
                            "RLE data underflow".into(),
                        ));
                    }
                    let code = data[pos];
                    pos += 1;
                    if code > 128 {
                        // Run: (code - 128) copies of next byte.
                        let count = (code - 128) as usize;
                        if pos >= data.len() {
                            return Err(BakeError::InvalidHdr(
                                "RLE run value missing".into(),
                            ));
                        }
                        let val = data[pos];
                        pos += 1;
                        for _ in 0..count {
                            out.push(val);
                        }
                    } else {
                        // Non-run: `code` literal bytes follow.
                        let count = code as usize;
                        if count == 0 {
                            return Err(BakeError::InvalidHdr("RLE zero-count literal".into()));
                        }
                        if pos + count > data.len() {
                            return Err(BakeError::InvalidHdr(
                                "RLE nonrun underflow".into(),
                            ));
                        }
                        out.extend_from_slice(&data[pos..pos + count]);
                        pos += count;
                    }
                }
                if out.len() != w {
                    return Err(BakeError::InvalidHdr(format!(
                        "channel {ch} decoded {} bytes, expected {w} at y={y}",
                        out.len()
                    )));
                }
            }

            for x in 0..w {
                pixels[y * w + x] = rgbe_to_f32(
                    ch_data[0][x],
                    ch_data[1][x],
                    ch_data[2][x],
                    ch_data[3][x],
                );
            }
        } else {
            // Uncompressed RGBE: 4 bytes per pixel (common in very small or
            // very old files, or when scan-line width ≤ 8).
            for x in 0..w {
                if pos + 4 > data.len() {
                    return Err(BakeError::InvalidHdr("raw pixel underflow".into()));
                }
                pixels[y * w + x] =
                    rgbe_to_f32(data[pos], data[pos + 1], data[pos + 2], data[pos + 3]);
                pos += 4;
            }
        }
    }

    Ok(pixels)
}

/// Convert an RGBE quad (Ward 1991) to linear f32 RGB.
/// Formula: RGB = (mantissa + 0.5) × 2^(exponent − 128 − 8).
#[inline]
fn rgbe_to_f32(r: u8, g: u8, b: u8, e: u8) -> [f32; 3] {
    if e == 0 {
        return [0.0, 0.0, 0.0];
    }
    let scale = 2.0f32.powi(e as i32 - 128 - 8);
    [
        (r as f32 + 0.5) * scale,
        (g as f32 + 0.5) * scale,
        (b as f32 + 0.5) * scale,
    ]
}

/// Load an OpenEXR file from `path` into linear f32 RGB.
/// Requires the `exr` feature: `cargo build --features exr`.
#[cfg(feature = "exr")]
pub fn load_exr_file(path: &Path) -> Result<HdrImage, BakeError> {
    use ::exr::prelude::*;
    struct Buf { w: usize, h: usize, pix: Vec<[f32; 3]> }
    let img = read()
        .no_deep_data()
        .largest_resolution_level()
        .rgba_channels(
            |res, _| Buf {
                w: res.x(),
                h: res.y(),
                pix: vec![[0.0f32; 3]; res.x() * res.y()],
            },
            |buf: &mut Buf, pos, (r, g, b, _): (f32, f32, f32, f32)| {
                buf.pix[pos.y() * buf.w + pos.x()] = [r, g, b];
            },
        )
        .first_valid_layer()
        .all_attributes()
        .from_file(path)
        .map_err(|e| BakeError::InvalidHdr(format!("EXR: {e}")))?;
    let Buf { w, h, pix } = img.layer_data.channel_data.pixels;
    Ok(HdrImage { width: w as u32, height: h as u32, pixels: pix })
}