indicatrix 0.5.0

//! wasm32-only async counterparts to [`GpuFrameRenderer::new`]/[`GpuFrameRenderer::accumulate`].
//!
//! A separate `impl` block, not `#[cfg]` branches inside the existing methods: a
//! browser's main thread must never block, so there is no OS thread for
//! `pollster::block_on` to park and no `Device::poll(Maintain::Wait)` to synchronously
//! drive a `Buffer::map_async` callback the way native's `compute::begin_map_read`/
//! `finish_map_read` do. Both methods below are genuinely `async fn`s that `.await`
//! `wgpu`'s own futures -- a different control-flow shape that would make non-wasm32
//! callers pay `async`/`.await` for no reason if forced into the same signatures.
//!
//! Also drops the double-buffered chunk overlap: `accumulate`'s "Overlapped chunk
//! pipeline" exists to keep the GPU busy while the CPU blocks mapping the PREVIOUS
//! chunk's readback. On wasm32 nothing blocks in the first place -- awaiting
//! [`map_read_async`] simply yields to the browser's microtask queue at zero CPU cost --
//! so there is no idle window to hide by queuing chunk i+1 early.
//! [`GpuFrameRenderer::accumulate_async`] dispatches, copies, and awaits each chunk in
//! turn; still chunked, just not pipelined two deep.

use glam::Vec3;
use wgpu::BufferUsages;

use crate::{
    optics::raytracer::{EnvironmentSource, environment::environment_white_balance},
    renderer::{
        buffers::{GpuCameraParams, GpuGemMaterial, GpuTransportParams, encode_facet_finishes},
        env_map_gpu::HdrEnvGpuData,
        gpu::{GpuContext, MEGAKERNEL_STORAGE_BUFFERS, compute},
    },
};

use super::{
    CHUNK_BUDGET_BYTES, GpuFrameError, GpuFrameScene, GpuPipelineKind, REDUCE_SHADER_SRC,
    SHADER_SRC, WORKGROUP_SIZE,
    bind_groups::{TransportDispatchArgs, build_bind_group},
    classify_material,
    dispatch::{chunk_pixels_for, environment_params},
    readback::GpuReduceParams,
    renderer::GpuFrameRenderer,
};

// clippy::future_not_send: every async fn below touches wgpu's web backend types, which
// hold browser-side Rc handles and are therefore !Send -- see
// GpuContext::acquire_async's identical allow. This impl block only exists on wasm32,
// where nothing is ever sent across threads, so the lint is meaningless here regardless.
#[allow(
    clippy::future_not_send,
    reason = "wasm32-unknown-unknown has no second thread to send a future to; see the \
              module-level comment above this impl block"
)]
impl GpuFrameRenderer {
    /// Async, wasm32-only counterpart to [`Self::new`]: acquires a GPU device and
    /// compiles `transport_main` without blocking the browser's main thread.
    ///
    /// Awaits [`GpuContext::acquire_async`] directly rather than going through
    /// [`GpuContext::acquire`]'s `pollster::block_on` wrapper -- see this `impl` block's
    /// doc comment. Pipeline compilation itself is synchronous on every target, so
    /// nothing else here needs `.await`.
    ///
    /// # Errors
    ///
    /// Same conditions as [`Self::new`]'s -- most commonly, the browser has no WebGPU
    /// support or no compatible adapter (Safari without the WebGPU flag, or a machine
    /// whose driver stack the browser declines to expose).
    pub async fn new_async() -> Result<Self, GpuFrameError> {
        let ctx = GpuContext::acquire_async()
            .await
            .map_err(GpuFrameError::Acquire)?;
        let available = ctx.device.limits().max_storage_buffers_per_shader_stage;
        if available < MEGAKERNEL_STORAGE_BUFFERS {
            return Err(GpuFrameError::DeviceLimits {
                needed: MEGAKERNEL_STORAGE_BUFFERS,
                available,
            });
        }
        let info = ctx.adapter.get_info();
        let adapter_label = format!("{} ({:?})", info.name, info.backend);
        let pipeline = compute::create_compute_pipeline(
            &ctx.device,
            "transport_main",
            SHADER_SRC,
            "transport_main",
        );
        let reduce_pipeline = compute::create_compute_pipeline(
            &ctx.device,
            "reduce_xyz_main",
            REDUCE_SHADER_SRC,
            "reduce_xyz_main",
        );
        Ok(Self {
            ctx,
            pipeline,
            pipeline_isotropic: None,
            pipeline_uniaxial: None,
            pipeline_biaxial: None,
            reduce_pipeline,
            adapter_label,
            outputs: [None, None],
            pixel_outputs: [None, None],
            staging: [None, None],
            params_buffers: [None, None],
            reduce_params_buffers: [None, None],
            scene_buffers: None,
            xyz_scratch: Vec::new(),
            chunk_budget_bytes: CHUNK_BUDGET_BYTES,
            ns_per_tuple_ema: None,
            pipeline_kind: GpuPipelineKind::default(),
            wavefront_pipelines: None,
            poisoned: false,
            turn_in_progress: false,
        })
    }

