environment/training-source/source-snapshot/src/render.rs

//! Drawing the coloured feature grid to an eye buffer.
//!
//! Deliberately tiny: a full-screen triangle and a `grid x grid x 3` float
//! buffer. The renderer runs at the headset's refresh rate regardless of
//! how slowly inference produces new grids, so the expensive part of the
//! frame is never on the display path.

use blade_graphics as gpu;

use crate::pca::COMPONENTS;

/// Maximum eyes we allocate parameter slots for.
pub const MAX_EYES: usize = 2;

#[repr(C)]
#[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
struct Params {
    scale: [f32; 2],
    offset: [f32; 2],
    grid: u32,
    pad: [u32; 3],
}

/// Where a flat overlay lands in one eye's clip space.
///
/// Purely 2D. The only thing an eye's frustum contributes to a head-locked
/// image is where its axis sits and how wide the image should be — a scale
/// and an offset — so that is all this carries.
#[derive(Clone, Copy, Debug)]
pub struct EyeTransform {
    pub scale: [f32; 2],
    pub offset: [f32; 2],
}

impl EyeTransform {
    /// Fills the eye buffer exactly. Right for a flat window, wrong for a
    /// headset, where it stretches the image over the whole field of view
    /// and ignores the frustum's asymmetry.
    pub const FULLSCREEN: Self = Self {
        scale: [1.0, 1.0],
        offset: [0.0, 0.0],
    };

    /// Fill an eye's field of view, centred on a given direction rather
    /// than on the eye's axis.
    ///
    /// `centre_tan` is where the image's middle should sit, in tangent
    /// units: `(x/-z, y/-z)` of the target direction expressed in the eye's
    /// space. Zero is straight ahead and reproduces [`filling`].
    ///
    /// This is how the image gets to stay put while the head moves. The
    /// content is as old as the last completed inference — a tenth of a
    /// second here — so labelling it with the current head pose tells the
    /// compositor to compensate for nothing, and it slides around with the
    /// head. Feeding the direction the camera was *pointing when the frame
    /// was captured* moves the image opposite to head motion, which is what
    /// world-locked looks like.
    pub fn filling_at(fov: [f32; 4], centre_tan: [f32; 2]) -> Self {
        let [left, right, up, down] = fov;
        let (tan_l, tan_r) = (left.tan(), right.tan());
        let (tan_u, tan_d) = (up.tan(), down.tan());
        let (span_x, span_y) = (tan_r - tan_l, tan_u - tan_d);
        Self {
            // Same size as filling the buffer.
            scale: [1.0, 1.0],
            offset: [
                (2.0 * centre_tan[0] - (tan_r + tan_l)) / span_x,
                (2.0 * centre_tan[1] - (tan_u + tan_d)) / span_y,
            ],
        }
    }

    /// Fill an eye's field of view, recentred on its axis.
    ///
    /// The right default for passthrough parity: the system's own
    /// passthrough maps these cameras across your whole view, so matching it
    /// means spanning the frustum rather than inventing an angular size.
    ///
    /// Note this is *not* the same as filling the eye buffer. Headset frusta
    /// are asymmetric and mirror between the eyes, so an image painted
    /// edge-to-edge has its centre at a different angle in each eye, and the
    /// two will not fuse. Spanning the frustum symmetrically about the axis
    /// is what makes them agree.
    pub fn filling(fov: [f32; 4]) -> Self {
        let [left, right, up, down] = fov;
        let half_tan = [
            (right.tan() - left.tan()) * 0.5,
            (up.tan() - down.tan()) * 0.5,
        ];
        Self::for_eye(fov, half_tan)
    }

    /// Place an image into an eye whose frustum is `[left, right, up, down]`
    /// radians, as OpenXR reports it.
    ///
    /// `half_tan` is the tangent of the image's half-angle on each axis.
    /// Separate axes matter: the camera frame is squashed into a square for
    /// the encoder, and giving x and y their true angular extents here is
    /// what unsquashes it, so the view keeps the camera's full field
    /// without looking stretched.
    ///
    /// A ray at angle θ from the view axis lands at clip
    /// `(2·tanθ − (tanR + tanL)) / (tanR − tanL)`, so the quad's ±1 edges
    /// map to a scale of `2·half_tan / (tanR − tanL)` about an offset of
    /// `−(tanR + tanL) / (tanR − tanL)`.
    ///
    /// Headset frusta are asymmetric and differ between eyes, which is why
    /// the offset is not zero and not shared: dropping it is what leaves the
    /// two views unfusable.
    pub fn for_eye(fov: [f32; 4], half_tan: [f32; 2]) -> Self {
        let [left, right, up, down] = fov;
        let (tan_l, tan_r) = (left.tan(), right.tan());
        let (tan_u, tan_d) = (up.tan(), down.tan());
        let (span_x, span_y) = (tan_r - tan_l, tan_u - tan_d);
        Self {
            scale: [
                2.0 * half_tan[0] / span_x,
                2.0 * half_tan[1] / span_y,
            ],
            offset: [-(tan_r + tan_l) / span_x, -(tan_u + tan_d) / span_y],
        }
    }
}

