DarkRoom/core/dr-gpu/tests/spot_removal.rs

//! TRACES: FR-DEV-8
//! Repairs, drawn on a real device.
//!
//! `dr-pipeline`'s tests assert the model and the record packing; nothing there
//! can say whether the disc lands where the photographer put it. That is what
//! this file is for, and the cases it covers are the ones where a repair goes
//! wrong *quietly*:
//!
//! - the mark is still there, because the disc landed beside it;
//! - the repair works on screen and not in the export, because a length was
//!   converted in the wrong units;
//! - the repair works until the photograph is cropped or rotated, because the
//!   framing was applied to the pixels and not to the spot.
//!
//! The frame is a flat grey field with one black mark on it, which makes every
//! assertion here countable: a repair either leaves dark pixels or it does not.

use dr_gpu::{AdjustPass, DemosaicedImage, GpuContext};
use dr_pipeline::spot::{Spot, SpotMode};
use dr_pipeline::{Affects, EditGraph};
use dr_types::ColourSpace;

const SIZE: u32 = 128;
const GREY: u8 = 128;

/// Where the mark is, in normalised coordinates, and how big it is in frame
/// units. Off-centre on both axes so that a repair landing on a mirrored or
/// transposed position fails rather than passing by symmetry.
const MARK: (f32, f32) = (0.3, 0.65);
const MARK_RADIUS: f32 = 0.03;

fn ctx() -> Option<GpuContext> {
    match pollster::block_on(GpuContext::new_headless()) {
        Ok(c) => Some(c),
        Err(e) => {
            eprintln!("skipping: no GPU adapter ({e})");
            None
        }
    }
}

/// A flat grey frame with one black mark on it — a dust spot, idealised.
fn marked_frame(ctx: &GpuContext) -> DemosaicedImage {
    let data: Vec<u8> = (0..SIZE * SIZE)
        .flat_map(|i| {
            let (x, y) = ((i % SIZE) as f32, (i / SIZE) as f32);
            let (cx, cy) = (MARK.0 * SIZE as f32, MARK.1 * SIZE as f32);
            let r = MARK_RADIUS * SIZE as f32;
            let v = if (x - cx).hypot(y - cy) <= r {
                0u8
            } else {
                GREY
            };
            [v, v, v, 255]
        })
        .collect();
    DemosaicedImage::from_rgba8(ctx, &data, SIZE, SIZE).expect("upload")
}

/// A repair covering the mark, reading from clean grey to its right.
///
/// The disc is twice the mark, so the mark sits entirely inside the solid core
/// and none of it falls in the feathered rim — otherwise this file would be
/// asserting a blend rather than a repair.
fn repair(mode: SpotMode) -> Spot {
    let mut spot = Spot::new(MARK, (0.3, 0.0), MARK_RADIUS * 2.0);
    spot.mode = mode;
    spot
}

fn render(pass: &mut AdjustPass, graph: &EditGraph, source: &DemosaicedImage, out: u32) -> Vec<u8> {
    let shader = graph.compose_for(ColourSpace::Srgb);
    let (w, h) = graph.output_size(source.size().0, source.size().1);
    let (w, h) = (w.min(out), h.min(out));
    let detail = graph.compose_detail_for(source.size(), (w, h), ColourSpace::Srgb);
    let key = graph.invalidation().through(Affects::Colour);
    pass.render_detailed(source, &shader, w, h, None, &detail, key)
        .expect("render");
    pass.export_pixels().expect("readback").0
}

/// How many pixels are darker than anything a grey field contains.
///
/// The mark is the only dark thing in the frame, so this counts what is left of
/// it — and counts it wherever it ended up, which is what makes the same
/// assertion work after a crop or a rotation.
fn dark_pixels(pixels: &[u8]) -> usize {
    pixels.chunks_exact(4).filter(|p| p[0] < GREY - 24).count()
}

/// The whole feature in one assertion: the mark is there, and then it is not.
#[test]
fn a_clone_removes_the_mark() {
    let Some(ctx) = ctx() else { return };
    let source = marked_frame(&ctx);
    let mut pass = AdjustPass::new(&ctx);

    let before = dark_pixels(&render(
        &mut pass,
        &EditGraph::default_chain(),
        &source,
        SIZE,
    ));
    assert!(before > 20, "the frame is supposed to have a mark on it");

    let mut graph = EditGraph::default_chain();
    graph.spots_mut().place(repair(SpotMode::Clone));

    let after = dark_pixels(&render(&mut pass, &graph, &source, SIZE));
    assert_eq!(after, 0, "{before} dark pixels before, {after} after");
}

/// The grey the repair lays down has to be the *photograph's* grey. A repair
/// that removes the mark by darkening or brightening the disc passes the count
/// above and is still visibly a disc.
#[test]
fn the_patch_is_the_photograph_and_not_an_approximation_of_it() {
    let Some(ctx) = ctx() else { return };
    let source = marked_frame(&ctx);
    let mut pass = AdjustPass::new(&ctx);

    let mut graph = EditGraph::default_chain();
    graph.spots_mut().place(repair(SpotMode::Clone));
    let pixels = render(&mut pass, &graph, &source, SIZE);

    let centre =
        ((MARK.1 * SIZE as f32) as u32 * SIZE + (MARK.0 * SIZE as f32) as u32) as usize * 4;
    assert!(
        pixels[centre].abs_diff(GREY) <= 2,
        "the repaired centre reads {}, the field is {GREY}",
        pixels[centre]
    );
}