    /// Async, wasm32-only counterpart to [`Self::accumulate`] -- same inputs, same
    /// summed-XYZ-into-`accum` contract, same decline conditions
    /// ([`GpuFrameError::UnsupportedMaterial`]/[`GpuFrameError::UnsupportedEnvironment`]),
    /// but chunk-at-a-time rather than two-deep pipelined -- see this `impl` block's own
    /// doc comment for why.
    ///
    /// # Errors
    ///
    /// Same conditions as [`Self::accumulate`]'s.
    ///
    /// # Panics
    ///
    /// Same conditions as [`Self::accumulate`]'s.
    pub async fn accumulate_async(
        &mut self,
        scene: &GpuFrameScene<'_>,
        sample_offset: u32,
        spp: u32,
        accum: &mut [Vec3],
    ) -> Result<(), GpuFrameError> {
        let pipeline_class = classify_material(scene.material);
        let num_pixels = scene.width as usize * scene.height as usize;
        assert_eq!(
            accum.len(),
            num_pixels,
            "accumulation buffer must have one entry per pixel"
        );
        if spp == 0 || num_pixels == 0 {
            return Ok(());
        }

        // See Self::accumulate_via_pipeline's identical check: unreachable today, kept as
        // defensive future-proofing.
        if !scene.material.gpu_supported() {
            return Err(GpuFrameError::UnsupportedMaterial(
                scene.material.name.clone(),
            ));
        }

        let (
            env_mode,
            temp_k,
            spot_mult,
            exposure,
            light_yaw,
            light_pitch,
            use_d65,
            studio_model,
            backdrop,
        ) = environment_params(scene.environment);
        // See `Self::prepare_turn`.
        let white_balance = environment_white_balance(scene.environment);

        let gpu_material = GpuGemMaterial::encode(scene.material);
        let gpu_finishes = encode_facet_finishes(scene.facet_finishes, scene.planes.len());
        // Built once for the whole call, like `gpu_material`/`gpu_finishes` above -- this
        // path never had `FrameSceneBuffers`'s cross-call persistence (see this `impl`
        // block's own doc comment), so there is no identity cache to consult here, only
        // "once per call rather than once per chunk".
        // Declines GpuFrameError::UnsupportedEnvironment for an oversized
        // HDR map -- mirrors `build_hdr_env`'s identical check (this path never had
        // `FrameSceneBuffers`'s cross-call persistence to share that helper through, see
        // this `impl` block's own doc comment, so the check is repeated here directly).
        let hdr_env = match scene.environment {
            EnvironmentSource::HdrMap(map) => {
                if !HdrEnvGpuData::fits_storage_binding(&self.ctx.device, map) {
                    return Err(GpuFrameError::UnsupportedEnvironment);
                }
                HdrEnvGpuData::upload(&self.ctx.device, map)
            }
            EnvironmentSource::Studio { .. } => HdrEnvGpuData::dummy(&self.ctx.device),
        };

        self.ensure_specialized_pipeline(pipeline_class);

        let chunk_pixels = chunk_pixels_for(self.chunk_budget_bytes, spp, num_pixels);
        self.ensure_capacity(chunk_pixels * spp as usize);
        self.ensure_pixel_capacity(chunk_pixels);

        let mut first_pixel = 0usize;
        let mut chunk_index = 0usize;
        while first_pixel < num_pixels {
            let pixels_this_chunk = chunk_pixels.min(num_pixels - first_pixel);
            let tuples = pixels_this_chunk * spp as usize;

            let camera_params = GpuCameraParams {
                origin: scene.camera.origin.to_array(),
                fov_tan: scene.camera.fov_tan,
                forward: scene.camera.forward.to_array(),
                width: scene.width as f32,
                right: scene.camera.right.to_array(),
                height: scene.height as f32,
                up: scene.camera.up.to_array(),
                num_samples: spp,
            };
            let params = GpuTransportParams::new(
                pixels_this_chunk as u32,
                scene.max_bounces,
                sample_offset,
                env_mode,
                0.0,
                temp_k,
                spot_mult,
                exposure,
                light_yaw,
                light_pitch,
                white_balance.to_array(),
            )
            .with_pixel_offset(first_pixel as u32)
            .with_debug_buffers_disabled()
            .with_studio_use_d65(use_d65)
            .with_studio_model(studio_model)
            .with_backdrop(backdrop);

            let outputs = self.outputs[chunk_index % 2]
                .as_ref()
                .expect("ensure_capacity just populated both slots");

            let bind_args = TransportDispatchArgs {
                ctx: &self.ctx,
                pipeline: self.pipeline_for_class(pipeline_class),
                camera_params: &camera_params,
                params: &params,
                material: &gpu_material,
                planes: scene.planes,
                facet_finishes: &gpu_finishes,
                outputs,
                hdr_env: &hdr_env,
            };

            // Non-blocking submit -- native's pipelined path builds its bind group via
            // build_chunk_bind_group instead (see the parent module's doc comment's
            // "Per-frame uploads, persistent staging" section), but this wasm32 loop
            // re-uploads every buffer every chunk via build_bind_group: it never had the
            // overlapped double-buffering the native path uses, so there is no per-frame
            // state to hoist these uploads out of here.
            let bind_group = build_bind_group(&bind_args);
            let workgroups = (tuples as u32).div_ceil(WORKGROUP_SIZE as u32);
            let _ = compute::dispatch(
                &self.ctx.device,
                &self.ctx.queue,
                bind_args.pipeline,
                &bind_group,
                (workgroups, 1, 1),
            );

            // GPU-side sample reduction, mirroring native's `dispatch_chunk` -- see that
            // function's "Why a second dispatch" doc comment. This wasm32
            // loop re-creates the tiny reduce-params buffer fresh every chunk (like every
            // other buffer here, per this function's own doc comment: it never had the
            // overlapped double-buffering the native path uses, so there is no per-frame
            // state to persist it in).
            let pixel_output = self.pixel_outputs[chunk_index % 2]
                .as_ref()
                .expect("ensure_pixel_capacity just populated both slots");
            let reduce_params = GpuReduceParams {
                num_pixels: pixels_this_chunk as u32,
                num_samples: spp,
                _pad0: 0,
                _pad1: 0,
            };
            let reduce_params_buf = compute::upload(
                &self.ctx.device,
                "reduce xyz params (wasm32 async)",
                std::slice::from_ref(&reduce_params),
                BufferUsages::UNIFORM,
            );
            let reduce_bind_group = compute::bind_buffers(
                &self.ctx.device,
                "reduce xyz bind group (wasm32 async)",
                &self.reduce_pipeline,
                &[
                    (0, &reduce_params_buf),
                    (1, outputs.xyz()),
                    (2, &pixel_output.buffer),
                ],
            );
            let reduce_workgroups = (pixels_this_chunk as u32).div_ceil(WORKGROUP_SIZE as u32);
            let _ = compute::dispatch(
                &self.ctx.device,
                &self.ctx.queue,
                &self.reduce_pipeline,
                &reduce_bind_group,
                (reduce_workgroups, 1, 1),
            );

            let (staging, _copy_index) = compute::copy_to_staging::<f32>(
                &self.ctx.device,
                &self.ctx.queue,
                &pixel_output.buffer,
                pixels_this_chunk * 3,
                "transport out pixel xyz staging (wasm32 async)",
            );

            // Awaits the browser's own resolution of map_async -- see map_read_async's
            // doc comment for why this needs no Device::poll call, unlike native's
            // finish_map_read. Returns GpuFrameError::DeviceLost on a mapping failure,
            // propagated here with `?`.
            let xyz: Vec<f32> = map_read_async(&staging, pixels_this_chunk * 3).await?;

            // xyz already holds reduce_xyz_main's GPU-summed per-pixel triples -- no
            // CPU-side per-sample summation loop needed any more.
            for local_pixel in 0..pixels_this_chunk {
                let base = local_pixel * 3;
                accum[first_pixel + local_pixel] +=
                    Vec3::new(xyz[base], xyz[base + 1], xyz[base + 2]);
            }

            first_pixel += pixels_this_chunk;
            chunk_index += 1;
        }

        Ok(())
    }
}

