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

//! Local adjustments, end to end on a device.
//!
//! The unit tests either side of this one check halves: `dr-pipeline` asserts
//! the generated WGSL says the right thing, and `dr-gpu`'s mask tests assert
//! an array of the right shape comes out. Neither would notice if the two
//! agreed with each other and both were wrong — a mask sampled with x and y
//! swapped satisfies both.
//!
//! So this renders a real frame and reads the pixels back: the masked region
//! must change, the rest must not, and the boundary must fall where the label
//! field says it does.

use dr_gpu::{AdjustPass, DemosaicedImage, GpuContext, LabelField, MaskPass};
use dr_pipeline::descriptor::ParamId;
use dr_pipeline::mask::{MaskLayer, MaskSource, MaskStack};
use dr_pipeline::operation::compose_full;
use dr_pipeline::spot::SpotSet;
use dr_pipeline::{ops, EditGraph, Framing};
use dr_types::ColourSpace;

const SIZE: u32 = 32;

fn ctx() -> Option<GpuContext> {
    pollster::block_on(GpuContext::new_headless()).ok()
}

/// A flat mid-grey JPEG-path image, so any change is the adjustment's.
fn grey(ctx: &GpuContext) -> DemosaicedImage {
    grey_at(ctx, SIZE, SIZE)
}

fn grey_at(ctx: &GpuContext, w: u32, h: u32) -> DemosaicedImage {
    let data: Vec<u8> = (0..w * h).flat_map(|_| [128, 128, 128, 255]).collect();
    DemosaicedImage::from_rgba8(ctx, &data, w, h).expect("upload")
}

/// Two regions: 0 is the left half, 1 the right.
fn split_field(ctx: &GpuContext) -> LabelField {
    let labels: Vec<u32> = (0..SIZE * SIZE)
        .map(|i| u32::from(i % SIZE >= SIZE / 2))
        .collect();
    LabelField::upload(ctx, &labels, SIZE, SIZE, 2).expect("label upload")
}

/// A layer brightening whatever it covers, by a lot, so it cannot be missed.
fn brighten(source: MaskSource) -> MaskLayer {
    let mut layer = MaskLayer::new("m1", source);
    layer.set_param("exposure", ParamId("exposure"), 2.0);
    layer
}

fn luma_at(pixels: &[u8], x: u32, y: u32) -> u8 {
    pixels[((y * SIZE + x) * 4) as usize]
}

/// Render `stack` over flat grey and hand back the RGBA8 result.
fn render(ctx: &GpuContext, stack: &MaskStack, field: Option<&LabelField>) -> Vec<u8> {
    render_at(ctx, stack, field, SIZE, SIZE)
}

fn render_at(
    ctx: &GpuContext,
    stack: &MaskStack,
    field: Option<&LabelField>,
    w: u32,
    h: u32,
) -> Vec<u8> {
    let source = grey_at(ctx, w, h);
    let shader = compose_full(
        &ops::chain(),
        &Framing::new(),
        ColourSpace::Srgb,
        stack,
        &SpotSet::new(),
        &[],
    );

    let mut masks = MaskPass::new(ctx).expect("mask pass");
    let array = masks
        .render(stack, field, None, None, w, h)
        .expect("rasterise");

    let mut adjust = AdjustPass::new(ctx);
    adjust
        .render_masked(&source, &shader, w, h, Some(array))
        .expect("render");
    adjust.export_pixels().expect("readback").0
}

#[test]
fn a_region_mask_changes_only_the_regions_it_names() {
    let Some(ctx) = ctx() else {
        eprintln!("no adapter; skipping");
        return;
    };
    let field = split_field(&ctx);

    let mut stack = MaskStack::new();
    stack.push(brighten(MaskSource::Regions {
        signature: 1,
        level: 2,
        ids: vec![0],
    }));

    let pixels = render(&ctx, &stack, Some(&field));

    // Sampled well inside each half, clear of the feathered boundary.
    let inside = luma_at(&pixels, 4, SIZE / 2);
    let outside = luma_at(&pixels, SIZE - 5, SIZE / 2);

    assert!(
        inside > outside + 40,
        "the masked half should be much brighter: {inside} vs {outside}"
    );
    assert!(
        (120..=136).contains(&outside),
        "the unmasked half must be untouched mid-grey, got {outside}"
    );
}