/// A repair is stored as a fraction of the frame, so it must land in the same
/// *place* whatever size the frame is drawn at — which is the difference
/// between a preview that tells the truth and an export that does not
/// (FR-DSP-1).
#[test]
fn a_proxy_and_an_export_repair_the_same_thing() {
    let Some(ctx) = ctx() else { return };
    let source = marked_frame(&ctx);
    let mut pass = AdjustPass::new(&ctx);

    let mut graph = EditGraph::default_chain();
    graph.spots_mut().place(repair(SpotMode::Clone));

    assert_eq!(dark_pixels(&render(&mut pass, &graph, &source, SIZE)), 0);
    assert_eq!(
        dark_pixels(&render(&mut pass, &graph, &source, SIZE / 4)),
        0,
        "the repair missed the mark at a quarter size"
    );
}

/// The framing is applied to the spot, not to the pixels afterwards. If the
/// centre were mapped and the offset were not, this is the test that fails: the
/// disc would land on the mark and read from the wrong side of the frame.
#[test]
fn a_rotated_photograph_carries_its_repairs_round_with_it() {
    let Some(ctx) = ctx() else { return };
    let source = marked_frame(&ctx);
    let mut pass = AdjustPass::new(&ctx);

    let mut graph = EditGraph::default_chain();
    graph.spots_mut().place(repair(SpotMode::Clone));
    graph.rotate_quarters(1);

    assert_eq!(
        dark_pixels(&render(&mut pass, &graph, &source, SIZE)),
        0,
        "the mark came back when the frame was turned"
    );
}

/// And a crop, which moves the origin and the scale at once.
#[test]
fn a_crop_carries_its_repairs_with_it() {
    let Some(ctx) = ctx() else { return };
    let source = marked_frame(&ctx);
    let mut pass = AdjustPass::new(&ctx);

    let mut graph = EditGraph::default_chain();
    graph.spots_mut().place(repair(SpotMode::Clone));
    graph.set_crop(dr_pipeline::CropRect {
        x: 0.1,
        y: 0.4,
        width: 0.5,
        height: 0.5,
    });

    assert_eq!(
        dark_pixels(&render(&mut pass, &graph, &source, SIZE)),
        0,
        "the mark is inside this crop and the repair no longer covers it"
    );
}

/// Opacity is a real control: at zero the repair is off, and the mark is
/// exactly as it was.
#[test]
fn a_transparent_repair_draws_nothing() {
    let Some(ctx) = ctx() else { return };
    let source = marked_frame(&ctx);
    let mut pass = AdjustPass::new(&ctx);

    let bare = dark_pixels(&render(
        &mut pass,
        &EditGraph::default_chain(),
        &source,
        SIZE,
    ));

    let mut graph = EditGraph::default_chain();
    let mut spot = repair(SpotMode::Clone);
    spot.set_opacity(0.0);
    graph.spots_mut().place(spot);

    assert_eq!(dark_pixels(&render(&mut pass, &graph, &source, SIZE)), bare);
}

/// Placing repairs must not recompile: the list is in a storage buffer and the
/// shader never learns how long it is, so a photographer working through a
/// dusty sky pays one compilation.
#[test]
fn placing_repairs_compiles_one_pipeline() {
    let Some(ctx) = ctx() else { return };
    let source = marked_frame(&ctx);
    let mut pass = AdjustPass::new(&ctx);

    let mut graph = EditGraph::default_chain();
    for i in 0..6 {
        let y = 0.1 + 0.1 * i as f32;
        graph
            .spots_mut()
            .place(Spot::new((0.5, y), (0.1, 0.0), 0.02));
        render(&mut pass, &graph, &source, SIZE);
    }

    assert_eq!(
        pass.cached_detail_pipelines(),
        1,
        "six repairs, one compiled pipeline"
    );
}

// ---------------------------------------------------------------------------
// Heal (FR-DEV-8)
// ---------------------------------------------------------------------------

/// A frame whose brightness ramps across it, with one black mark on it.
///
/// This is the case that separates the two modes. A clone copies a patch from
/// somewhere else on the ramp, so it arrives at the wrong level and leaves a
/// disc of the right texture and the wrong tone; a heal carries the difference
/// across from the boundary and leaves nothing.
fn ramped_frame(ctx: &GpuContext) -> DemosaicedImage {
    let data: Vec<u8> = (0..SIZE * SIZE)
        .flat_map(|i| {
            let (x, y) = ((i % SIZE) as f32, (i / SIZE) as f32);
            let (cx, cy) = (MARK.0 * SIZE as f32, MARK.1 * SIZE as f32);
            let r = MARK_RADIUS * SIZE as f32;
            // 64 at the left edge to 192 at the right: a ramp steep enough that
            // a clone from a third of a frame away is unmistakably wrong, and
            // shallow enough to stay well inside the range.
            let ramp = 64.0 + 128.0 * x / SIZE as f32;
            let v = if (x - cx).hypot(y - cy) <= r {
                0u8
            } else {
                ramp as u8
            };
            [v, v, v, 255]
        })
        .collect();
    DemosaicedImage::from_rgba8(ctx, &data, SIZE, SIZE).expect("upload")
}