#[derive(blade_macros::ShaderData)]
struct ViewData {
    params: gpu::BufferPiece,
    cells: gpu::BufferPiece,
}

/// Renders the latest feature grid full-screen.
pub struct GridView {
    pipeline: gpu::RenderPipeline,
    params_buf: gpu::Buffer,
    cells_buf: gpu::Buffer,
    grid: usize,
}

impl GridView {
    /// `color_format` must match the eye buffer being rendered into —
    /// on an XR surface, `XrSurface::format()`.
    pub fn new(context: &gpu::Context, color_format: gpu::TextureFormat, grid: usize) -> Self {
        let shader = context.create_shader(gpu::ShaderDesc {
            source: include_str!("shaders/grid_view.wgsl"),
            naga_module: None,
        });
        let data_layout = <ViewData as gpu::ShaderData>::layout();
        let pipeline = context.create_render_pipeline(gpu::RenderPipelineDesc {
            name: "grid-view",
            data_layouts: &[&data_layout],
            vertex: shader.at("vs_main"),
            vertex_fetches: &[],
            primitive: gpu::PrimitiveState::default(),
            // The view covers every pixel, so there is nothing to depth-test
            // against and no depth buffer to allocate.
            depth_stencil: None,
            fragment: Some(shader.at("fs_main")),
            color_targets: &[gpu::ColorTargetState::from(color_format)],
            multisample_state: gpu::MultisampleState::default(),
        });

        // One slot per eye: the projection differs between them, and both
        // are in flight within a single frame's submission.
        let params_buf = context.create_buffer(gpu::BufferDesc {
            name: "grid-view-params",
            size: (std::mem::size_of::<Params>() * MAX_EYES) as u64,
            memory: gpu::Memory::Shared,
        });
        let cells_buf = context.create_buffer(gpu::BufferDesc {
            name: "grid-view-cells",
            size: (grid * grid * COMPONENTS * std::mem::size_of::<f32>()) as u64,
            memory: gpu::Memory::Shared,
        });

        let mut view = Self {
            pipeline,
            params_buf,
            cells_buf,
            grid,
        };
        // Full-screen until the caller supplies per-eye transforms, which
        // is right for a window and wrong for a headset.
        for eye in 0..MAX_EYES {
            view.set_transform(context, eye, EyeTransform::FULLSCREEN);
        }
        // A mid-grey field, so a failure to ever produce a result looks
        // like "no data" rather than like a working black-and-nothing view.
        view.upload(context, &vec![0.5; grid * grid * COMPONENTS]);
        view
    }

    /// Replace the displayed grid. `rgb` is `[grid * grid, 3]`, row-major.
    pub fn upload(&mut self, context: &gpu::Context, rgb: &[f32]) {
        let expected = self.grid * self.grid * COMPONENTS;
        debug_assert_eq!(rgb.len(), expected, "grid payload is the wrong size");
        let n = rgb.len().min(expected);
        unsafe {
            std::ptr::copy_nonoverlapping(rgb.as_ptr(), self.cells_buf.data() as *mut f32, n);
        }
        context.sync_buffer(self.cells_buf);
    }

    /// Set where the overlay lands in one eye.
    pub fn set_transform(&mut self, context: &gpu::Context, eye: usize, t: EyeTransform) {
        debug_assert!(eye < MAX_EYES);
        unsafe {
            let slot = (self.params_buf.data() as *mut Params).add(eye);
            *slot = Params {
                scale: t.scale,
                offset: t.offset,
                grid: self.grid as u32,
                pad: [0; 3],
            };
        }
        context.sync_buffer(self.params_buf);
    }

