molgfx-render 0.3.2

//! The per-frame uniform block.
//!
//! Byte-identical to the shader's frame uniforms; uploading it is the only
//! per-frame write a camera-only change performs.

use molgfx_core::{ClipSet, MAX_CLIP_PLANES, Material, Representation, ScalarVolume, SurfaceKind};
use molgfx_gpu::{Device, Queue};
use molgfx_math::{Aabb, Camera, Mat4, Projection, Vec3};

use super::probe_offsets::PROBE_SAMPLE_COUNT;
use overlay::overlay_uniforms;

mod overlay;

pub(super) const SURFACE_GRID_MAX_DIMENSION: u32 = 192;
/// Bounded hybrid traversal budget for a persistent grid.
///
/// Each iteration crosses at least one cell and empty-space distances skip
/// farther. Capping the rare near-surface miss at the grid's longest axis
/// prevents grazing fragments from dominating frame time.
fn march_steps(dimensions: [u32; 3]) -> u32 {
    dimensions
        .iter()
        .copied()
        .fold(2, u32::max)
        .min(SURFACE_GRID_MAX_DIMENSION)
}

/// Per-structure placement, shared by every representation of that structure.
#[repr(C)]
#[derive(Clone, Copy, Debug, bytemuck::Pod, bytemuck::Zeroable)]
pub(crate) struct ModelUniforms {
    /// World-from-model transform.
    pub model_to_world: Mat4,
    /// Model-from-world transform for local-space hierarchy traversal.
    pub world_to_model: Mat4,
    /// Previous world-from-model transform for object motion.
    pub previous_model_to_world: Mat4,
    /// Resident page per `EntityKind`, followed by alignment padding.
    pub pick_pages: [[u32; 4]; 3],
}

impl ModelUniforms {
    pub(super) fn new(
        model_to_world: Mat4,
        previous_model_to_world: Mat4,
        pick_pages: [u32; 9],
    ) -> Self {
        let mut aligned_pages = [[u32::MAX; 4]; 3];
        for (index, page) in pick_pages.into_iter().enumerate() {
            aligned_pages[index / 4][index % 4] = page;
        }
        Self {
            model_to_world,
            world_to_model: model_to_world.inverse(),
            previous_model_to_world,
            pick_pages: aligned_pages,
        }
    }
}

/// The per-frame camera block, laid out exactly as the shader declares it.
#[repr(C)]
#[derive(Clone, Copy, Debug, bytemuck::Pod, bytemuck::Zeroable)]
pub(crate) struct FrameUniforms {
    /// View-from-world.
    pub view: Mat4,
    /// World-from-view, used by fullscreen analytic ray generation.
    pub inv_view: Mat4,
    /// Clip-from-view, reversed depth.
    pub proj: Mat4,
    /// Clip-from-world.
    pub view_proj: Mat4,
    /// View-from-clip, used by screen-space lighting passes.
    pub inv_proj: Mat4,
    /// Previous jittered clip-from-current jittered clip.
    pub reprojection: Mat4,
    /// Previous jittered clip-from-world, used for true geometry motion.
    pub previous_view_proj: Mat4,
    /// Width, height, 1/width, 1/height.
    pub viewport: [f32; 4],
    /// History, quality and publication flags followed by sample index.
    pub temporal: [f32; 4],
    /// Silhouette, cavity and legacy relative depth-cue strengths followed by focus distance.
    pub illustration: [f32; 4],
    /// Explicit near distance, far distance, strength and reserved lane.
    pub depth_cue: [f32; 4],
    /// Non-photorealistic lane: cel-shading band count in x, spare in yzw.
    pub npr: [f32; 4],
    /// Focus distance, aperture scale, maximum blur radius and blade count.
    pub optics: [f32; 4],
    /// Shutter fraction and maximum motion-blur radius in pixels.
    pub motion_blur: [f32; 4],
    /// Projection kind in x and lateral sphere/frustum factors in yz.
    pub projection_kind: [f32; 4],
    /// Light-view transform used by the scene-fit shadow pass.
    pub shadow_view: Mat4,
    /// World-from-light-view transform used by analytic shadow intersections.
    pub shadow_inv_view: Mat4,
    /// Reversed-depth orthographic light projection.
    pub shadow_projection: Mat4,
    /// Light clip-from-world transform used by lighting PCF.
    pub shadow_view_proj: Mat4,
    /// Background top, bottom, glow and display-grade controls.
    pub atmosphere: [[f32; 4]; 6],
    /// Environment, key and fill illumination controls.
    pub lighting: [[f32; 4]; 8],
}