/// The failure a swapped axis or an inverted comparison would produce, and
/// which the "inside is brighter" assertion alone would not catch.
#[test]
fn inverting_a_region_mask_swaps_which_half_moves() {
    let Some(ctx) = ctx() else {
        eprintln!("no adapter; skipping");
        return;
    };
    let field = split_field(&ctx);

    let mut layer = brighten(MaskSource::Regions {
        signature: 1,
        level: 2,
        ids: vec![0],
    });
    layer.invert = true;

    let mut stack = MaskStack::new();
    stack.push(layer);

    let pixels = render(&ctx, &stack, Some(&field));
    let left = luma_at(&pixels, 4, SIZE / 2);
    let right = luma_at(&pixels, SIZE - 5, SIZE / 2);

    assert!(
        right > left + 40,
        "inverted, the *other* half should brighten: left {left}, right {right}"
    );
}

#[test]
fn opacity_scales_the_effect() {
    let Some(ctx) = ctx() else {
        eprintln!("no adapter; skipping");
        return;
    };
    let field = split_field(&ctx);
    let source = MaskSource::Regions {
        signature: 1,
        level: 2,
        ids: vec![0],
    };

    let mut full = MaskStack::new();
    full.push(brighten(source.clone()));

    let mut half = MaskStack::new();
    let mut layer = brighten(source);
    layer.opacity = 0.5;
    half.push(layer);

    let at_full = luma_at(&render(&ctx, &full, Some(&field)), 4, SIZE / 2);
    let at_half = luma_at(&render(&ctx, &half, Some(&field)), 4, SIZE / 2);
    let untouched = 128;

    assert!(
        at_half > untouched && at_half < at_full,
        "half opacity should land between neutral and full: {untouched} < {at_half} < {at_full}"
    );
}

#[test]
fn a_linear_gradient_ramps_across_the_frame() {
    let Some(ctx) = ctx() else {
        eprintln!("no adapter; skipping");
        return;
    };

    let mut stack = MaskStack::new();
    stack.push(brighten(MaskSource::Linear {
        centre: (0.5, 0.5),
        angle: 0.0,
        width: 1.0,
    }));

    let pixels = render(&ctx, &stack, None);
    let left = luma_at(&pixels, 1, SIZE / 2);
    let middle = luma_at(&pixels, SIZE / 2, SIZE / 2);
    let right = luma_at(&pixels, SIZE - 2, SIZE / 2);

    assert!(
        left < middle && middle < right,
        "a horizontal ramp should increase left to right: {left}, {middle}, {right}"
    );
}

#[test]
fn a_radial_mask_is_strongest_at_its_centre() {
    let Some(ctx) = ctx() else {
        eprintln!("no adapter; skipping");
        return;
    };

    let mut stack = MaskStack::new();
    stack.push(brighten(MaskSource::Radial {
        centre: (0.5, 0.5),
        radii: (0.3, 0.3),
        angle: 0.0,
        feather: 0.5,
    }));

    let pixels = render(&ctx, &stack, None);
    let centre = luma_at(&pixels, SIZE / 2, SIZE / 2);
    let corner = luma_at(&pixels, 1, 1);

    assert!(
        centre > corner + 40,
        "the centre should carry the effect: {centre} vs corner {corner}"
    );
    assert!(
        (120..=136).contains(&corner),
        "outside the radius must be untouched, got {corner}"
    );
}

/// Two layers must not read each other's slice.
#[test]
fn stacked_layers_use_their_own_masks() {
    let Some(ctx) = ctx() else {
        eprintln!("no adapter; skipping");
        return;
    };
    let field = split_field(&ctx);

    let mut stack = MaskStack::new();
    // Left half up.
    stack.push(brighten(MaskSource::Regions {
        signature: 1,
        level: 2,
        ids: vec![0],
    }));
    // Right half down.
    let mut darken = MaskLayer::new(
        "m2",
        MaskSource::Regions {
            signature: 1,
            level: 2,
            ids: vec![1],
        },
    );
    darken.set_param("exposure", ParamId("exposure"), -2.0);
    stack.push(darken);

    let pixels = render(&ctx, &stack, Some(&field));
    let left = luma_at(&pixels, 4, SIZE / 2);
    let right = luma_at(&pixels, SIZE - 5, SIZE / 2);

    assert!(left > 150, "left should have brightened, got {left}");
    assert!(right < 100, "right should have darkened, got {right}");
}

