DarkRoom/core/dr-pipeline/src/coverage.rs

//! TRACES: FR-DEV-3 | FR-CAT-8
//! A model's mask, in a form a sidecar can carry.
//!
//! [`MaskSource::Subject`](crate::mask::MaskSource::Subject) and
//! [`MaskSource::Category`](crate::mask::MaskSource::Category) name what they
//! cover — an index, a class, a category — and naming is enough only while
//! the run that produced the numbers is still in memory. Reopen the
//! photograph, or export it from the grid, and there is no run: the layer
//! resolves to nothing and the local adjustment silently is not applied. That
//! is what this module exists to stop. It stores the *pixels* the layer
//! covered, beside the identity rather than instead of it, so a second model
//! pass is an optimisation rather than a precondition.
//!
//! # Two levels, and why that is not a compromise
//!
//! The model hands out a byte per pixel, but nothing downstream reads more
//! than one bit of it. A subject or category layer becomes a mask by way of
//! an exact Euclidean distance field, and that field is measured from
//! `coverage >= threshold` — the soft shoulder the model produced is
//! discarded on the first line of the transform. Everything soft about the
//! rendered edge comes afterwards, from [`MaskLayer::feather`] and
//! [`MaskLayer::falloff`], which are read off the *distance*.
//!
//! [`MaskLayer::feather`]: crate::mask::MaskLayer::feather
//! [`MaskLayer::falloff`]: crate::mask::MaskLayer::falloff
//!
//! So [`RENDERED_LEVELS`] is two, and the result is not an approximation of
//! what the model said: it is exactly the part of what the model said that
//! reaches a pixel. Storing all 256 levels would be storing 1.7 MB of
//! interpolation to reconstruct a predicate — and it would not even compress,
//! because a model mask is a bilinear upsample of a coarse grid and therefore
//! has almost no two adjacent bytes alike. Measured on a simulated sky and a
//! simulated figure at 1600x1067, against 1.71 MB raw: **4.0 kB and 6.5 kB at
//! two levels**, 46 kB and 76 kB at sixteen, and 835 kB and 1.43 MB at all
//! 256 — the last two being over [`MAX_PAYLOAD`] and therefore not storable
//! at all.
//!
//! [`Coverage::encode`] still takes the level count, and it is written into
//! the line, so a later build that finds a use for the shoulder can write
//! sixteen levels and this one will read them back correctly rather than
//! misreading a stream of lengths as pairs.
//!
//! # One line, because a node is a line
//!
//! FR-NC-9 merges the edit graph per node and [`Version::merge`] does that by
//! comparing lines, so a stored mask is one `coverage = ...` line inside the
//! layer's block — the same shape the `regions = ...` line already had, for
//! the same reason. Splitting it over many lines would put a single opaque
//! blob into the merge as several independently-winnable keys, which is a
//! merge that can produce a mask neither device ever had.
//!
//! [`Version::merge`]: crate::sidecar::Version::merge
//!
//! # Hand-rolled, and it has to be
//!
//! `dr-pipeline` links nothing (ARCH §6.5a), which is what lets the descriptor
//! and codegen logic be tested without a device. That rules out `flate2`,
//! `serde` and `base64`, so the run-length coder and the digits below are
//! written out. It is forty lines, and the alternative was a dependency in the
//! one crate that has none.

use std::fmt::Write as _;

/// The number of coverage levels the renderer can actually tell apart.
///
/// Two. See the module header: the distance field is built from a threshold,
/// so a second level is the whole of the information that survives into a
/// rendered frame. Named rather than written as `2` at the call site because
/// the number is a *claim about the render path*, and a claim wants somewhere
/// to be explained.
pub const RENDERED_LEVELS: u32 = 2;