    /// Draw into an already-started render pass, using `eye`'s transform.
    pub fn draw(&self, pass: &mut gpu::RenderCommandEncoder, eye: usize) {
        let offset = (eye.min(MAX_EYES - 1) * std::mem::size_of::<Params>()) as u64;
        let mut encoder = pass.with(&self.pipeline);
        encoder.bind(
            0,
            &ViewData {
                params: self.params_buf.at(offset),
                cells: self.cells_buf.into(),
            },
        );
        encoder.draw(0, 6, 0, 1);
    }

    pub fn destroy(mut self, context: &gpu::Context) {
        context.destroy_buffer(self.cells_buf);
        context.destroy_buffer(self.params_buf);
        context.destroy_render_pipeline(&mut self.pipeline);
    }
}

/// Rotate a vector by a quaternion `[x, y, z, w]`.
fn rotate(q: [f32; 4], v: [f32; 3]) -> [f32; 3] {
    let (qx, qy, qz, qw) = (q[0], q[1], q[2], q[3]);
    // t = 2 * (q_vec × v); v' = v + qw * t + q_vec × t
    let tx = 2.0 * (qy * v[2] - qz * v[1]);
    let ty = 2.0 * (qz * v[0] - qx * v[2]);
    let tz = 2.0 * (qx * v[1] - qy * v[0]);
    [
        v[0] + qw * tx + qy * tz - qz * ty,
        v[1] + qw * ty + qz * tx - qx * tz,
        v[2] + qw * tz + qx * ty - qy * tx,
    ]
}

/// Where a direction that was straight ahead under `then` appears under
/// `now`, in tangent units suitable for [`EyeTransform::filling_at`].
///
/// Both quaternions are `[x, y, z, w]` in the same reference space. The
/// result is the image's centre after the head has moved, so a rotation to
/// the right pushes the content left — the image staying put in the world
/// while the view sweeps across it.
///
/// Returns `None` when the old direction has swung behind the viewer, where
/// a tangent-plane shift stops meaning anything.
pub fn reprojection_offset(then: [f32; 4], now: [f32; 4]) -> Option<[f32; 2]> {
    // Forward under the capture pose, brought into the current eye's space
    // by the inverse of the current pose.
    let forward = rotate(then, [0.0, 0.0, -1.0]);
    let inverse_now = [-now[0], -now[1], -now[2], now[3]];
    let v = rotate(inverse_now, forward);
    if v[2] >= -1e-3 {
        return None;
    }
    Some([v[0] / -v[2], v[1] / -v[2]])
}

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

    /// Straight ahead stays straight ahead when the head has not moved.
    #[test]
    fn no_motion_means_no_shift() {
        let identity = [0.0, 0.0, 0.0, 1.0];
        let offset = reprojection_offset(identity, identity).unwrap();
        assert!(offset[0].abs() < 1e-6 && offset[1].abs() < 1e-6);
    }

    /// Turning the head right must push the content left, so the image
    /// appears to stay where it was in the world.
    #[test]
    fn turning_right_pushes_content_left() {
        let identity = [0.0, 0.0, 0.0, 1.0];
        // +15 degrees about Y is a leftward yaw in a right-handed system
        // looking down -Z, so -15 turns the view to the right.
        let half = (-15.0f32).to_radians() * 0.5;
        let turned_right = [0.0, half.sin(), 0.0, half.cos()];
        let offset = reprojection_offset(identity, turned_right).unwrap();
        assert!(
            offset[0] < -0.1,
            "expected the image to move left, got {offset:?}"
        );
        // The magnitude should be tan(15 deg) = 0.268.
        assert!((offset[0] + 15.0f32.to_radians().tan()).abs() < 1e-3);
        assert!(offset[1].abs() < 1e-6, "yaw should not shift vertically");
    }

    /// Looking up must push the content down.
    #[test]
    fn looking_up_pushes_content_down() {
        let identity = [0.0, 0.0, 0.0, 1.0];
        let half = 10.0f32.to_radians() * 0.5;
        let looked_up = [half.sin(), 0.0, 0.0, half.cos()];
        let offset = reprojection_offset(identity, looked_up).unwrap();
        assert!(offset[1] < -0.1, "expected downward shift, got {offset:?}");
        assert!(offset[0].abs() < 1e-6);
    }

    /// A view swung right round has nothing sensible to show.
    #[test]
    fn facing_away_has_no_offset() {
        let identity = [0.0, 0.0, 0.0, 1.0];
        let behind = [0.0, 1.0, 0.0, 0.0];
        assert!(reprojection_offset(identity, behind).is_none());
    }
}