#[derive(Clone, Copy, Debug)]
pub(crate) struct TemporalFrame {
    pub(crate) jitter_pixels: [f32; 2],
    pub(crate) previous_view_proj: Option<Mat4>,
    pub(crate) shadow_view: Mat4,
    pub(crate) shadow_projection: Mat4,
    pub(crate) shadow_view_proj: Mat4,
    pub(crate) sample_index: u32,
    pub(crate) quality: bool,
    pub(crate) publication: bool,
    pub(crate) illustration: [f32; 4],
    pub(crate) depth_cue: [f32; 4],
    pub(crate) npr: [f32; 4],
    pub(crate) optics: [f32; 4],
    pub(crate) motion_blur: [f32; 4],
    pub(crate) atmosphere: [[f32; 4]; 6],
    pub(crate) lighting: [[f32; 4]; 8],
}

impl FrameUniforms {
    /// Builds the block for one frame.
    #[must_use]
    pub(crate) fn new(camera: &Camera, width: u32, height: u32, temporal: &TemporalFrame) -> Self {
        let view = camera.view();
        let base_proj = camera.projection.matrix();
        // Viewport dimensions fit in 16 bits on every supported device, so
        // the float conversion is exact.
        let dim = |d: u32| match u16::try_from(d) {
            Ok(v) => f32::from(v),
            Err(_) => f32::from(u16::MAX),
        };
        let (w, h) = (dim(width), dim(height));
        let jitter_ndc = molgfx_math::Vec3::new(
            2.0 * temporal.jitter_pixels[0] / w.max(1.0),
            -2.0 * temporal.jitter_pixels[1] / h.max(1.0),
            0.0,
        );
        let proj = Mat4::from_translation(jitter_ndc) * base_proj;
        let view_proj = proj * view;
        let history_valid = temporal.previous_view_proj.is_some();
        let previous = match temporal.previous_view_proj {
            Some(previous) => previous,
            None => view_proj,
        };
        Self {
            view,
            inv_view: view.inverse(),
            proj,
            view_proj,
            inv_proj: proj.inverse(),
            reprojection: previous * view_proj.inverse(),
            previous_view_proj: previous,
            shadow_view: temporal.shadow_view,
            shadow_inv_view: temporal.shadow_view.inverse(),
            shadow_projection: temporal.shadow_projection,
            shadow_view_proj: temporal.shadow_view_proj,
            viewport: [w, h, 1.0 / w.max(1.0), 1.0 / h.max(1.0)],
            temporal: [
                if history_valid { 1.0 } else { 0.0 },
                if temporal.quality { 1.0 } else { 0.0 },
                if temporal.publication { 1.0 } else { 0.0 },
                f32::from(crate::fallback(
                    u16::try_from(temporal.sample_index.min(u32::from(u16::MAX))),
                    u16::MAX,
                )),
            ],
            illustration: temporal.illustration,
            depth_cue: temporal.depth_cue,
            npr: temporal.npr,
            optics: temporal.optics,
            motion_blur: temporal.motion_blur,
            projection_kind: projection_parameters(camera.projection, proj),
            atmosphere: temporal.atmosphere,
            lighting: temporal.lighting,
        }
    }
}

fn projection_parameters(projection: Projection, matrix: Mat4) -> [f32; 4] {
    let orthographic = matches!(projection, Projection::Orthographic { .. });
    let factor = |scale: f32| {
        let scale = scale.abs();
        if orthographic {
            scale
        } else {
            scale.hypot(1.0)
        }
    };
    [
        if orthographic { 1.0 } else { 0.0 },
        factor(matrix.x_axis.x),
        factor(matrix.y_axis.y),
        0.0,
    ]
}

/// Per-representation parameters consumed by procedural shaders.
#[repr(C)]
#[derive(Clone, Copy, Debug, bytemuck::Pod, bytemuck::Zeroable)]
pub(super) struct RepresentationUniforms {
    /// Probe radius, level set, hit epsilon and minimum marching step.
    pub(super) surface: [f32; 4],
    /// Local-space grid minimum and padding.
    pub(super) grid_min: [f32; 4],
    /// Local-space cell size and padding.
    pub(super) grid_cell: [f32; 4],
    /// Surface kind, maximum steps, reentrant directions and presentation.
    pub(super) options: [u32; 4],
    /// Three grid dimensions followed by total cell count.
    pub(super) grid_size: [u32; 4],
    /// Point, surface-pattern and bond-wire dimensions in physical pixels.
    pub(super) visual: [f32; 4],
    /// World-space clipping plane equations.
    pub(super) clip_planes: [[f32; 4]; MAX_CLIP_PLANES],
    /// Active plane count followed by future cap flags.
    pub(super) clip_meta: [u32; 4],
    /// World-space position to caller scalar-grid coordinates.
    pub(super) overlay_world_to_voxel: [[f32; 4]; 4],
    /// The reciprocal contour interval in `x`, or zero for no contours.
    pub(super) overlay_contour: [f32; 4],
    /// The caller's scalar ramp, baked to a lookup table.
    pub(super) overlay_ramp: super::ramp_lut::RampLut,
    /// Grid dimensions followed by an enabled sentinel.
    pub(super) overlay_size: [u32; 4],
    /// Contour half-width and normal sampling offset.
    pub(super) overlay_visual: [f32; 4],
    /// Perceptual roughness, dielectric specular strength and reserved lanes.
    pub(super) material: [f32; 4],
    /// Global opacity followed by reserved presentation lanes.
    pub(super) presentation: [f32; 4],
}