// ---------------------------------------------------------------------------
// Brush strokes (ARCH §5.4)
// ---------------------------------------------------------------------------

const UNTOUCHED: u8 = 128;

fn luma_in(pixels: &[u8], w: u32, x: u32, y: u32) -> u8 {
    pixels[((y * w + x) * 4) as usize]
}

/// One gesture: whether it erases, its radius, its flow, and its path.
type Gesture = (bool, f32, f32, Vec<(f32, f32)>);

/// A brightening layer with the given gestures already painted onto it.
fn painted(gestures: &[Gesture]) -> MaskLayer {
    let mut layer = brighten(MaskSource::brush());
    for (erase, radius, flow, path) in gestures {
        layer.begin_stroke(*erase, *radius, 0.9, *flow);
        for &(x, y) in path {
            layer.extend_stroke(x, y);
        }
        layer.end_stroke();
    }
    layer
}

fn stack_of(layer: MaskLayer) -> MaskStack {
    let mut stack = MaskStack::new();
    stack.push(layer);
    stack
}

/// The whole feature, at its simplest: paint somewhere, and that is where the
/// adjustment lands.
///
/// Painted across the top rather than down the middle, because a mask drawn
/// upside down is symmetric about the middle and a centred stroke would not
/// notice — and the vertex shader that draws a stroke has to flip y to reach
/// clip space, which is exactly the kind of thing that is wrong once.
#[test]
fn a_stroke_paints_where_it_was_drawn_and_nowhere_else() {
    let Some(ctx) = ctx() else {
        eprintln!("no adapter; skipping");
        return;
    };

    let stack = stack_of(painted(&[(
        false,
        0.1,
        1.0,
        vec![(0.2, 0.25), (0.8, 0.25)],
    )]));
    let pixels = render(&ctx, &stack, None);

    let under = luma_in(&pixels, SIZE, SIZE / 2, SIZE / 4);
    let below = luma_in(&pixels, SIZE, SIZE / 2, SIZE * 3 / 4);

    assert!(
        under > UNTOUCHED + 40,
        "the stroke should have brightened the upper quarter, got {under}"
    );
    assert!(
        (120..=136).contains(&below),
        "the lower half was never painted and must be untouched, got {below}"
    );
}

/// The failure a bounding box that is not grown by the radius produces: a tap
/// has no extent at all, so its quad has no area and nothing is drawn. Silent,
/// and it looks exactly like a brush that ignores short gestures.
#[test]
fn a_tap_paints_a_dab() {
    let Some(ctx) = ctx() else {
        eprintln!("no adapter; skipping");
        return;
    };

    let stack = stack_of(painted(&[(false, 0.2, 1.0, vec![(0.5, 0.5)])]));
    let pixels = render(&ctx, &stack, None);

    let centre = luma_in(&pixels, SIZE, SIZE / 2, SIZE / 2);
    let corner = luma_in(&pixels, SIZE, 1, 1);
    assert!(centre > 180, "the dab should be there, got {centre}");
    assert!(
        (120..=136).contains(&corner),
        "and only there, got {corner}"
    );
}

/// Painting must be able to erase, or a mask is one mistake away from being
/// started again.
#[test]
fn an_erasing_stroke_takes_back_what_was_painted() {
    let Some(ctx) = ctx() else {
        eprintln!("no adapter; skipping");
        return;
    };

    let stack = stack_of(painted(&[
        (false, 0.25, 1.0, vec![(0.15, 0.5), (0.85, 0.5)]),
        (true, 0.12, 1.0, vec![(0.5, 0.5)]),
    ]));
    let pixels = render(&ctx, &stack, None);

    let erased = luma_in(&pixels, SIZE, SIZE / 2, SIZE / 2);
    let kept = luma_in(&pixels, SIZE, 3, SIZE / 2);

    assert!(
        (120..=136).contains(&erased),
        "the erased middle should be back to untouched grey, got {erased}"
    );
    assert!(
        kept > 180,
        "the ends of the stroke are still painted, got {kept}"
    );
}