/// The most encoded payload a stored coverage may take, in bytes.
///
/// Sidecars sync over WebDAV and are read whole by every device that opens the
/// photograph, so a mask that will not compress must not be allowed to make
/// the file enormous — it is a *cache* of something a model can produce again,
/// and no cache is worth a megabyte of sync traffic per layer.
///
/// A realistic mask lands between 4 and 7 kB, so this is roughly ten times the
/// worst case anyone has measured: enough for genuinely awkward subjects —
/// foliage, chain-link, hair against a busy background — and far short of a
/// file a human cannot open. Past it [`Coverage::encode`] returns `None`, the
/// layer stores nothing, and the behaviour falls back to what it was before
/// this module existed: the mask needs the model run. Refusing rather than
/// truncating, because half a mask renders as a *wrong* mask, which is the
/// failure that announces itself to nobody.
pub const MAX_PAYLOAD: usize = 64 * 1024;

/// The most pixels a coverage read from a file may claim.
///
/// A file is not trusted. The proxy a mask is built at is bounded by the long
/// edge the segmentation runs on — under three megapixels — so this is ample
/// headroom, and it is here so that `width * height` from a corrupt line
/// cannot ask for an allocation measured in gigabytes.
pub const MAX_PIXELS: usize = 16 << 20;

/// Digits of the payload's base-64 varint. Ordered so the alphabet is stable
/// and contains nothing a line-oriented format would have to escape — no
/// whitespace, no `=`, no `#`.
const DIGITS: &[u8; 64] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_";

/// Bit set in a digit that means "another digit follows".
const CONTINUE: u32 = 32;

/// Value bits carried by one digit.
const CHUNK: u32 = 5;

const fn reverse_digits() -> [u8; 256] {
    let mut table = [255u8; 256];
    let mut i = 0;
    while i < 64 {
        table[DIGITS[i] as usize] = i as u8;
        i += 1;
    }
    table
}

/// Digit value by byte, `255` for anything that is not a digit.
const REVERSE: [u8; 256] = reverse_digits();

/// TRACES: FR-DEV-3
/// One layer's pixel coverage, held in the form it is stored in.
///
/// **Encoded, not expanded.** The struct owns the payload text rather than the
/// 1.7 MB raster it decodes to, because that raster is wanted exactly once —
/// when a distance field is built — and is held by nothing afterwards. Keeping
/// it expanded would put a megabyte and a half per layer into every undo
/// snapshot the history stack holds, to save a decode that costs far less than
/// the exact Euclidean transform immediately following it.
///
/// It also makes the round trip byte-identical for free: a coverage read from
/// a file and written back is the same characters, which is the property that
/// lets a caller skip an upload by comparing content
/// ([`Sidecar`](crate::Sidecar)).
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Coverage {
    width: usize,
    height: usize,
    levels: u32,
    payload: String,
}

impl Coverage {
    /// Encode one byte-per-pixel coverage, or `None` where it will not fit.
    ///
    /// `levels` is what the bytes are quantised to on the way in; see
    /// [`RENDERED_LEVELS`] for why two is the honest answer for a mask that is
    /// going to be thresholded.
    ///
    /// `None` for a mismatched length, a nonsensical level count, or a payload
    /// over [`MAX_PAYLOAD`] — all three meaning "do not store this", which the
    /// caller can act on identically because the fallback is the same in every
    /// case.
    pub fn encode(values: &[u8], width: usize, height: usize, levels: u32) -> Option<Self> {
        if width == 0 || height == 0 || values.len() != width.checked_mul(height)? {
            return None;
        }
        if !(2..=256).contains(&levels) {
            return None;
        }

        let top = levels - 1;
        let mut payload = String::new();
        let mut index = 0;
        while index < values.len() {
            let level = quantise(values[index], top);
            let mut run = 1;
            while index + run < values.len() && quantise(values[index + run], top) == level {
                run += 1;
            }
            push_varint(&mut payload, level as u64);
            push_varint(&mut payload, run as u64);
            index += run;

            // Checked inside the loop rather than after it: the pathological
            // input is one that runs to a payload larger than the raster, and
            // building the whole of that before deciding to throw it away is
            // the allocation this bound exists to prevent.
            if payload.len() > MAX_PAYLOAD {
                return None;
            }
        }

        Some(Self {
            width,
            height,
            levels,
            payload,
        })
    }

