molgfx-render 0.3.2

//! Colour is presentation state, so changing a scheme never touches a record.

use super::tests::{camera, engine, represented_scene};
use crate::engine::Engine;
use crate::scene_gpu::color_uniforms;
use crate::testing::MockDevice;
use molgfx_core::ColorScheme;

use molgfx_math::Rgba8;

/// Buffer writes observed since `before`, with their creation labels.
fn writes_since(engine: &Engine<MockDevice>, before: usize) -> Vec<&'static str> {
    let Ok(writes) = engine.device.log.writes.lock() else {
        panic!("write log lock")
    };
    let Ok(buffers) = engine.device.log.buffers.lock() else {
        panic!("buffer log lock")
    };
    writes[before..]
        .iter()
        .map(|(buffer, _, _, _)| {
            buffers
                .iter()
                .find(|(id, _, _)| id == buffer)
                .map_or("unknown", |(_, label, _)| *label)
        })
        .collect()
}

/// Renders one frame, returning the writes it performed.
fn frame_writes(engine: &mut Engine<MockDevice>, scene: &molgfx_core::Scene) -> Vec<&'static str> {
    let Ok(writes) = engine.device.log.writes.lock() else {
        panic!("write log lock")
    };
    let before = writes.len();
    drop(writes);
    if let Err(error) = engine.render(scene, &camera()) {
        panic!("frame renders: {error}")
    }
    writes_since(engine, before)
}

/// A category column of zeros over the fixture's first structure.
fn category_column(scene: &mut molgfx_core::Scene) -> molgfx_core::AtomPropertyHandle {
    let Some((structure, placed)) = scene.structures().next() else {
        panic!("fixture has a structure")
    };
    let atoms = placed.atoms.len() as usize;
    let column = match molgfx_core::AtomProperty::new(
        structure,
        std::sync::Arc::<str>::from("categories"),
        vec![0.0; atoms].into(),
        molgfx_core::AtomPropertyMeaning::Generic,
        molgfx_core::ScalarFieldSemantics::UncalibratedRank,
    ) {
        Ok(property) => property,
        Err(error) => panic!("category column builds: {error}"),
    };
    match scene.add_atom_property(column) {
        Ok(handle) => handle,
        Err(error) => panic!("category column binds: {error}"),
    }
}

#[test]
fn a_colour_scheme_change_writes_only_fixed_size_uniforms() {
    // The scheme is resolved on the GPU from the record's element colour and
    // a property column, so changing it rewrites the presentation uniforms and
    // nothing else: no repack, no record upload, no new buffer.
    let mut scene = represented_scene(1, 1);
    let Some((representation, _)) = scene.representations().next() else {
        panic!("fixture has a representation")
    };
    let column = category_column(&mut scene);
    let mut engine = engine();
    if let Err(error) = engine.render(&scene, &camera()) {
        panic!("initial frame renders: {error}")
    }
    let Ok(buffers_before) = engine.device.log.buffers.lock().map(|value| value.len()) else {
        panic!("buffer log lock")
    };

    for scheme in [
        ColorScheme::category(column, molgfx_core::CategoryPalette::Kelly),
        ColorScheme::category(column, molgfx_core::CategoryPalette::MoleculeType),
        ColorScheme::Uniform(Rgba8::opaque(12, 34, 56)),
    ] {
        let Some(value) = scene.representation_mut(representation) else {
            panic!("representation resolves")
        };
        value.color = scheme;
        let changed = frame_writes(&mut engine, &scene);
        assert!(
            // A newly referenced column uploads once into the property table;
            // records, culling state and indirect arguments are untouched.
            changed.iter().all(|label| *label == "frame uniforms"
                || *label == "representation uniforms"
                || *label == "colour scheme uniforms"
                || *label == "visual property table"),
            "{scheme:?} wrote unexpected buffers: {changed:?}"
        );
    }

    let Ok(buffers_after) = engine.device.log.buffers.lock().map(|value| value.len()) else {
        panic!("buffer log lock")
    };
    assert_eq!(
        buffers_after, buffers_before,
        "a scheme change must not allocate"
    );
}

#[test]
fn representations_differing_only_in_colour_share_one_record_set() {
    // Before colour moved to the GPU this was impossible: the scheme was baked
    // into every atom record, so two differently-coloured representations could
    // not share one packed set.
    let mut counts = Vec::new();
    for (index, scheme) in [
        ColorScheme::ByElement,
        ColorScheme::Uniform(Rgba8::opaque(200, 30, 40)),
    ]
    .into_iter()
    .enumerate()
    {
        let mut scene = represented_scene(1, 1);
        let Some((representation, _)) = scene.representations().next() else {
            panic!("fixture has a representation")
        };
        let Some(value) = scene.representation_mut(representation) else {
            panic!("representation resolves")
        };
        value.color = scheme;
        value.order = u16::try_from(index).unwrap_or(0);
        let mut engine = engine();
        if let Err(error) = engine.render(&scene, &camera()) {
            panic!("frame renders: {error}")
        }
        let Ok(buffers) = engine.device.log.buffers.lock() else {
            panic!("buffer log lock")
        };
        counts.push(
            buffers
                .iter()
                .filter(|(_, label, _)| *label == "atom instances")
                .count(),
        );
    }
    assert!(
        counts.iter().all(|count| *count == 1),
        "each scene allocates exactly one atom record set: {counts:?}"
    );
}