/// Order is the mask. The same two gestures the other way round leave the
/// paint alone, and a rasteriser that composited by kind rather than by
/// sequence would give the same answer to both.
#[test]
fn erasing_before_painting_removes_nothing() {
    let Some(ctx) = ctx() else {
        eprintln!("no adapter; skipping");
        return;
    };

    let stack = stack_of(painted(&[
        (true, 0.12, 1.0, vec![(0.5, 0.5)]),
        (false, 0.25, 1.0, vec![(0.15, 0.5), (0.85, 0.5)]),
    ]));
    let pixels = render(&ctx, &stack, None);

    let middle = luma_in(&pixels, SIZE, SIZE / 2, SIZE / 2);
    assert!(
        middle > 180,
        "an erase before the paint has nothing to take away, got {middle}"
    );
}

/// A stroke that crosses itself must not build up where it did. Summing the
/// segments instead of taking the nearest would make every circle and every
/// scribble blotchy — and at full flow it would not show at all, which is why
/// this paints at half.
#[test]
fn a_stroke_that_doubles_back_does_not_build_up() {
    let Some(ctx) = ctx() else {
        eprintln!("no adapter; skipping");
        return;
    };

    let once = stack_of(painted(&[(false, 0.15, 0.5, vec![(0.1, 0.5), (0.9, 0.5)])]));
    let twice = stack_of(painted(&[(
        false,
        0.15,
        0.5,
        // Out to the right and back over the last third of itself.
        vec![(0.1, 0.5), (0.9, 0.5), (0.65, 0.5)],
    )]));

    let single = luma_in(&render(&ctx, &once, None), SIZE, SIZE * 3 / 4, SIZE / 2);
    let crossed = luma_in(&render(&ctx, &twice, None), SIZE, SIZE * 3 / 4, SIZE / 2);

    assert_eq!(
        single, crossed,
        "one pass of the brush, however many times the path went over it"
    );
}

/// Between gestures, though, paint does build up — that is what a flow below
/// one is for, and it is the same blend that lets an erase work.
#[test]
fn two_gestures_at_half_flow_build_up() {
    let Some(ctx) = ctx() else {
        eprintln!("no adapter; skipping");
        return;
    };

    let dab = (false, 0.2, 0.5, vec![(0.5, 0.5)]);
    let once = stack_of(painted(std::slice::from_ref(&dab)));
    let twice = stack_of(painted(&[dab.clone(), dab]));

    let single = luma_in(&render(&ctx, &once, None), SIZE, SIZE / 2, SIZE / 2);
    let doubled = luma_in(&render(&ctx, &twice, None), SIZE, SIZE / 2, SIZE / 2);

    assert!(
        doubled > single,
        "a second pass should deposit more: {single} then {doubled}"
    );
}

/// A brush whose dab is an ellipse is not a brush. The radius is a fraction of
/// the *shorter* edge, so on a frame twice as wide as it is tall a circle in
/// normalised coordinates would come out twice as wide as it is high.
#[test]
fn a_dab_is_round_on_a_frame_that_is_not_square() {
    let Some(ctx) = ctx() else {
        eprintln!("no adapter; skipping");
        return;
    };
    const W: u32 = 64;
    const H: u32 = 32;

    let stack = stack_of(painted(&[(false, 0.25, 1.0, vec![(0.5, 0.5)])]));
    let pixels = render_at(&ctx, &stack, None, W, H);

    let lit = |v: u8| v > 160;
    let across = (0..W)
        .filter(|&x| lit(luma_in(&pixels, W, x, H / 2)))
        .count();
    let down = (0..H)
        .filter(|&y| lit(luma_in(&pixels, W, W / 2, y)))
        .count();

    assert!(
        across > 4 && down > 4,
        "the dab should exist: {across}x{down}"
    );
    assert!(
        across.abs_diff(down) <= 2,
        "a dab must be as wide as it is tall, got {across} across and {down} down"
    );
}