    /// Read the value of a sidecar `coverage` line.
    ///
    /// `None` for anything that does not describe a complete raster. A
    /// coverage is a cache, so refusing it costs a model run; accepting a
    /// partial one costs a photograph rendered with a mask that is wrong in a
    /// way nothing reports.
    pub fn parse(value: &str) -> Option<Self> {
        let mut tokens = value.split_whitespace();
        let width: usize = tokens.next()?.parse().ok()?;
        let height: usize = tokens.next()?.parse().ok()?;
        let levels: u32 = tokens.next()?.parse().ok()?;
        let payload = tokens.next()?;

        let pixels = width.checked_mul(height)?;
        if pixels == 0 || pixels > MAX_PIXELS || !(2..=256).contains(&levels) {
            return None;
        }
        if payload.len() > MAX_PAYLOAD {
            return None;
        }

        // Measured rather than expanded. The payload has to be checked here —
        // failing at the point of use would put the error in the renderer,
        // where there is no longer a file to name in the message — but a
        // library scan parses thousands of sidecars, and materialising a
        // megabyte and a half per layer to establish that the arithmetic adds
        // up would make opening the grid pay for masks nobody is rendering.
        if measure(payload, levels)? != pixels {
            return None;
        }

        Some(Self {
            width,
            height,
            levels,
            payload: payload.to_string(),
        })
    }

    /// The value to write after `coverage = `.
    pub fn to_text(&self) -> String {
        format!(
            "{} {} {} {}",
            self.width, self.height, self.levels, self.payload
        )
    }

    pub fn width(&self) -> usize {
        self.width
    }

    pub fn height(&self) -> usize {
        self.height
    }

    pub fn levels(&self) -> u32 {
        self.levels
    }

    /// The encoded payload's length in bytes — what this costs a sidecar.
    pub fn encoded_len(&self) -> usize {
        self.payload.len()
    }

    /// Expand back to one byte per pixel, at the size it was stored at.
    pub fn decode(&self) -> Vec<u8> {
        decode(&self.payload, self.width * self.height, self.levels)
            .expect("a Coverage only exists once its payload has been decoded once")
    }

    /// Expand to `width` x `height`, resampling if that is not the size it was
    /// stored at.
    ///
    /// Nearest neighbour, and deliberately: the values are a threshold's two
    /// sides, so interpolating between them would invent coverage levels that
    /// mean nothing and move the boundary by a rounding rule rather than by a
    /// measurement. The resample only runs at all when a build reads a mask
    /// stored against a different proxy edge — in the ordinary case the sizes
    /// match and this is the decode.
    pub fn decode_at(&self, width: usize, height: usize) -> Vec<u8> {
        let source = self.decode();
        if (width, height) == (self.width, self.height) {
            return source;
        }
        if width == 0 || height == 0 {
            return Vec::new();
        }

        let mut out = vec![0u8; width * height];
        for y in 0..height {
            let sy = ((y * self.height) / height).min(self.height - 1);
            let row = sy * self.width;
            for x in 0..width {
                let sx = ((x * self.width) / width).min(self.width - 1);
                out[y * width + x] = source[row + sx];
            }
        }
        out
    }
}

/// One byte to its level, rounding to nearest.
fn quantise(value: u8, top: u32) -> u32 {
    ((value as u32 * top) + 127) / 255
}

/// One level back to a byte, so that the top level is exactly 255.
fn dequantise(level: u32, top: u32) -> u8 {
    (((level * 255) + top / 2) / top).min(255) as u8
}

/// Little-endian base-64 varint: five value bits per digit, the sixth saying
/// whether another follows.
fn push_varint(out: &mut String, mut value: u64) {
    loop {
        let chunk = (value & (CONTINUE - 1) as u64) as u32;
        value >>= CHUNK;
        let more = if value != 0 { CONTINUE } else { 0 };
        let _ = out.write_char(DIGITS[(chunk | more) as usize] as char);
        if value == 0 {
            return;
        }
    }
}