/// Awaits `staging`'s `Buffer::map_async` callback and reads back `count` `T`s, wasm32
/// only.
///
/// Native's `compute::finish_map_read` drives the identical callback to completion via
/// `Device::poll(Maintain::Wait)`, blocking the calling OS thread. `wgpu`'s WebGPU
/// backend needs no such call: a browser resolves `GPUBuffer.mapAsync`'s promise on its
/// own microtask queue as soon as ready -- the callback fires whenever this function
/// `.await`s the channel below, which is why this can be a genuine `async fn` instead of
/// needing a `Device::poll` equivalent.
///
/// # Errors
///
/// Returns [`GpuFrameError::DeviceLost`] on a dropped callback, a failed mapping, or an
/// unreadable mapped range, for [`GpuFrameRenderer::accumulate_async`] to propagate with
/// `?`. There is no bounded-wait timeout here the way `compute::GPU_WAIT_TIMEOUT` bounds
/// the native path: nothing here blocks an OS thread, so there is no thread to free by
/// timing out; a browser tab that never resolves `mapAsync` leaves this `Future` pending.
#[allow(
    clippy::future_not_send,
    reason = "wgpu's web backend's buffer-mapping state is inherently !Send (browser-side \
              Rc handles); wasm32-unknown-unknown has no second thread to send this future \
              to regardless -- see GpuContext::acquire_async's identical allow"
)]
async fn map_read_async<T: bytemuck::Pod>(
    staging: &wgpu::Buffer,
    count: usize,
) -> Result<Vec<T>, GpuFrameError> {
    let (tx, rx) = futures_channel::oneshot::channel();
    staging
        .slice(..)
        .map_async(wgpu::MapMode::Read, move |result| {
            // The receiver can only be dropped by this function returning early, which it
            // never does before this send -- a failed send would mean the callback
            // outlived its own awaiting task, not a race this single-threaded target can
            // produce.
            let _ = tx.send(result);
        });
    rx.await
        .map_err(|_| {
            GpuFrameError::DeviceLost(
                "wgpu buffer-mapping callback was dropped before it fired".to_string(),
            )
        })?
        .map_err(|e| GpuFrameError::DeviceLost(format!("wgpu buffer mapping failed: {e}")))?;
    // Mirrors compute::finish_map_read: staging was sized by copy_to_staging::<T> for
    // precisely `count` T's, so the mapped range's length already equals `count`; `count`
    // is taken as a parameter so the assertion below can catch a size mismatch.
    let out = {
        let data = staging.slice(..).get_mapped_range().map_err(|e| {
            GpuFrameError::DeviceLost(format!("mapped buffer range could not be read: {e}"))
        })?;
        bytemuck::cast_slice::<u8, T>(&data).to_vec()
    };
    debug_assert_eq!(
        out.len(),
        count,
        "staging buffer length must match copy_to_staging's count"
    );
    staging.unmap();
    Ok(out)
}