/// Hardness is the edge, and the edge is what a brush is judged on. A hard
/// brush that faded like a soft one would make the control do nothing anyone
/// could see.
#[test]
fn hardness_decides_how_quickly_the_edge_falls_away() {
    let Some(ctx) = ctx() else {
        eprintln!("no adapter; skipping");
        return;
    };

    let edge = |hardness: f32| {
        let mut layer = brighten(MaskSource::brush());
        layer.begin_stroke(false, 0.4, hardness, 1.0);
        layer.extend_stroke(0.5, 0.5);
        layer.end_stroke();

        let pixels = render(&ctx, &stack_of(layer), None);
        // How many pixels along the centre row are neither fully painted nor
        // fully clear — the width of the transition. "Fully painted" is read
        // from the middle of the dab rather than assumed: +2 EV over mid grey
        // lands wherever the output transform puts it.
        let solid = luma_in(&pixels, SIZE, SIZE / 2, SIZE / 2);
        (0..SIZE)
            .filter(|&x| {
                let v = luma_in(&pixels, SIZE, x, SIZE / 2);
                v > UNTOUCHED + 8 && v < solid - 8
            })
            .count()
    };

    let soft = edge(0.0);
    let hard = edge(1.0);
    assert!(
        hard < soft,
        "a hard brush should transition in fewer pixels: hard {hard}, soft {soft}"
    );
    assert!(hard <= 4, "and it should be nearly a step, got {hard}");
}

/// The loud failure an unpainted mask can produce: empty inverts to
/// everything, so a layer created with invert already set would apply its
/// adjustment to the whole photograph before a stroke was made.
#[test]
fn an_inverted_brush_layer_with_no_strokes_changes_nothing() {
    let Some(ctx) = ctx() else {
        eprintln!("no adapter; skipping");
        return;
    };

    let mut layer = brighten(MaskSource::brush());
    layer.invert = true;
    let pixels = render(&ctx, &stack_of(layer), None);

    for (x, y) in [(1, 1), (SIZE / 2, SIZE / 2), (SIZE - 2, SIZE - 2)] {
        let v = luma_in(&pixels, SIZE, x, y);
        assert!(
            (120..=136).contains(&v),
            "an unpainted mask covers nothing, inverted or not; got {v} at {x},{y}"
        );
    }
}

/// A painted layer and a gradient in one stack must not read each other's
/// slice — the brush writes its slot through a different pipeline, which is
/// exactly where a slot could be got wrong without either alone noticing.
#[test]
fn a_brush_layer_and_a_gradient_keep_their_own_slices() {
    let Some(ctx) = ctx() else {
        eprintln!("no adapter; skipping");
        return;
    };

    let mut stack = MaskStack::new();
    stack.push(painted(&[(false, 0.15, 1.0, vec![(0.5, 0.15)])]));

    let mut darken = MaskLayer::new(
        "m2",
        MaskSource::Radial {
            centre: (0.5, 0.85),
            radii: (0.15, 0.15),
            angle: 0.0,
            feather: 0.1,
        },
    );
    darken.set_param("exposure", ParamId("exposure"), -2.0);
    stack.push(darken);

    let pixels = render(&ctx, &stack, None);
    let top = luma_in(&pixels, SIZE, SIZE / 2, SIZE * 3 / 20);
    let bottom = luma_in(&pixels, SIZE, SIZE / 2, SIZE * 17 / 20);

    assert!(top > 180, "the painted dab should have brightened: {top}");
    assert!(bottom < 100, "the radial should have darkened: {bottom}");
}

/// A neutral edit must render identically whether or not masks are bound —
/// otherwise merely *having* the feature would alter every unedited image.
#[test]
fn an_empty_stack_renders_exactly_as_the_unmasked_path() {
    let Some(ctx) = ctx() else {
        eprintln!("no adapter; skipping");
        return;
    };

    let plain = {
        let source = grey(&ctx);
        let mut adjust = AdjustPass::new(&ctx);
        let shader = EditGraph::default_chain().compose();
        adjust.render(&source, &shader, SIZE, SIZE).expect("render");
        adjust.export_pixels().expect("readback").0
    };

    let masked = render(&ctx, &MaskStack::new(), None);
    assert_eq!(plain, masked, "an empty mask stack must be a no-op");
}