impl RepresentationUniforms {
    #[cfg(test)]
    pub(super) fn new(
        representation: &Representation,
        bounds: Aabb,
        overlay_volume: Option<&ScalarVolume>,
    ) -> Self {
        Self::for_spacing(
            representation,
            bounds,
            overlay_volume,
            super::detail::FINEST_SURFACE_SPACING,
        )
    }

    /// Uniforms for a surface field sampled at `target_spacing` ångström, coarsened
    /// only where the grid would exceed its dimension cap.
    pub(super) fn for_spacing(
        representation: &Representation,
        bounds: Aabb,
        overlay_volume: Option<&ScalarVolume>,
        target_spacing: f32,
    ) -> Self {
        let gaussian = representation.params.surface_kind == SurfaceKind::Gaussian;
        let sigma = representation.params.gaussian_sigma.max(0.05);
        let isolevel = if gaussian {
            if representation.params.isolevel > 0.0 {
                representation.params.isolevel.max(0.001)
            } else {
                0.5
            }
        } else {
            representation.params.isolevel
        };
        let probe = if representation.params.surface_kind == SurfaceKind::VanDerWaals {
            0.0
        } else if gaussian {
            sigma * 4.0
        } else {
            representation.params.probe_radius.max(0.0)
        };
        let minimum = bounds.min - Vec3::splat(probe);
        let maximum = bounds.max + Vec3::splat(probe);
        let extent = maximum - minimum;
        let max_divisions = dimension_f32(SURFACE_GRID_MAX_DIMENSION.saturating_sub(1));
        let cell = (extent.max_element() / max_divisions).max(target_spacing);
        let dimensions = [
            axis_cells(extent.x, cell),
            axis_cells(extent.y, cell),
            axis_cells(extent.z, cell),
        ];
        let total = dimensions.iter().copied().fold(1u32, u32::saturating_mul);
        let overlay = overlay_uniforms(representation, overlay_volume);
        Self {
            surface: [probe, isolevel, if gaussian { sigma } else { 0.02 }, 0.02],
            grid_min: [minimum.x, minimum.y, minimum.z, 0.0],
            grid_cell: [cell, cell, cell, 0.0],
            options: [
                representation.params.surface_kind as u32,
                march_steps(dimensions),
                u32::from(PROBE_SAMPLE_COUNT),
                representation.params.surface_style as u32,
            ],
            grid_size: [dimensions[0], dimensions[1], dimensions[2], total],
            visual: [
                representation.params.point_size_pixels.max(1.0),
                representation.params.surface_pattern_spacing.max(0.05),
                representation.params.surface_pattern_width_pixels.max(0.25),
                if representation.kind == molgfx_core::RepresentationKind::Surface {
                    if gaussian {
                        0.0
                    } else {
                        representation.params.radius_scale.max(0.0)
                    }
                } else if representation.kind == molgfx_core::RepresentationKind::Lines {
                    representation.params.line_width_pixels.max(0.5)
                } else {
                    0.0
                },
            ],
            clip_planes: clip_planes(&representation.clipping),
            clip_meta: clip_meta(&representation.clipping),
            overlay_world_to_voxel: overlay.world_to_voxel,
            overlay_contour: overlay.contour,
            overlay_ramp: overlay.ramp,
            overlay_size: overlay.size,
            overlay_visual: overlay.visual,
            material: material_uniforms(representation.material),
            presentation: presentation_uniforms(representation),
        }
    }
}

pub(super) fn write_representation_uniforms<D: Device>(
    queue: &D::Queue,
    buffer: &D::Buffer,
    representation: &Representation,
    bounds: Aabb,
    overlay_volume: Option<&ScalarVolume>,
    target_spacing: f32,
) {
    let value =
        RepresentationUniforms::for_spacing(representation, bounds, overlay_volume, target_spacing);
    queue.write_buffer(buffer, 0, bytemuck::bytes_of(&value));
}