/// Read one varint, returning it and how many digits it took.
fn read_varint(bytes: &[u8]) -> Option<(u64, usize)> {
    let mut value: u64 = 0;
    let mut shift = 0;
    for (taken, &byte) in bytes.iter().enumerate() {
        let digit = REVERSE[byte as usize];
        if digit == 255 {
            return None;
        }
        // A run cannot exceed MAX_PIXELS and a level cannot exceed 255, so a
        // varint past this width is a corrupt line rather than a large number.
        if shift >= 64 {
            return None;
        }
        value |= ((digit as u64) & (CONTINUE - 1) as u64) << shift;
        if digit as u32 & CONTINUE == 0 {
            return Some((value, taken + 1));
        }
        shift += CHUNK;
    }
    None
}

/// Walk a payload's runs, handing each `(level, length)` to `take`.
///
/// Returns the total length, or `None` for a payload that is not well formed:
/// a digit that is not one, a truncated varint, a zero-length run, or a level
/// the declared count does not contain.
fn walk(payload: &str, levels: u32, mut take: impl FnMut(u32, usize)) -> Option<usize> {
    let top = levels - 1;
    let bytes = payload.as_bytes();
    let mut total: usize = 0;
    let mut at = 0;

    while at < bytes.len() {
        let (level, used) = read_varint(&bytes[at..])?;
        at += used;
        let (run, used) = read_varint(&bytes[at..])?;
        at += used;

        let level = u32::try_from(level).ok()?;
        if level > top {
            return None;
        }
        let run = usize::try_from(run).ok()?;
        // A zero-length run is not a shorter way of saying anything, so it is
        // a corrupt line rather than a run to skip — and left in, two of them
        // would encode the same raster two ways and break the byte-identical
        // round trip the sidecar relies on.
        if run == 0 {
            return None;
        }
        total = total.checked_add(run)?;
        if total > MAX_PIXELS {
            return None;
        }
        take(level, run);
    }

    Some(total)
}

/// How many pixels a payload covers, without building any of them.
fn measure(payload: &str, levels: u32) -> Option<usize> {
    walk(payload, levels, |_, _| {})
}

