cloudflare-soos 2.3.1

pub use crate::error::*;
use crate::Scans;

pub(crate) mod byteorder;
pub(crate) mod decoder;
pub(crate) mod marker;
pub(crate) mod parser;

pub use decoder::Decoder;
pub use decoder::MarkerData;
pub use decoder::MarkerPosition;
pub use marker::Marker;
pub use parser::CodingProcess;
pub use parser::Component;
pub use parser::Dimensions;
pub use parser::FrameInfo;
pub use parser::ScanInfo;

pub fn scans(input_file: &[u8]) -> Result<Scans> {
    let mut decoder = Decoder::new(input_file);
    let markers = decoder.decode()?;

    let mut found = Scans {
        file_size: input_file.len(),
        ..Scans::default()
    };
    let mut seen_sof = false;
    let mut seen_dc = false;
    let mut number_of_components = 3;
    let mut is_progressive = false;
    let mut dc_components_seen = [false; 3];
    let mut luma_coeff_bits_missing = [255u8; 16]; // there are 64 coefs, but only the first 16 are important for us

    for m in markers {
        // markers specify only their starting position, so to know the end position,
        // we need to wait until the next marker.
        if seen_sof {
            if found.metadata_end.is_none() {
                found.metadata_end = Some(m.position);
                if !is_progressive {
                    break;
                }
            } else if seen_dc {
                if found.first_scan_end.is_none() {
                    found.first_scan_end = Some(m.position);
                }
                // mark scan as reaching good-enough quality when either enough bits were sent (the cuttoff value is empirical),
                // or we've crossed half of the file, so we have to stop searching anyway.
                else if found.good_scan_end.is_none() && (m.position >= input_file.len()/2 || coeff_fill_factor(&luma_coeff_bits_missing) >= 91) {
                    found.good_scan_end = Some(m.position);
                    break;
                }
            }
        }

        match m.marker {
            // Image dimensions
            MarkerData::Frame(frame) => {
                number_of_components = frame.components.len();
                is_progressive = frame.coding_process == CodingProcess::DctProgressive;
                seen_sof = true;
            },
            // Progressive scan
            MarkerData::Scan(scan) => {
                if found.frame_render_start.is_none() {
                    found.frame_render_start = Some(m.position + 32); // 32 is a fudge factor to give it some data to decode
                }

                // DC is 0th component
                if *scan.spectral_selection.start() == 0 {
                    for comp in scan.component_indices.iter().copied() {
                        if let Some(seen) = dc_components_seen.get_mut(comp) {
                            *seen = true;
                        }
                    }
                }
                // DC may be split into multiple scans. Don't show incomplete DC to avoid green faces problem.
                if dc_components_seen.iter().take(number_of_components).all(|&v| v) {
                    // We don't care how many bits of DC were sent. As soon as there is something in full color, we're good.
                    seen_dc = true;
                }

                // Measure only quality of luma, assuming encoders know how to sensibly interleave chroma
                if scan.component_indices.contains(&0) {
                    for bit in scan.spectral_selection {
                        if (bit as usize) < luma_coeff_bits_missing.len() {
                            luma_coeff_bits_missing[bit as usize] = luma_coeff_bits_missing[bit as usize].min(scan.successive_approximation_low);
                        }
                    }
                }
            },
            MarkerData::Other(_) => {},
        }
    }

    // libjpeg-turbo needs a few more bytes of the next scan to render the previous one
    if let Some(first) = &mut found.first_scan_end {
        *first = (8 + *first).min(found.good_scan_end.unwrap_or(found.file_size));
    }
    if let Some(good) = &mut found.first_scan_end {
        *good = (4 + *good).min(found.file_size);
    }

    Ok(found)
}

/// Checks if there are enough coefficient bits sent for the image to look good enough.
/// This is entirely speculative. Allows either few precise ACs or many OK-ish ACs.
/// It is *not* JPEG's quality scale!
/// FIXME: this should be scaled by quantization table
fn coeff_fill_factor(bits_missing: &[u8; 16]) -> u16 {
    let mut coeff_fill_factor = 0;
    // includes DC, because posterization is also an issue
    for missing in bits_missing {
        coeff_fill_factor += match missing {
            0 => 12,
            1 => 8,
            2 => 5,
            3 => 1,
            _ => 0,
        }
    }
    coeff_fill_factor
}