#[repr(C)]
#[derive(Clone, Copy, Debug, bytemuck::Pod, bytemuck::Zeroable)]
pub(super) struct ClipUniforms {
    pub(super) planes: [[f32; 4]; MAX_CLIP_PLANES],
    pub(super) meta: [u32; 4],
    pub(super) material: [f32; 4],
    pub(super) presentation: [f32; 4],
    pub(super) tube_mapping: [f32; 4],
    pub(super) tube: [f32; 4],
}

impl ClipUniforms {
    pub(super) fn for_material_and_clipping(material: Material, clipping: &ClipSet) -> Self {
        Self {
            planes: clip_planes(clipping),
            meta: clip_meta(clipping),
            material: material_uniforms(material),
            presentation: presentation_uniforms_for_material(material),
            tube_mapping: [0.0; 4],
            tube: [0.0; 4],
        }
    }

    pub(super) fn for_mesh(mesh: &molgfx_core::Mesh) -> Self {
        let mut value = Self::for_material_and_clipping(mesh.material(), &mesh.clipping());
        value.meta[2] = match mesh.face_visibility() {
            molgfx_core::FaceVisibility::DoubleSided => 0,
            molgfx_core::FaceVisibility::FrontOnly => 1,
            molgfx_core::FaceVisibility::BackOnly => 2,
        };
        value
    }

    pub(super) fn new(representation: &Representation) -> Self {
        let mut value = Self {
            planes: clip_planes(&representation.clipping),
            meta: clip_meta(&representation.clipping),
            material: material_uniforms(representation.material),
            presentation: presentation_uniforms(representation),
            tube_mapping: [0.0; 4],
            tube: [
                representation.params.tube_radius.abs().max(1.0e-6),
                0.0,
                0.0,
                0.0,
            ],
        };
        if matches!(
            representation.kind,
            molgfx_core::RepresentationKind::Trace | molgfx_core::RepresentationKind::Tube
        ) && let Some((domain, radii)) = representation
            .params
            .tube_radius_mapping
            .b_factor_parameters()
        {
            value.meta[3] = 1;
            value.tube_mapping = [domain[0], domain[1], radii[0], radii[1]];
        }
        value
    }
}

pub(super) fn material_uniforms(material: Material) -> [f32; 4] {
    let model = material.model_lanes();
    [
        material.perceptual_roughness(),
        material.specular_strength(),
        model[0],
        model[1],
    ]
}

/// Presentation lanes for a representation: opacity, then the soft-union blend
/// span when the style reads it.
///
/// The span rides in a presentation lane rather than a surface lane so the
/// surface-uniform block keeps its four-lane shape and every existing shader
/// layout is unchanged.
fn presentation_uniforms(representation: &Representation) -> [f32; 4] {
    let material = representation.material;
    let span = if representation.params.surface_style == molgfx_core::SurfaceStyle::SoftUnion {
        representation.params.blob_spread.max(0.0)
    } else {
        0.0
    };
    let opacity = f32::from(material.opacity_unorm8()) / 255.0;
    let mut lanes = presentation_lanes(material);
    lanes[0] = opacity;
    lanes[1] = span;
    lanes
}

/// Presentation lanes for a material alone, with no representation to read a
/// style-specific control from.
fn presentation_uniforms_for_material(material: Material) -> [f32; 4] {
    presentation_lanes(material)
}

fn presentation_lanes(material: Material) -> [f32; 4] {
    [f32::from(material.opacity_unorm8()) / 255.0, 0.0, 0.0, 0.0]
}

pub(super) fn clip_planes(clipping: &ClipSet) -> [[f32; 4]; MAX_CLIP_PLANES] {
    let mut output = [[0.0; 4]; MAX_CLIP_PLANES];
    for (index, plane) in clipping.planes().iter().enumerate() {
        output[index] = [plane.normal.x, plane.normal.y, plane.normal.z, plane.offset];
    }
    output
}

pub(super) fn clip_meta(clipping: &ClipSet) -> [u32; 4] {
    [
        crate::fallback(u32::try_from(clipping.planes().len()), 0),
        clipping.cap() as u32,
        0,
        0,
    ]
}

fn axis_cells(extent: f32, cell: f32) -> u32 {
    let mut cells = 2u32;
    while cells < SURFACE_GRID_MAX_DIMENSION
        && dimension_f32(cells.saturating_sub(1)) * cell < extent
    {
        cells += 1;
    }
    cells
}

fn dimension_f32(value: u32) -> f32 {
    f32::from(crate::fallback(u16::try_from(value), u16::MAX))
}

#[cfg(test)]
#[path = "uniforms_tests.rs"]
mod tests;