/// Expand a payload to `pixels` bytes, or `None` if it does not describe
/// exactly that many.
fn decode(payload: &str, pixels: usize, levels: u32) -> Option<Vec<u8>> {
    let top = levels - 1;
    let mut out = Vec::with_capacity(pixels);
    let total = walk(payload, levels, |level, run| {
        out.resize(out.len() + run, dequantise(level, top));
    })?;
    (total == pixels).then_some(out)
}

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

    /// Round trip at the level count the renderer actually uses.
    fn round_trip(values: &[u8], width: usize, height: usize) -> Vec<u8> {
        let coverage = Coverage::encode(values, width, height, RENDERED_LEVELS)
            .expect("this mask should encode");
        let text = coverage.to_text();
        let read = Coverage::parse(&text).expect("what was written should parse");
        assert_eq!(read, coverage, "the round trip changed the encoding");
        assert_eq!(read.to_text(), text, "re-writing must be byte-identical");
        read.decode()
    }

    /// Two levels is exactly the predicate the distance transform applies, so
    /// the round trip must agree with it on every pixel.
    fn thresholded(values: &[u8]) -> Vec<bool> {
        values.iter().map(|&v| v >= 128).collect()
    }

    #[test]
    fn an_empty_mask_round_trips() {
        let values = vec![0u8; 64 * 32];
        let back = round_trip(&values, 64, 32);
        assert_eq!(back, values);
    }

    #[test]
    fn a_full_mask_round_trips() {
        let values = vec![255u8; 64 * 32];
        let back = round_trip(&values, 64, 32);
        assert_eq!(back, values);
    }

    #[test]
    fn a_single_pixel_mask_round_trips() {
        let mut values = vec![0u8; 64 * 32];
        values[17 * 64 + 33] = 255;
        let back = round_trip(&values, 64, 32);
        assert_eq!(back, values);
    }

    #[test]
    fn a_one_pixel_raster_round_trips() {
        assert_eq!(round_trip(&[255], 1, 1), vec![255]);
        assert_eq!(round_trip(&[0], 1, 1), vec![0]);
    }

    /// The whole of the fidelity claim: two levels loses nothing the renderer
    /// could have used, because the renderer thresholds.
    #[test]
    fn two_levels_preserve_the_threshold_exactly() {
        let values: Vec<u8> = (0..=255u8).collect();
        let back = round_trip(&values, 16, 16);
        assert_eq!(thresholded(&back), thresholded(&values));
        // And the shoulder really is gone, which is the cost being paid.
        assert!(back.iter().all(|&v| v == 0 || v == 255));
    }

    /// A soft edge quantised to sixteen levels stays within one step of what
    /// went in, so a later build that wants the shoulder can have it.
    #[test]
    fn sixteen_levels_are_within_one_step() {
        let values: Vec<u8> = (0..256).map(|i| i as u8).collect();
        let coverage = Coverage::encode(&values, 16, 16, 16).expect("should encode");
        let back = Coverage::parse(&coverage.to_text())
            .expect("should parse")
            .decode();
        for (a, b) in values.iter().zip(&back) {
            assert!(
                (*a as i32 - *b as i32).abs() <= 255 / 15 / 2 + 1,
                "{a} came back as {b}"
            );
        }
    }

    /// The case run-length coding is worst at. It must refuse rather than
    /// write a payload larger than the raster it came from.
    #[test]
    fn alternating_detail_is_refused_rather_than_expanded() {
        let (w, h) = (512, 512);
        let values: Vec<u8> = (0..w * h)
            .map(|i| if i % 2 == 0 { 0 } else { 255 })
            .collect();
        assert!(
            Coverage::encode(&values, w, h, RENDERED_LEVELS).is_none(),
            "a checkerboard must not be stored"
        );
    }

    /// Small enough to fit, and still exact — the bound is on size, not on
    /// shape, so awkward detail that *does* fit must survive intact.
    #[test]
    fn alternating_detail_that_fits_is_exact() {
        let (w, h) = (64, 64);
        let values: Vec<u8> = (0..w * h)
            .map(|i| if i % 2 == 0 { 0 } else { 255 })
            .collect();
        let back = round_trip(&values, w, h);
        assert_eq!(back, values);
    }

    #[test]
    fn a_mask_of_the_wrong_length_is_refused() {
        assert!(Coverage::encode(&[0u8; 10], 4, 4, RENDERED_LEVELS).is_none());
        assert!(Coverage::encode(&[], 0, 0, RENDERED_LEVELS).is_none());
    }

    #[test]
    fn a_payload_that_does_not_cover_the_raster_is_refused() {
        let values = vec![0u8; 32];
        let coverage = Coverage::encode(&values, 8, 4, RENDERED_LEVELS).expect("should encode");
        let payload = coverage.to_text();
        let payload = payload.rsplit_once(' ').expect("a payload").1;
        // The same payload, against a raster twice the size it covers.
        assert!(Coverage::parse(&format!("8 4 2 {payload}")).is_some());
        assert!(Coverage::parse(&format!("8 8 2 {payload}")).is_none());
    }

    #[test]
    fn nonsense_is_refused_rather_than_guessed_at() {
        assert!(Coverage::parse("").is_none());
        assert!(Coverage::parse("8 4 2").is_none(), "no payload");
        assert!(
            Coverage::parse("8 4 1 AA").is_none(),
            "one level is not a mask"
        );
        assert!(Coverage::parse("8 4 2 ****").is_none(), "not digits");
        assert!(
            Coverage::parse(&format!("{} {} 2 AA", usize::MAX, usize::MAX)).is_none(),
            "a size that overflows must not be believed"
        );
        assert!(
            Coverage::parse("100000 100000 2 A_____").is_none(),
            "a raster past the cap must not be allocated"
        );
    }

    /// A level a payload is not allowed to name, in a file that names it.
    #[test]
    fn a_level_outside_the_range_is_refused() {
        // "BB" is level 1, run 1 — the shortest legal payload there is.
        assert_eq!(decode("BB", 1, 2), Some(vec![255]));
        // "DB" is level 3, run 1, and a two-level coverage has no level 3.
        assert!(decode("DB", 1, 2).is_none());
        // A run of zero says nothing and is refused rather than skipped.
        assert!(decode("BA", 1, 2).is_none());
    }

    #[test]
    fn resampling_lands_on_the_same_shape() {
        let (w, h) = (32, 32);
        let mut values = vec![0u8; w * h];
        for y in 8..24 {
            for x in 8..24 {
                values[y * w + x] = 255;
            }
        }
        let coverage = Coverage::encode(&values, w, h, RENDERED_LEVELS).expect("should encode");

        let same = coverage.decode_at(w, h);
        assert_eq!(same, values, "the matching size must not resample at all");

        let half = coverage.decode_at(16, 16);
        assert_eq!(half.len(), 256);
        assert_eq!(half.iter().filter(|&&v| v == 255).count(), 64);

        let double = coverage.decode_at(64, 64);
        assert_eq!(double.len(), 4096);
        assert_eq!(double.iter().filter(|&&v| v == 255).count(), 1024);
    }

    /// Long runs cross rows, which is what makes a flat mask cost almost
    /// nothing: a 1600x1067 empty frame is two numbers.
    #[test]
    fn a_flat_mask_costs_almost_nothing() {
        let values = vec![0u8; 1600 * 1067];
        let coverage =
            Coverage::encode(&values, 1600, 1067, RENDERED_LEVELS).expect("should encode");
        assert!(
            coverage.encoded_len() < 16,
            "an empty mask took {} bytes",
            coverage.encoded_len()
        );
    }

    /// What a real one costs. The shape is a bilinear upsample of a coarse
    /// grid, which is what both models produce, so the run structure is the
    /// one a photograph actually gives.
    #[test]
    fn a_realistic_mask_fits_in_a_sidecar() {
        let (w, h) = (1600usize, 1067usize);
        let (gw, gh) = (160usize, 107usize);
        let mut grid = vec![0f32; gw * gh];
        for y in 0..gh {
            for x in 0..gw {
                let dx = (x as f32 - 80.0) / 26.0;
                let dy = (y as f32 - 60.0) / 42.0;
                let r = (dx * dx + dy * dy).sqrt()
                    + 0.06 * ((y as f32 * 0.9).sin() * (x as f32 * 0.7).cos());
                grid[y * gw + x] = 1.0 / (1.0 + ((r - 1.0) * 9.0).exp());
            }
        }
        let mut values = vec![0u8; w * h];
        for y in 0..h {
            let fy = ((y as f32 + 0.5) / h as f32 * gh as f32 - 0.5).max(0.0);
            let (y0, ty) = (fy.floor() as usize, fy.fract());
            let y1 = (y0 + 1).min(gh - 1);
            for x in 0..w {
                let fx = ((x as f32 + 0.5) / w as f32 * gw as f32 - 0.5).max(0.0);
                let (x0, tx) = (fx.floor() as usize, fx.fract());
                let x1 = (x0 + 1).min(gw - 1);
                let a = grid[y0 * gw + x0] * (1.0 - tx) + grid[y0 * gw + x1] * tx;
                let b = grid[y1 * gw + x0] * (1.0 - tx) + grid[y1 * gw + x1] * tx;
                values[y * w + x] = ((a * (1.0 - ty) + b * ty) * 255.0) as u8;
            }
        }

        let coverage = Coverage::encode(&values, w, h, RENDERED_LEVELS).expect("should encode");
        let expected: Vec<u8> = values
            .iter()
            .map(|&v| if v >= 128 { 255 } else { 0 })
            .collect();
        assert_eq!(
            coverage.decode(),
            expected,
            "the stored mask must threshold identically to the model's"
        );
        // Measured at 6,464 bytes; the bound is loose enough not to fail over
        // a change of rounding and tight enough to catch a coder that has
        // stopped coding. Against 1,707,200 bytes raw.
        assert!(
            coverage.encoded_len() < 8 * 1024,
            "a realistic subject took {} bytes",
            coverage.encoded_len()
        );
    }
}