/// What the ramp says at a column, as the byte the renderer should produce.
fn ramp_at(x: u32) -> u8 {
    (64.0 + 128.0 * x as f32 / SIZE as f32) as u8
}

/// The worst a repair is wrong by, over the disc it covers.
///
/// Measured against the ramp the photograph would have had if the mark had
/// never been there, which is the only definition of "repaired" worth
/// asserting: a repair that removes the mark and leaves the wrong tone has not
/// repaired anything, it has drawn a different mark.
fn worst_error(pixels: &[u8]) -> u8 {
    let (cx, cy) = (MARK.0 * SIZE as f32, MARK.1 * SIZE as f32);
    let r = MARK_RADIUS * SIZE as f32;
    let mut worst = 0u8;
    for y in 0..SIZE {
        for x in 0..SIZE {
            if (x as f32 - cx).hypot(y as f32 - cy) > r {
                continue;
            }
            let got = pixels[((y * SIZE + x) * 4) as usize];
            worst = worst.max(got.abs_diff(ramp_at(x)));
        }
    }
    worst
}

/// The measurement the mode exists for, and the one that decides
/// `RIM_SAMPLES`: on a gradient, a heal is right and a clone is not.
#[test]
fn a_heal_takes_the_tone_from_the_hole_it_fills() {
    let Some(ctx) = ctx() else { return };
    let source = ramped_frame(&ctx);
    let mut pass = AdjustPass::new(&ctx);

    let mut cloned = EditGraph::default_chain();
    cloned.spots_mut().place(repair(SpotMode::Clone));
    let clone_error = worst_error(&render(&mut pass, &cloned, &source, SIZE));

    let mut healed = EditGraph::default_chain();
    healed.spots_mut().place(repair(SpotMode::Heal));
    let heal_error = worst_error(&render(&mut pass, &healed, &source, SIZE));

    // The clone is wrong by roughly the ramp across the source offset — about
    // 38 levels here — and it is wrong across the whole disc.
    assert!(
        clone_error > 20,
        "the clone was supposed to be visibly wrong, and is off by {clone_error}"
    );
    // Measured at 0 on the reference device — the ramp is linear, so the
    // boundary difference is constant all the way round and the membrane
    // reproduces it exactly. The tolerance is for the rounding a different
    // driver may do, not for the method being approximate here.
    assert!(
        heal_error <= 1,
        "the heal is off by {heal_error} levels; the clone it has to beat is off by {clone_error}"
    );
}

/// And the repair still has to remove the mark: a membrane that matched the
/// boundary while leaving the black disc underneath would pass the measurement
/// above by averaging its way past it.
#[test]
fn a_heal_removes_the_mark_as_well_as_matching_the_tone() {
    let Some(ctx) = ctx() else { return };
    let source = ramped_frame(&ctx);
    let mut pass = AdjustPass::new(&ctx);

    let mut graph = EditGraph::default_chain();
    graph.spots_mut().place(repair(SpotMode::Heal));
    let pixels = render(&mut pass, &graph, &source, SIZE);

    let mark = (MARK.0 * SIZE as f32) as u32;
    for x in mark - 2..=mark + 2 {
        let got = pixels[(((MARK.1 * SIZE as f32) as u32 * SIZE + x) * 4) as usize];
        assert!(
            got.abs_diff(ramp_at(x)) <= 3,
            "column {x} reads {got}, the ramp says {}",
            ramp_at(x)
        );
    }
}

/// A heal on a flat field is a clone: the boundary difference is zero all the
/// way round, so the membrane is zero and neither mode has anything to add.
/// Worth pinning, because a membrane that quietly tinted a flat repair would
/// be invisible in the gradient test above.
#[test]
fn a_heal_on_a_flat_field_changes_nothing_a_clone_would_not() {
    let Some(ctx) = ctx() else { return };
    let source = marked_frame(&ctx);
    let mut pass = AdjustPass::new(&ctx);

    let mut cloned = EditGraph::default_chain();
    cloned.spots_mut().place(repair(SpotMode::Clone));
    let a = render(&mut pass, &cloned, &source, SIZE);

    let mut healed = EditGraph::default_chain();
    healed.spots_mut().place(repair(SpotMode::Heal));
    let b = render(&mut pass, &healed, &source, SIZE);

    let worst = a
        .iter()
        .zip(&b)
        .map(|(x, y)| x.abs_diff(*y))
        .max()
        .unwrap_or(0);
    assert!(
        worst <= 1,
        "heal and clone differ by {worst} on a flat field"
    );
}