#[test]
fn the_gpu_palette_bank_matches_the_cpu_palettes() {
    // The shader resolves categories from the bank the block uploads. This pins
    // the arrangement it indexes: every palette back to back in `ALL` order, so
    // getting a base wrong silently draws chains in another palette's colours.
    let scene = represented_scene(1, 1);
    let Some((representation, _)) = scene.representations().next() else {
        panic!("fixture has a representation")
    };
    let Some(value) = scene.representation(representation) else {
        panic!("representation resolves")
    };
    let uniforms =
        color_uniforms::ColorUniforms::new(value, &color_uniforms::ResolvedColumns::NONE);
    for palette in molgfx_core::CategoryPalette::ALL {
        let base = color_uniforms::bank_base(palette);
        for (index, color) in palette.colors().iter().enumerate() {
            assert_eq!(
                uniforms.bank_probe(base + index),
                u32::from_le_bytes([color.r, color.g, color.b, color.a]),
                "{palette:?} colour {index}"
            );
        }
    }
    let total: usize = molgfx_core::CategoryPalette::ALL
        .into_iter()
        .map(molgfx_core::CategoryPalette::len)
        .sum();
    assert!(
        total <= color_uniforms::BANK_COLORS,
        "every palette fits the bank"
    );
}

#[test]
fn the_colour_block_is_written_with_the_active_scheme() {
    // The shader reads the scheme from the colour block, so a block that never
    // receives the selector leaves every scheme drawing the element colour.
    let mut scene = represented_scene(1, 1);
    let Some((representation, _)) = scene.representations().next() else {
        panic!("fixture has a representation")
    };
    let Some(value) = scene.representation_mut(representation) else {
        panic!("representation resolves")
    };
    value.color = ColorScheme::Uniform(Rgba8::opaque(12, 34, 56));
    let mut engine = engine();
    if let Err(error) = engine.render(&scene, &camera()) {
        panic!("frame renders: {error}")
    }
    let Ok(buffers) = engine.device.log.buffers.lock() else {
        panic!("buffer log lock")
    };
    let color_buffer = buffers
        .iter()
        .find(|(_, label, _)| *label == "colour scheme uniforms")
        .map(|(id, _, _)| *id);
    let Some(color_buffer) = color_buffer else {
        panic!("the colour block is allocated")
    };
    drop(buffers);
    let Ok(writes) = engine.device.log.write_payloads.lock() else {
        panic!("payload log lock")
    };
    let Ok(write_log) = engine.device.log.writes.lock() else {
        panic!("write log lock")
    };
    let Some(index) = write_log
        .iter()
        .position(|(buffer, _, _, _)| *buffer == color_buffer)
    else {
        panic!("the colour block is written at least once")
    };
    let Some(payload) = writes.get(index) else {
        panic!("the colour block write captured its bytes")
    };
    // The base rule comes first: tag, packed colour, column offset, stride.
    let selector_at = 0;
    let scheme = u32::from_le_bytes([
        payload[selector_at],
        payload[selector_at + 1],
        payload[selector_at + 2],
        payload[selector_at + 3],
    ]);
    let packed = u32::from_le_bytes([
        payload[selector_at + 4],
        payload[selector_at + 5],
        payload[selector_at + 6],
        payload[selector_at + 7],
    ]);
    assert_eq!(scheme, 1, "a uniform scheme selects the uniform rule");
    assert_eq!(
        packed,
        u32::from_le_bytes([12, 34, 56, 255]),
        "the colour travels packed"
    );
}

#[test]
fn an_overlay_packs_its_column_and_scheme_table_for_the_shader() {
    let mut scene = represented_scene(1, 1);
    let Some((structure, placed)) = scene.structures().next() else {
        panic!("fixture has a structure")
    };
    let atoms = placed.atoms.len() as usize;
    let classes = match molgfx_core::AtomProperty::new(
        structure,
        std::sync::Arc::<str>::from("classes"),
        vec![0.0; atoms].into(),
        molgfx_core::AtomPropertyMeaning::Generic,
        molgfx_core::ScalarFieldSemantics::UncalibratedRank,
    ) {
        Ok(property) => property,
        Err(error) => panic!("class column builds: {error}"),
    };
    let Ok(handle) = scene.add_atom_property(classes) else {
        panic!("class column binds")
    };
    let red = molgfx_math::Rgba8::opaque(255, 0, 0);
    let Ok(overlay) = molgfx_core::ColorOverlay::new(
        handle,
        &[
            ColorScheme::category(handle, molgfx_core::CategoryPalette::Kelly),
            ColorScheme::Uniform(red),
        ],
    ) else {
        panic!("overlay builds")
    };
    let Some((representation, _)) = scene.representations().next() else {
        panic!("fixture has a representation")
    };
    let Some(value) = scene.representation_mut(representation) else {
        panic!("representation resolves")
    };
    value.color_overlay = Some(overlay);
    let mut columns = color_uniforms::ResolvedColumns::NONE;
    columns.overlay = [40, 1];
    columns.schemes[0] = [48, 1];
    let packed = color_uniforms::ColorUniforms::new(value, &columns);
    let (header, rules) = packed.overlay_probe();
    assert_eq!(header, [40, 1, 2, 0]);
    // Class one is the category rule over its own column, class two the packed
    // uniform red.
    assert_eq!(rules[0], [3, 0, 48, 1]);
    assert_eq!(rules[1][0], 1);
    assert_eq!(rules[1][1], u32::from_le_bytes([255, 0, 0, 255]));
    // An unplanned class column disables the overlay instead of reading offset
    // zero.
    let disabled =
        color_uniforms::ColorUniforms::new(value, &color_uniforms::ResolvedColumns::NONE);
    assert_eq!(disabled.overlay_probe().0[0], 0);
}