spatialrust_render_wgpu/
render.rs

1use std::collections::HashMap;
2
3use bytemuck::{Pod, Zeroable};
4use spatialrust_gpu::WgpuRuntime;
5use spatialrust_viz::{
6    Camera, ColorMap, DeviceIdentity, LinearRgba, PointColor, Projection, TransferDirection,
7    TransferEvent, TransferReceipt, VisualResidency, VisualStyle,
8};
9use wgpu::util::DeviceExt;
10
11use crate::{GpuGeometry, GpuGeometryKind, RenderError, RenderResult, WgpuRenderer};
12
13/// Portable color format used by headless render targets.
14pub const RENDER_TARGET_FORMAT: wgpu::TextureFormat = wgpu::TextureFormat::Rgba8Unorm;
15const DEPTH_FORMAT: wgpu::TextureFormat = wgpu::TextureFormat::Depth32Float;
16
17/// Configuration for one device-resident headless render.
18#[derive(Clone, Debug, PartialEq)]
19pub struct RenderOptions {
20    /// Output width in pixels.
21    pub width: u32,
22    /// Output height in pixels.
23    pub height: u32,
24    /// Camera used for this render.
25    pub camera: Camera,
26    /// Primitive presentation style.
27    pub style: VisualStyle,
28    /// Linear clear color.
29    pub clear_color: LinearRgba,
30}
31
32impl RenderOptions {
33    /// Creates options with a validated non-zero target size.
34    pub fn try_new(
35        width: u32,
36        height: u32,
37        camera: Camera,
38        style: VisualStyle,
39        clear_color: LinearRgba,
40    ) -> RenderResult<Self> {
41        if width == 0 || height == 0 {
42            return Err(RenderError::GeometrySize(
43                "render target dimensions must be non-zero".into(),
44            ));
45        }
46        Ok(Self { width, height, camera, style, clear_color })
47    }
48}
49
50/// Named device-side execution evidence for one render.
51#[derive(Clone, Debug, PartialEq, Eq)]
52pub struct RenderReceipt {
53    /// Exact renderer adapter identity.
54    pub adapter: DeviceIdentity,
55    /// Named stages in submission order.
56    pub stages: Vec<&'static str>,
57    /// Number of submitted draw calls.
58    pub draw_calls: u32,
59    /// Number of vertices or point instances consumed.
60    pub element_count: u32,
61}
62
63/// Device-resident color and depth attachments.
64pub struct GpuRenderTarget {
65    pub(crate) renderer_id: u64,
66    pub(crate) color: wgpu::Texture,
67    pub(crate) depth: wgpu::Texture,
68    point_ids: Option<wgpu::Texture>,
69    width: u32,
70    height: u32,
71    residency: VisualResidency,
72}
73
74/// Caller-requested tightly packed RGBA8 readback.
75#[derive(Debug)]
76pub struct ReadbackImage {
77    /// Image width.
78    pub width: u32,
79    /// Image height.
80    pub height: u32,
81    /// Tightly packed row-major RGBA8 pixels.
82    pub rgba: Vec<u8>,
83    /// Exact device-to-host transfer receipt.
84    pub transfers: TransferReceipt,
85}
86
87/// Caller-requested point-picking result.
88#[derive(Debug)]
89pub struct PickResult {
90    /// Uploaded point index, or `None` for the cleared background.
91    pub point_index: Option<u32>,
92    /// Exact four-byte device-to-host transfer receipt.
93    pub transfers: TransferReceipt,
94}
95
96impl GpuRenderTarget {
97    /// Target width in pixels.
98    #[must_use]
99    pub const fn width(&self) -> u32 {
100        self.width
101    }
102
103    /// Target height in pixels.
104    #[must_use]
105    pub const fn height(&self) -> u32 {
106        self.height
107    }
108
109    /// Device residency of both attachments.
110    #[must_use]
111    pub const fn residency(&self) -> &VisualResidency {
112        &self.residency
113    }
114
115    /// Returns whether this target belongs to the supplied renderer runtime.
116    #[must_use]
117    pub fn is_owned_by(&self, renderer: &WgpuRenderer) -> bool {
118        self.renderer_id == renderer.id
119    }
120}
121
122/// Result of a device-resident headless render.
123pub struct HeadlessRender {
124    /// Device-resident target. No readback has occurred.
125    pub target: GpuRenderTarget,
126    /// Explicit transfers required to submit the render.
127    pub transfers: TransferReceipt,
128    /// Named render stages and draw counts.
129    pub receipt: RenderReceipt,
130}
131
132#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
133enum ColorMode {
134    Uniform,
135    Rgb,
136    Scalar,
137}
138
139#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
140struct PipelineKey {
141    kind: GpuGeometryKind,
142    color: ColorMode,
143}
144
145#[derive(Default)]
146pub(crate) struct RenderPipelines {
147    bind_group_layout: Option<wgpu::BindGroupLayout>,
148    pipelines: HashMap<PipelineKey, wgpu::RenderPipeline>,
149    point_id_pipeline: Option<wgpu::RenderPipeline>,
150}
151
152#[repr(C)]
153#[derive(Clone, Copy, Debug, Pod, Zeroable)]
154struct RenderUniform {
155    view_projection: [[f32; 4]; 4],
156    color: [f32; 4],
157    scalar_range: [f32; 2],
158    viewport: [f32; 2],
159    point_size: f32,
160    color_map: u32,
161    _padding: [u32; 2],
162}
163
164impl WgpuRenderer {
165    /// Renders uploaded geometry into device-resident color and depth textures.
166    ///
167    /// The method never reads the target back to the host. Its only host/device
168    /// crossing is the returned, byte-accounted uniform upload.
169    pub fn render_headless(
170        &self,
171        geometry: &GpuGeometry,
172        options: &RenderOptions,
173    ) -> RenderResult<HeadlessRender> {
174        if geometry.renderer_id != self.id {
175            return Err(RenderError::RuntimeMismatch(
176                "geometry was uploaded by another renderer".into(),
177            ));
178        }
179        let max_dimension = self.runtime.device().limits().max_texture_dimension_2d;
180        if options.width > max_dimension || options.height > max_dimension {
181            return Err(RenderError::GeometrySize(format!(
182                "render target {}x{} exceeds adapter limit {max_dimension}",
183                options.width, options.height
184            )));
185        }
186        let (color_mode, uniform_color, scalar_range, point_size, color_map) =
187            resolve_style(geometry, &options.style)?;
188        let key = PipelineKey { kind: geometry.kind, color: color_mode };
189        let view_projection =
190            view_projection(options.camera, options.width as f32 / options.height as f32);
191        let uniform = RenderUniform {
192            view_projection,
193            color: [
194                uniform_color.red,
195                uniform_color.green,
196                uniform_color.blue,
197                uniform_color.alpha,
198            ],
199            scalar_range,
200            viewport: [options.width as f32, options.height as f32],
201            point_size,
202            color_map,
203            _padding: [0; 2],
204        };
205        let uniform_bytes = bytemuck::bytes_of(&uniform);
206        let uniform_buffer =
207            self.runtime.device().create_buffer_init(&wgpu::util::BufferInitDescriptor {
208                label: Some("spatialrust render uniform"),
209                contents: uniform_bytes,
210                usage: wgpu::BufferUsages::UNIFORM,
211            });
212        let mut pipelines = self.render_pipelines.lock().expect("render pipeline cache poisoned");
213        let RenderPipelines { bind_group_layout, pipelines: pipeline_map, point_id_pipeline } =
214            &mut *pipelines;
215        if bind_group_layout.is_none() {
216            *bind_group_layout = Some(self.runtime.device().create_bind_group_layout(
217                &wgpu::BindGroupLayoutDescriptor {
218                    label: Some("spatialrust render uniform layout"),
219                    entries: &[wgpu::BindGroupLayoutEntry {
220                        binding: 0,
221                        visibility: wgpu::ShaderStages::VERTEX_FRAGMENT,
222                        ty: wgpu::BindingType::Buffer {
223                            ty: wgpu::BufferBindingType::Uniform,
224                            has_dynamic_offset: false,
225                            min_binding_size: None,
226                        },
227                        count: None,
228                    }],
229                },
230            ));
231        }
232        let bind_group_layout = bind_group_layout.as_ref().expect("bind group layout was inserted");
233        let bind_group = self.runtime.device().create_bind_group(&wgpu::BindGroupDescriptor {
234            label: Some("spatialrust render uniform bind group"),
235            layout: bind_group_layout,
236            entries: &[wgpu::BindGroupEntry {
237                binding: 0,
238                resource: uniform_buffer.as_entire_binding(),
239            }],
240        });
241
242        let target = create_target(
243            self.runtime.device(),
244            self.id,
245            options.width,
246            options.height,
247            self.device_identity().clone(),
248            geometry.kind == GpuGeometryKind::Points,
249        );
250        let color_view = target.color.create_view(&wgpu::TextureViewDescriptor::default());
251        let depth_view = target.depth.create_view(&wgpu::TextureViewDescriptor::default());
252        let pipeline = pipeline_map
253            .entry(key)
254            .or_insert_with(|| create_pipeline(self.runtime.device(), bind_group_layout, key));
255        if geometry.kind == GpuGeometryKind::Points && point_id_pipeline.is_none() {
256            *point_id_pipeline =
257                Some(create_point_id_pipeline(self.runtime.device(), bind_group_layout));
258        }
259
260        let mut encoder =
261            self.runtime.device().create_command_encoder(&wgpu::CommandEncoderDescriptor {
262                label: Some("spatialrust headless render encoder"),
263            });
264        {
265            let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
266                label: Some("spatialrust headless render pass"),
267                color_attachments: &[Some(wgpu::RenderPassColorAttachment {
268                    view: &color_view,
269                    resolve_target: None,
270                    ops: wgpu::Operations {
271                        load: wgpu::LoadOp::Clear(wgpu::Color {
272                            r: options.clear_color.red as f64,
273                            g: options.clear_color.green as f64,
274                            b: options.clear_color.blue as f64,
275                            a: options.clear_color.alpha as f64,
276                        }),
277                        store: wgpu::StoreOp::Store,
278                    },
279                })],
280                depth_stencil_attachment: Some(wgpu::RenderPassDepthStencilAttachment {
281                    view: &depth_view,
282                    depth_ops: Some(wgpu::Operations {
283                        load: wgpu::LoadOp::Clear(1.0),
284                        store: wgpu::StoreOp::Store,
285                    }),
286                    stencil_ops: None,
287                }),
288                timestamp_writes: None,
289                occlusion_query_set: None,
290            });
291            pass.set_pipeline(pipeline);
292            pass.set_bind_group(0, &bind_group, &[]);
293            pass.set_vertex_buffer(
294                0,
295                geometry
296                    .positions
297                    .as_ref()
298                    .expect("uploaded geometry always owns positions")
299                    .buffer
300                    .slice(..),
301            );
302            match color_mode {
303                ColorMode::Uniform => {}
304                ColorMode::Rgb => pass.set_vertex_buffer(
305                    1,
306                    geometry.rgb.as_ref().expect("RGB style was validated").buffer.slice(..),
307                ),
308                ColorMode::Scalar => pass.set_vertex_buffer(
309                    1,
310                    geometry.scalar.as_ref().expect("scalar style was validated").buffer.slice(..),
311                ),
312            }
313            match geometry.kind {
314                GpuGeometryKind::Points => pass.draw(0..6, 0..geometry.vertex_count),
315                GpuGeometryKind::Lines => pass.draw(0..geometry.vertex_count, 0..1),
316                GpuGeometryKind::Triangles => {
317                    let indices =
318                        geometry.indices.as_ref().expect("triangle geometry owns indices");
319                    pass.set_index_buffer(indices.buffer.slice(..), wgpu::IndexFormat::Uint32);
320                    pass.draw_indexed(0..geometry.index_count, 0, 0..1);
321                }
322            }
323        }
324        if geometry.kind == GpuGeometryKind::Points {
325            let id_view = target
326                .point_ids
327                .as_ref()
328                .expect("point target owns an ID texture")
329                .create_view(&wgpu::TextureViewDescriptor::default());
330            let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
331                label: Some("spatialrust point ID render pass"),
332                color_attachments: &[Some(wgpu::RenderPassColorAttachment {
333                    view: &id_view,
334                    resolve_target: None,
335                    ops: wgpu::Operations {
336                        load: wgpu::LoadOp::Clear(wgpu::Color::TRANSPARENT),
337                        store: wgpu::StoreOp::Store,
338                    },
339                })],
340                depth_stencil_attachment: Some(wgpu::RenderPassDepthStencilAttachment {
341                    view: &depth_view,
342                    depth_ops: Some(wgpu::Operations {
343                        load: wgpu::LoadOp::Load,
344                        store: wgpu::StoreOp::Store,
345                    }),
346                    stencil_ops: None,
347                }),
348                timestamp_writes: None,
349                occlusion_query_set: None,
350            });
351            pass.set_pipeline(
352                point_id_pipeline.as_ref().expect("point ID pipeline was initialized"),
353            );
354            pass.set_bind_group(0, &bind_group, &[]);
355            pass.set_vertex_buffer(
356                0,
357                geometry
358                    .positions
359                    .as_ref()
360                    .expect("uploaded geometry always owns positions")
361                    .buffer
362                    .slice(..),
363            );
364            pass.draw(0..6, 0..geometry.vertex_count);
365        }
366        self.runtime.queue().submit(Some(encoder.finish()));
367        drop(pipelines);
368
369        let mut transfers = TransferReceipt::new();
370        transfers.push(
371            TransferEvent::try_new(
372                "render-uniform-upload",
373                TransferDirection::Upload,
374                VisualResidency::Host,
375                VisualResidency::Device(self.device_identity().clone()),
376                uniform_bytes.len() as u64,
377            )
378            .map_err(|error| RenderError::Transfer(error.to_string()))?,
379        );
380        let stage = match geometry.kind {
381            GpuGeometryKind::Points => "draw-points",
382            GpuGeometryKind::Lines => "draw-lines",
383            GpuGeometryKind::Triangles => "draw-triangles",
384        };
385        let element_count = match geometry.kind {
386            GpuGeometryKind::Triangles => geometry.index_count,
387            _ => geometry.vertex_count,
388        };
389        Ok(HeadlessRender {
390            target,
391            transfers,
392            receipt: RenderReceipt {
393                adapter: self.device_identity().clone(),
394                stages: if geometry.kind == GpuGeometryKind::Points {
395                    vec!["clear-color-depth", stage, "draw-point-ids", "submit"]
396                } else {
397                    vec!["clear-color-depth", stage, "submit"]
398                },
399                draw_calls: if geometry.kind == GpuGeometryKind::Points { 2 } else { 1 },
400                element_count,
401            },
402        })
403    }
404
405    /// Number of initialized render-pipeline variants.
406    #[must_use]
407    pub fn initialized_render_pipeline_count(&self) -> usize {
408        let pipelines = self.render_pipelines.lock().expect("render pipeline cache poisoned");
409        pipelines.pipelines.len() + usize::from(pipelines.point_id_pipeline.is_some())
410    }
411
412    /// Explicitly reads one device-resident target into tightly packed RGBA8.
413    pub fn readback_rgba(&self, target: &GpuRenderTarget) -> RenderResult<ReadbackImage> {
414        self.validate_target(target)?;
415        let unpadded_row_bytes = target
416            .width
417            .checked_mul(4)
418            .ok_or_else(|| RenderError::Readback("RGBA row byte count overflowed".into()))?;
419        let padded_row_bytes = align_up(unpadded_row_bytes, wgpu::COPY_BYTES_PER_ROW_ALIGNMENT)?;
420        let buffer_bytes = u64::from(padded_row_bytes)
421            .checked_mul(u64::from(target.height))
422            .ok_or_else(|| RenderError::Readback("padded readback byte count overflowed".into()))?;
423        let buffer = self.runtime.device().create_buffer(&wgpu::BufferDescriptor {
424            label: Some("spatialrust RGBA readback"),
425            size: buffer_bytes,
426            usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
427            mapped_at_creation: false,
428        });
429        let mut encoder =
430            self.runtime.device().create_command_encoder(&wgpu::CommandEncoderDescriptor {
431                label: Some("spatialrust RGBA readback encoder"),
432            });
433        encoder.copy_texture_to_buffer(
434            wgpu::TexelCopyTextureInfo {
435                texture: &target.color,
436                mip_level: 0,
437                origin: wgpu::Origin3d::ZERO,
438                aspect: wgpu::TextureAspect::All,
439            },
440            wgpu::TexelCopyBufferInfo {
441                buffer: &buffer,
442                layout: wgpu::TexelCopyBufferLayout {
443                    offset: 0,
444                    bytes_per_row: Some(padded_row_bytes),
445                    rows_per_image: Some(target.height),
446                },
447            },
448            wgpu::Extent3d { width: target.width, height: target.height, depth_or_array_layers: 1 },
449        );
450        self.runtime.queue().submit(Some(encoder.finish()));
451        let mapped = map_buffer(&self.runtime, &buffer)?;
452        let logical_bytes = usize::try_from(unpadded_row_bytes)
453            .ok()
454            .and_then(|row| row.checked_mul(target.height as usize))
455            .ok_or_else(|| RenderError::Readback("logical RGBA byte count overflowed".into()))?;
456        let mut rgba = Vec::with_capacity(logical_bytes);
457        for row in mapped.chunks_exact(padded_row_bytes as usize).take(target.height as usize) {
458            rgba.extend_from_slice(&row[..unpadded_row_bytes as usize]);
459        }
460        drop(mapped);
461        buffer.unmap();
462        let mut transfers = TransferReceipt::new();
463        transfers.push(self.readback_event("rgba-readback", logical_bytes as u64)?);
464        Ok(ReadbackImage { width: target.width, height: target.height, rgba, transfers })
465    }
466
467    /// Explicitly reads the point ID at one target pixel.
468    pub fn pick_point(&self, target: &GpuRenderTarget, x: u32, y: u32) -> RenderResult<PickResult> {
469        self.validate_target(target)?;
470        if x >= target.width || y >= target.height {
471            return Err(RenderError::Readback("pick coordinates are outside the target".into()));
472        }
473        let texture = target.point_ids.as_ref().ok_or_else(|| {
474            RenderError::Readback("point picking requires a target rendered from points".into())
475        })?;
476        let buffer = self.runtime.device().create_buffer(&wgpu::BufferDescriptor {
477            label: Some("spatialrust point ID readback"),
478            size: 4,
479            usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
480            mapped_at_creation: false,
481        });
482        let mut encoder =
483            self.runtime.device().create_command_encoder(&wgpu::CommandEncoderDescriptor {
484                label: Some("spatialrust point ID readback encoder"),
485            });
486        encoder.copy_texture_to_buffer(
487            wgpu::TexelCopyTextureInfo {
488                texture,
489                mip_level: 0,
490                origin: wgpu::Origin3d { x, y, z: 0 },
491                aspect: wgpu::TextureAspect::All,
492            },
493            wgpu::TexelCopyBufferInfo {
494                buffer: &buffer,
495                layout: wgpu::TexelCopyBufferLayout {
496                    offset: 0,
497                    bytes_per_row: None,
498                    rows_per_image: None,
499                },
500            },
501            wgpu::Extent3d { width: 1, height: 1, depth_or_array_layers: 1 },
502        );
503        self.runtime.queue().submit(Some(encoder.finish()));
504        let mapped = map_buffer(&self.runtime, &buffer)?;
505        let encoded = u32::from_le_bytes(mapped[..4].try_into().expect("mapped ID is four bytes"));
506        drop(mapped);
507        buffer.unmap();
508        let mut transfers = TransferReceipt::new();
509        transfers.push(self.readback_event("point-id-readback", 4)?);
510        Ok(PickResult { point_index: encoded.checked_sub(1), transfers })
511    }
512
513    fn validate_target(&self, target: &GpuRenderTarget) -> RenderResult<()> {
514        if target.renderer_id != self.id {
515            return Err(RenderError::RuntimeMismatch(
516                "render target belongs to another renderer".into(),
517            ));
518        }
519        Ok(())
520    }
521
522    fn readback_event(&self, stage: &str, bytes: u64) -> RenderResult<TransferEvent> {
523        TransferEvent::try_new(
524            stage,
525            TransferDirection::Readback,
526            VisualResidency::Device(self.device_identity().clone()),
527            VisualResidency::Host,
528            bytes,
529        )
530        .map_err(|error| RenderError::Transfer(error.to_string()))
531    }
532}
533
534fn resolve_style(
535    geometry: &GpuGeometry,
536    style: &VisualStyle,
537) -> RenderResult<(ColorMode, LinearRgba, [f32; 2], f32, u32)> {
538    match (geometry.kind, style) {
539        (GpuGeometryKind::Points, VisualStyle::Points(point_style)) => match point_style.color {
540            PointColor::Uniform(color) => {
541                Ok((ColorMode::Uniform, color, [0.0, 1.0], point_style.size, 0))
542            }
543            PointColor::Rgb if geometry.rgb.is_some() => {
544                Ok((ColorMode::Rgb, LinearRgba::WHITE, [0.0, 1.0], point_style.size, 0))
545            }
546            PointColor::Scalar { min, max, map } if geometry.scalar.is_some() => Ok((
547                ColorMode::Scalar,
548                LinearRgba::WHITE,
549                [min, max],
550                point_style.size,
551                color_map_id(map),
552            )),
553            PointColor::Rgb => {
554                Err(RenderError::GeometrySize("RGB style requires an uploaded RGB buffer".into()))
555            }
556            PointColor::Scalar { .. } => Err(RenderError::GeometrySize(
557                "scalar style requires an uploaded scalar buffer".into(),
558            )),
559        },
560        (GpuGeometryKind::Points, VisualStyle::Uniform(color)) => {
561            Ok((ColorMode::Uniform, *color, [0.0, 1.0], 1.0, 0))
562        }
563        (_, VisualStyle::Uniform(color)) => Ok((ColorMode::Uniform, *color, [0.0, 1.0], 1.0, 0)),
564        (_, VisualStyle::Points(_)) => Err(RenderError::GeometrySize(
565            "point style cannot be applied to line or triangle geometry".into(),
566        )),
567    }
568}
569
570const fn color_map_id(map: ColorMap) -> u32 {
571    match map {
572        ColorMap::Viridis => 0,
573        ColorMap::Turbo => 1,
574        ColorMap::Gray => 2,
575    }
576}
577
578fn create_target(
579    device: &wgpu::Device,
580    renderer_id: u64,
581    width: u32,
582    height: u32,
583    identity: DeviceIdentity,
584    point_picking: bool,
585) -> GpuRenderTarget {
586    let size = wgpu::Extent3d { width, height, depth_or_array_layers: 1 };
587    let color = device.create_texture(&wgpu::TextureDescriptor {
588        label: Some("spatialrust headless color"),
589        size,
590        mip_level_count: 1,
591        sample_count: 1,
592        dimension: wgpu::TextureDimension::D2,
593        format: RENDER_TARGET_FORMAT,
594        usage: wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::COPY_SRC,
595        view_formats: &[],
596    });
597    let depth = device.create_texture(&wgpu::TextureDescriptor {
598        label: Some("spatialrust headless depth"),
599        size,
600        mip_level_count: 1,
601        sample_count: 1,
602        dimension: wgpu::TextureDimension::D2,
603        format: DEPTH_FORMAT,
604        usage: wgpu::TextureUsages::RENDER_ATTACHMENT,
605        view_formats: &[],
606    });
607    let point_ids = point_picking.then(|| {
608        device.create_texture(&wgpu::TextureDescriptor {
609            label: Some("spatialrust point IDs"),
610            size,
611            mip_level_count: 1,
612            sample_count: 1,
613            dimension: wgpu::TextureDimension::D2,
614            format: wgpu::TextureFormat::R32Uint,
615            usage: wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::COPY_SRC,
616            view_formats: &[],
617        })
618    });
619    GpuRenderTarget {
620        renderer_id,
621        color,
622        depth,
623        point_ids,
624        width,
625        height,
626        residency: VisualResidency::Device(identity),
627    }
628}
629
630fn create_point_id_pipeline(
631    device: &wgpu::Device,
632    bind_group_layout: &wgpu::BindGroupLayout,
633) -> wgpu::RenderPipeline {
634    let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
635        label: Some("spatialrust point ID shader"),
636        source: wgpu::ShaderSource::Wgsl(POINT_ID_SHADER.into()),
637    });
638    let layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
639        label: Some("spatialrust point ID pipeline layout"),
640        bind_group_layouts: &[bind_group_layout],
641        push_constant_ranges: &[],
642    });
643    device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
644        label: Some("spatialrust point ID pipeline"),
645        layout: Some(&layout),
646        vertex: wgpu::VertexState {
647            module: &shader,
648            entry_point: Some("vs_main"),
649            compilation_options: wgpu::PipelineCompilationOptions::default(),
650            buffers: &[wgpu::VertexBufferLayout {
651                array_stride: 12,
652                step_mode: wgpu::VertexStepMode::Instance,
653                attributes: &[wgpu::VertexAttribute {
654                    format: wgpu::VertexFormat::Float32x3,
655                    offset: 0,
656                    shader_location: 0,
657                }],
658            }],
659        },
660        primitive: wgpu::PrimitiveState::default(),
661        depth_stencil: Some(wgpu::DepthStencilState {
662            format: DEPTH_FORMAT,
663            depth_write_enabled: false,
664            depth_compare: wgpu::CompareFunction::LessEqual,
665            stencil: wgpu::StencilState::default(),
666            bias: wgpu::DepthBiasState::default(),
667        }),
668        multisample: wgpu::MultisampleState::default(),
669        fragment: Some(wgpu::FragmentState {
670            module: &shader,
671            entry_point: Some("fs_main"),
672            compilation_options: wgpu::PipelineCompilationOptions::default(),
673            targets: &[Some(wgpu::ColorTargetState {
674                format: wgpu::TextureFormat::R32Uint,
675                blend: None,
676                write_mask: wgpu::ColorWrites::ALL,
677            })],
678        }),
679        multiview: None,
680        cache: None,
681    })
682}
683
684fn map_buffer<'a>(
685    runtime: &WgpuRuntime,
686    buffer: &'a wgpu::Buffer,
687) -> RenderResult<wgpu::BufferView<'a>> {
688    let slice = buffer.slice(..);
689    let (sender, receiver) = std::sync::mpsc::channel();
690    slice.map_async(wgpu::MapMode::Read, move |result| {
691        let _ = sender.send(result);
692    });
693    runtime.device().poll(wgpu::Maintain::Wait);
694    receiver
695        .recv()
696        .map_err(|error| RenderError::Readback(format!("map callback dropped: {error}")))?
697        .map_err(|error| RenderError::Readback(format!("buffer map: {error}")))?;
698    Ok(slice.get_mapped_range())
699}
700
701fn align_up(value: u32, alignment: u32) -> RenderResult<u32> {
702    let remainder = value % alignment;
703    if remainder == 0 {
704        return Ok(value);
705    }
706    value
707        .checked_add(alignment - remainder)
708        .ok_or_else(|| RenderError::Readback("row alignment overflowed".into()))
709}
710
711fn create_pipeline(
712    device: &wgpu::Device,
713    bind_group_layout: &wgpu::BindGroupLayout,
714    key: PipelineKey,
715) -> wgpu::RenderPipeline {
716    let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
717        label: Some("spatialrust headless shader"),
718        source: wgpu::ShaderSource::Wgsl(shader_source(key).into()),
719    });
720    let layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
721        label: Some("spatialrust headless pipeline layout"),
722        bind_group_layouts: &[bind_group_layout],
723        push_constant_ranges: &[],
724    });
725    let position_step = if key.kind == GpuGeometryKind::Points {
726        wgpu::VertexStepMode::Instance
727    } else {
728        wgpu::VertexStepMode::Vertex
729    };
730    let position = wgpu::VertexBufferLayout {
731        array_stride: 12,
732        step_mode: position_step,
733        attributes: &[wgpu::VertexAttribute {
734            format: wgpu::VertexFormat::Float32x3,
735            offset: 0,
736            shader_location: 0,
737        }],
738    };
739    let attribute = match key.color {
740        ColorMode::Uniform => None,
741        ColorMode::Rgb => Some(wgpu::VertexBufferLayout {
742            array_stride: 4,
743            step_mode: position_step,
744            attributes: &[wgpu::VertexAttribute {
745                format: wgpu::VertexFormat::Unorm8x4,
746                offset: 0,
747                shader_location: 1,
748            }],
749        }),
750        ColorMode::Scalar => Some(wgpu::VertexBufferLayout {
751            array_stride: 4,
752            step_mode: position_step,
753            attributes: &[wgpu::VertexAttribute {
754                format: wgpu::VertexFormat::Float32,
755                offset: 0,
756                shader_location: 1,
757            }],
758        }),
759    };
760    let mut buffers = vec![position];
761    if let Some(attribute) = attribute {
762        buffers.push(attribute);
763    }
764    device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
765        label: Some("spatialrust headless render pipeline"),
766        layout: Some(&layout),
767        vertex: wgpu::VertexState {
768            module: &shader,
769            entry_point: Some("vs_main"),
770            compilation_options: wgpu::PipelineCompilationOptions::default(),
771            buffers: &buffers,
772        },
773        primitive: wgpu::PrimitiveState {
774            topology: match key.kind {
775                GpuGeometryKind::Points => wgpu::PrimitiveTopology::TriangleList,
776                GpuGeometryKind::Lines => wgpu::PrimitiveTopology::LineList,
777                GpuGeometryKind::Triangles => wgpu::PrimitiveTopology::TriangleList,
778            },
779            ..Default::default()
780        },
781        depth_stencil: Some(wgpu::DepthStencilState {
782            format: DEPTH_FORMAT,
783            depth_write_enabled: true,
784            depth_compare: wgpu::CompareFunction::Less,
785            stencil: wgpu::StencilState::default(),
786            bias: wgpu::DepthBiasState::default(),
787        }),
788        multisample: wgpu::MultisampleState::default(),
789        fragment: Some(wgpu::FragmentState {
790            module: &shader,
791            entry_point: Some("fs_main"),
792            compilation_options: wgpu::PipelineCompilationOptions::default(),
793            targets: &[Some(wgpu::ColorTargetState {
794                format: RENDER_TARGET_FORMAT,
795                blend: Some(wgpu::BlendState::ALPHA_BLENDING),
796                write_mask: wgpu::ColorWrites::ALL,
797            })],
798        }),
799        multiview: None,
800        cache: None,
801    })
802}
803
804fn shader_source(key: PipelineKey) -> &'static str {
805    match (key.kind, key.color) {
806        (GpuGeometryKind::Points, ColorMode::Uniform) => POINT_UNIFORM_SHADER,
807        (GpuGeometryKind::Points, ColorMode::Rgb) => POINT_RGB_SHADER,
808        (GpuGeometryKind::Points, ColorMode::Scalar) => POINT_SCALAR_SHADER,
809        (_, ColorMode::Uniform) => PRIMITIVE_UNIFORM_SHADER,
810        (_, _) => unreachable!("non-point styles are validated as uniform"),
811    }
812}
813
814fn view_projection(camera: Camera, aspect: f32) -> [[f32; 4]; 4] {
815    multiply4(projection_matrix(camera.projection, aspect), view_matrix(camera))
816}
817
818fn view_matrix(camera: Camera) -> [[f32; 4]; 4] {
819    let forward = (camera.target - camera.eye).normalize();
820    let side = forward.cross(camera.up).normalize();
821    let up = side.cross(forward);
822    [
823        [side.x, up.x, -forward.x, 0.0],
824        [side.y, up.y, -forward.y, 0.0],
825        [side.z, up.z, -forward.z, 0.0],
826        [-side.dot(camera.eye), -up.dot(camera.eye), forward.dot(camera.eye), 1.0],
827    ]
828}
829
830fn projection_matrix(projection: Projection, aspect: f32) -> [[f32; 4]; 4] {
831    match projection {
832        Projection::Perspective { vertical_fov_radians, near, far } => {
833            let f = 1.0 / (vertical_fov_radians * 0.5).tan();
834            [
835                [f / aspect, 0.0, 0.0, 0.0],
836                [0.0, f, 0.0, 0.0],
837                [0.0, 0.0, far / (near - far), -1.0],
838                [0.0, 0.0, far * near / (near - far), 0.0],
839            ]
840        }
841        Projection::Orthographic { vertical_span, near, far } => {
842            let horizontal_span = vertical_span * aspect;
843            [
844                [2.0 / horizontal_span, 0.0, 0.0, 0.0],
845                [0.0, 2.0 / vertical_span, 0.0, 0.0],
846                [0.0, 0.0, 1.0 / (near - far), 0.0],
847                [0.0, 0.0, near / (near - far), 1.0],
848            ]
849        }
850    }
851}
852
853fn multiply4(left: [[f32; 4]; 4], right: [[f32; 4]; 4]) -> [[f32; 4]; 4] {
854    let mut output = [[0.0; 4]; 4];
855    for column in 0..4 {
856        for row in 0..4 {
857            output[column][row] = (0..4).map(|index| left[index][row] * right[column][index]).sum();
858        }
859    }
860    output
861}
862
863const POINT_ID_SHADER: &str = r#"
864struct Uniforms {
865    view_projection: mat4x4<f32>, color: vec4<f32>, scalar_range: vec2<f32>,
866    viewport: vec2<f32>, point_size: f32, color_map: u32, padding: vec2<u32>,
867};
868@group(0) @binding(0) var<uniform> uniforms: Uniforms;
869struct VertexOutput {
870    @builtin(position) position: vec4<f32>,
871    @location(0) @interpolate(flat) encoded_id: u32,
872};
873const OFFSETS = array<vec2<f32>, 6>(
874    vec2<f32>(-1.0, -1.0), vec2<f32>(1.0, -1.0), vec2<f32>(1.0, 1.0),
875    vec2<f32>(-1.0, -1.0), vec2<f32>(1.0, 1.0), vec2<f32>(-1.0, 1.0)
876);
877@vertex fn vs_main(
878    @builtin(vertex_index) corner: u32,
879    @builtin(instance_index) point_id: u32,
880    @location(0) center: vec3<f32>
881) -> VertexOutput {
882    var output: VertexOutput;
883    output.position = uniforms.view_projection * vec4<f32>(center, 1.0);
884    let offset = OFFSETS[corner] * uniforms.point_size / uniforms.viewport * 2.0 * output.position.w;
885    output.position = vec4<f32>(output.position.xy + offset, output.position.zw);
886    output.encoded_id = point_id + 1u;
887    return output;
888}
889@fragment fn fs_main(input: VertexOutput) -> @location(0) u32 {
890    return input.encoded_id;
891}
892"#;
893
894const POINT_UNIFORM_SHADER: &str = r#"
895struct Uniforms {
896    view_projection: mat4x4<f32>,
897    color: vec4<f32>,
898    scalar_range: vec2<f32>,
899    viewport: vec2<f32>,
900    point_size: f32,
901    color_map: u32,
902    padding: vec2<u32>,
903};
904@group(0) @binding(0) var<uniform> uniforms: Uniforms;
905struct VertexOutput { @builtin(position) position: vec4<f32>, @location(0) color: vec4<f32> };
906const OFFSETS = array<vec2<f32>, 6>(
907    vec2<f32>(-1.0, -1.0), vec2<f32>(1.0, -1.0), vec2<f32>(1.0, 1.0),
908    vec2<f32>(-1.0, -1.0), vec2<f32>(1.0, 1.0), vec2<f32>(-1.0, 1.0)
909);
910@vertex fn vs_main(@builtin(vertex_index) corner: u32, @location(0) center: vec3<f32>) -> VertexOutput {
911    var output: VertexOutput;
912    output.position = uniforms.view_projection * vec4<f32>(center, 1.0);
913    let offset = OFFSETS[corner] * uniforms.point_size / uniforms.viewport * 2.0 * output.position.w;
914    output.position = vec4<f32>(output.position.xy + offset, output.position.zw);
915    output.color = uniforms.color;
916    return output;
917}
918@fragment fn fs_main(input: VertexOutput) -> @location(0) vec4<f32> { return input.color; }
919"#;
920
921const POINT_RGB_SHADER: &str = r#"
922struct Uniforms {
923    view_projection: mat4x4<f32>, color: vec4<f32>, scalar_range: vec2<f32>,
924    viewport: vec2<f32>, point_size: f32, color_map: u32, padding: vec2<u32>,
925};
926@group(0) @binding(0) var<uniform> uniforms: Uniforms;
927struct VertexOutput { @builtin(position) position: vec4<f32>, @location(0) color: vec4<f32> };
928const OFFSETS = array<vec2<f32>, 6>(
929    vec2<f32>(-1.0, -1.0), vec2<f32>(1.0, -1.0), vec2<f32>(1.0, 1.0),
930    vec2<f32>(-1.0, -1.0), vec2<f32>(1.0, 1.0), vec2<f32>(-1.0, 1.0)
931);
932@vertex fn vs_main(
933    @builtin(vertex_index) corner: u32,
934    @location(0) center: vec3<f32>,
935    @location(1) color: vec4<f32>
936) -> VertexOutput {
937    var output: VertexOutput;
938    output.position = uniforms.view_projection * vec4<f32>(center, 1.0);
939    let offset = OFFSETS[corner] * uniforms.point_size / uniforms.viewport * 2.0 * output.position.w;
940    output.position = vec4<f32>(output.position.xy + offset, output.position.zw);
941    output.color = color;
942    return output;
943}
944@fragment fn fs_main(input: VertexOutput) -> @location(0) vec4<f32> { return input.color; }
945"#;
946
947const POINT_SCALAR_SHADER: &str = r#"
948struct Uniforms {
949    view_projection: mat4x4<f32>, color: vec4<f32>, scalar_range: vec2<f32>,
950    viewport: vec2<f32>, point_size: f32, color_map: u32, padding: vec2<u32>,
951};
952@group(0) @binding(0) var<uniform> uniforms: Uniforms;
953struct VertexOutput { @builtin(position) position: vec4<f32>, @location(0) color: vec4<f32> };
954const OFFSETS = array<vec2<f32>, 6>(
955    vec2<f32>(-1.0, -1.0), vec2<f32>(1.0, -1.0), vec2<f32>(1.0, 1.0),
956    vec2<f32>(-1.0, -1.0), vec2<f32>(1.0, 1.0), vec2<f32>(-1.0, 1.0)
957);
958fn color_map(t_in: f32, map: u32) -> vec3<f32> {
959    let t = clamp(t_in, 0.0, 1.0);
960    if map == 2u { return vec3<f32>(t); }
961    if map == 1u {
962        return clamp(vec3<f32>(
963            1.5 - abs(4.0 * t - 3.0),
964            1.5 - abs(4.0 * t - 2.0),
965            1.5 - abs(4.0 * t - 1.0)
966        ), vec3<f32>(0.0), vec3<f32>(1.0));
967    }
968    return vec3<f32>(
969        0.267 + t * (0.993 - 0.267),
970        0.005 + t * (0.906 - 0.005),
971        0.329 + t * (0.144 - 0.329)
972    );
973}
974@vertex fn vs_main(
975    @builtin(vertex_index) corner: u32,
976    @location(0) center: vec3<f32>,
977    @location(1) scalar: f32
978) -> VertexOutput {
979    var output: VertexOutput;
980    output.position = uniforms.view_projection * vec4<f32>(center, 1.0);
981    let offset = OFFSETS[corner] * uniforms.point_size / uniforms.viewport * 2.0 * output.position.w;
982    output.position = vec4<f32>(output.position.xy + offset, output.position.zw);
983    let t = (scalar - uniforms.scalar_range.x) /
984        (uniforms.scalar_range.y - uniforms.scalar_range.x);
985    output.color = vec4<f32>(color_map(t, uniforms.color_map), 1.0);
986    return output;
987}
988@fragment fn fs_main(input: VertexOutput) -> @location(0) vec4<f32> { return input.color; }
989"#;
990
991const PRIMITIVE_UNIFORM_SHADER: &str = r#"
992struct Uniforms {
993    view_projection: mat4x4<f32>, color: vec4<f32>, scalar_range: vec2<f32>,
994    viewport: vec2<f32>, point_size: f32, color_map: u32, padding: vec2<u32>,
995};
996@group(0) @binding(0) var<uniform> uniforms: Uniforms;
997struct VertexOutput { @builtin(position) position: vec4<f32>, @location(0) color: vec4<f32> };
998@vertex fn vs_main(@location(0) position: vec3<f32>) -> VertexOutput {
999    var output: VertexOutput;
1000    output.position = uniforms.view_projection * vec4<f32>(position, 1.0);
1001    output.color = uniforms.color;
1002    return output;
1003}
1004@fragment fn fs_main(input: VertexOutput) -> @location(0) vec4<f32> { return input.color; }
1005"#;
1006
1007#[cfg(test)]
1008mod tests {
1009    use std::sync::Arc;
1010
1011    use spatialrust_gpu::WgpuRuntime;
1012    use spatialrust_math::Vec3;
1013    use spatialrust_viz::{
1014        Camera, ColorMap, LineListView, LinearRgba, PointCloudView, PointColor, PointStyle,
1015        PositionColumns3, Projection, Rgb8Columns, ScalarColumn, TriangleMeshView, VisualPrimitive,
1016        VisualStyle,
1017    };
1018
1019    use super::{RenderError, RenderOptions, WgpuRenderer};
1020
1021    fn camera() -> Camera {
1022        Camera::try_new(
1023            Vec3::new(0.0, 0.0, 3.0),
1024            Vec3::new(0.0, 0.0, 0.0),
1025            Vec3::new(0.0, 1.0, 0.0),
1026            Projection::Perspective { vertical_fov_radians: 1.0, near: 0.1, far: 10.0 },
1027        )
1028        .unwrap()
1029    }
1030
1031    #[test]
1032    fn rejects_zero_target_size() {
1033        assert!(RenderOptions::try_new(
1034            0,
1035            64,
1036            camera(),
1037            VisualStyle::Uniform(LinearRgba::WHITE),
1038            LinearRgba::BLACK,
1039        )
1040        .is_err());
1041    }
1042
1043    #[test]
1044    fn renders_all_topologies_and_point_color_modes() {
1045        let Ok(runtime) = WgpuRuntime::new_headless() else {
1046            eprintln!("skipping headless render test: no adapter");
1047            return;
1048        };
1049        let runtime = Arc::new(runtime);
1050        let renderer = WgpuRenderer::new(Arc::clone(&runtime));
1051        let other_renderer = WgpuRenderer::new(runtime);
1052        let positions = PositionColumns3::try_new(&[0.0, 0.5], &[0.0, 0.0], &[0.0, 0.0]).unwrap();
1053        let rgb = Rgb8Columns::try_new(&[255, 0], &[0, 255], &[0, 0], 2).unwrap();
1054        let scalar = ScalarColumn::try_new("intensity", &[0.0, 1.0], 2).unwrap();
1055        let points = PointCloudView::positions_only(positions)
1056            .with_rgb(rgb)
1057            .unwrap()
1058            .with_scalar(scalar)
1059            .unwrap();
1060        let (gpu_points, _) = renderer.upload(VisualPrimitive::Points(points)).unwrap();
1061
1062        let styles = [
1063            VisualStyle::Points(
1064                PointStyle::try_new(
1065                    9.0,
1066                    PointColor::Uniform(LinearRgba::try_new(1.0, 0.0, 0.0, 1.0).unwrap()),
1067                )
1068                .unwrap(),
1069            ),
1070            VisualStyle::Points(PointStyle::try_new(5.0, PointColor::Rgb).unwrap()),
1071            VisualStyle::Points(
1072                PointStyle::try_new(
1073                    5.0,
1074                    PointColor::Scalar { min: 0.0, max: 1.0, map: ColorMap::Viridis },
1075                )
1076                .unwrap(),
1077            ),
1078        ];
1079        for (style_index, style) in styles.into_iter().enumerate() {
1080            let options =
1081                RenderOptions::try_new(64, 64, camera(), style, LinearRgba::BLACK).unwrap();
1082            let output = renderer.render_headless(&gpu_points, &options).unwrap();
1083            assert_eq!(output.target.width(), 64);
1084            assert_eq!(
1085                output.receipt.stages,
1086                ["clear-color-depth", "draw-points", "draw-point-ids", "submit"]
1087            );
1088            assert_eq!(output.transfers.events()[0].stage, "render-uniform-upload");
1089            assert_eq!(output.transfers.total_bytes().unwrap(), 112);
1090            if style_index == 0 {
1091                let image = renderer.readback_rgba(&output.target).unwrap();
1092                assert_eq!(image.rgba.len(), 64 * 64 * 4);
1093                let center = (32 * 64 + 32) * 4;
1094                assert_eq!(&image.rgba[center..center + 4], &[255, 0, 0, 255]);
1095                assert_eq!(&image.rgba[..4], &[0, 0, 0, 255]);
1096                assert_eq!(image.transfers.total_bytes().unwrap(), 64 * 64 * 4);
1097                let pick = renderer.pick_point(&output.target, 32, 32).unwrap();
1098                assert_eq!(pick.point_index, Some(0));
1099                assert_eq!(pick.transfers.total_bytes().unwrap(), 4);
1100                assert!(other_renderer.readback_rgba(&output.target).is_err());
1101                assert_eq!(
1102                    renderer.pick_point(&output.target, 64, 32).unwrap_err(),
1103                    RenderError::Readback("pick coordinates are outside the target".into())
1104                );
1105            }
1106        }
1107
1108        let padded_options = RenderOptions::try_new(
1109            17,
1110            19,
1111            camera(),
1112            VisualStyle::Points(
1113                PointStyle::try_new(
1114                    7.0,
1115                    PointColor::Uniform(LinearRgba::try_new(1.0, 0.0, 0.0, 1.0).unwrap()),
1116                )
1117                .unwrap(),
1118            ),
1119            LinearRgba::BLACK,
1120        )
1121        .unwrap();
1122        let padded_output = renderer.render_headless(&gpu_points, &padded_options).unwrap();
1123        let padded_image = renderer.readback_rgba(&padded_output.target).unwrap();
1124        assert_eq!(padded_image.rgba.len(), 17 * 19 * 4);
1125        assert!(renderer.pick_point(&padded_output.target, 8, 9).unwrap().point_index.is_some());
1126
1127        let lines = LineListView::try_new(&[-0.5, 0.0, 0.0, 0.5, 0.0, 0.0]).unwrap();
1128        let (gpu_lines, _) = renderer.upload(VisualPrimitive::Lines(lines)).unwrap();
1129        let uniform = VisualStyle::Uniform(LinearRgba::WHITE);
1130        let options =
1131            RenderOptions::try_new(64, 64, camera(), uniform.clone(), LinearRgba::BLACK).unwrap();
1132        let line_output = renderer.render_headless(&gpu_lines, &options).unwrap();
1133        assert_eq!(line_output.receipt.stages[1], "draw-lines");
1134        assert!(renderer.pick_point(&line_output.target, 32, 32).is_err());
1135
1136        let mesh = TriangleMeshView::try_new(
1137            &[-0.5, -0.5, 0.0, 0.5, -0.5, 0.0, 0.0, 0.5, 0.0],
1138            &[0, 1, 2],
1139        )
1140        .unwrap();
1141        let (gpu_mesh, _) = renderer.upload(VisualPrimitive::Triangles(mesh)).unwrap();
1142        let options = RenderOptions::try_new(64, 64, camera(), uniform, LinearRgba::BLACK).unwrap();
1143        assert_eq!(
1144            renderer.render_headless(&gpu_mesh, &options).unwrap().receipt.stages[1],
1145            "draw-triangles"
1146        );
1147        renderer.runtime.wait_idle();
1148        assert_eq!(renderer.initialized_render_pipeline_count(), 6);
1149    }
1150}
spatialrust_render_wgpu/render.rs

spatialrust_render_wgpu/
render.rs
