src/lib.rs
与本次数学页面一起发布的源码快照。SHA-256: 3ebd65568c04ab84ac77393afd69ae3f8187ff125e4caca3bcade37c3567bc2f
1#![cfg_attr(not(target_arch = "wasm32"), allow(dead_code))]
2
3pub mod gpu_prepare;
4mod inspection;
5pub mod relight;
6pub mod stg_pass;
7pub mod stg_prepare;
8pub mod workshop;
9pub mod workshop_surface;
10#[cfg(target_arch = "wasm32")]
11mod workshop_web;
12
13use bytemuck::{Pod, Zeroable};
14use glam::{Mat4, Vec3};
15
16use std::sync::atomic::{AtomicBool, AtomicU64, Ordering as AtomicOrdering};
17use std::sync::Arc;
18use wasm_bindgen::prelude::*;
19use web_time::Instant;
20use wgpu::util::DeviceExt;
21
22const LOOP_SECONDS: f32 = 10.0;
23const FLOATS_PER_SPLAT: usize = 32;
24const BYTES_PER_SPLAT: usize = FLOATS_PER_SPLAT * std::mem::size_of::<f32>();
25const MAX_SPLATS: usize = 4_000_000;
26const GPU_TIMESTAMP_BYTES: u64 = 4 * std::mem::size_of::<u64>() as u64;
27const GPU_TIMESTAMP_READBACK_SLOTS: usize = 3;
28const SCENE_TARGET: Vec3 = Vec3::new(0.0, 3.5, 14.0);
29const STG_PROPERTIES: [&str; FLOATS_PER_SPLAT] = [
30 "x",
31 "y",
32 "z",
33 "trbf_center",
34 "trbf_scale",
35 "nx",
36 "ny",
37 "nz",
38 "motion_0",
39 "motion_1",
40 "motion_2",
41 "motion_3",
42 "motion_4",
43 "motion_5",
44 "motion_6",
45 "motion_7",
46 "motion_8",
47 "f_dc_0",
48 "f_dc_1",
49 "f_dc_2",
50 "opacity",
51 "scale_0",
52 "scale_1",
53 "scale_2",
54 "rot_0",
55 "rot_1",
56 "rot_2",
57 "rot_3",
58 "omega_0",
59 "omega_1",
60 "omega_2",
61 "omega_3",
62];
63
64#[repr(C)]
65#[derive(Clone, Copy, Pod, Zeroable)]
66struct CameraUniform {
67 view_proj: [[f32; 4]; 4],
68 viewport: [f32; 4],
69 scene: [f32; 4],
70}
71
72#[derive(Clone, Copy)]
73struct ResearchSplat {
74 base: Vec3,
75 temporal_center: f32,
76 temporal_scale: f32,
77 motion_linear: Vec3,
78 motion_quadratic: Vec3,
79 motion_cubic: Vec3,
80 opacity_logit: f32,
81}
82
83impl ResearchSplat {
84 fn sample(self, time: f32) -> (Vec3, f32) {
85 let dt = time - self.temporal_center;
86 let dt2 = dt * dt;
87 let position = self.base
88 + self.motion_linear * dt
89 + self.motion_quadratic * dt2
90 + self.motion_cubic * (dt2 * dt);
91 let temporal = (-(dt / self.temporal_scale).powi(2)).exp();
92 (position, sigmoid(self.opacity_logit) * temporal)
93 }
94}
95
96#[derive(Default)]
97struct FrameTelemetry {
98 prepare_ms: f64,
99 sort_ms: f64,
100 visible: u32,
101 upload_bytes: u32,
102}
103
104struct TimestampReadback {
105 buffer: wgpu::Buffer,
106 pending: Arc<AtomicBool>,
107}
108
109struct GpuTimer {
110 query_set: wgpu::QuerySet,
111 resolve_buffer: wgpu::Buffer,
112 readbacks: Vec<TimestampReadback>,
113 next_readback: usize,
114 latest_ms_bits: Arc<AtomicU64>,
115 latest_prepare_ms_bits: Arc<AtomicU64>,
116 latest_frame_ms_bits: Arc<AtomicU64>,
117 timestamp_period_ns: f64,
118}
119
120impl GpuTimer {
121 fn new(device: &wgpu::Device, queue: &wgpu::Queue) -> Self {
122 let query_set = device.create_query_set(&wgpu::QuerySetDescriptor {
123 label: Some("stg-render-timestamps"),
124 ty: wgpu::QueryType::Timestamp,
125 count: 4,
126 });
127 let resolve_buffer = device.create_buffer(&wgpu::BufferDescriptor {
128 label: Some("stg-timestamp-resolve"),
129 size: GPU_TIMESTAMP_BYTES,
130 usage: wgpu::BufferUsages::QUERY_RESOLVE | wgpu::BufferUsages::COPY_SRC,
131 mapped_at_creation: false,
132 });
133 let readbacks = (0..GPU_TIMESTAMP_READBACK_SLOTS)
134 .map(|_| TimestampReadback {
135 buffer: device.create_buffer(&wgpu::BufferDescriptor {
136 label: Some("stg-timestamp-readback"),
137 size: GPU_TIMESTAMP_BYTES,
138 usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
139 mapped_at_creation: false,
140 }),
141 pending: Arc::new(AtomicBool::new(false)),
142 })
143 .collect();
144 Self {
145 query_set,
146 resolve_buffer,
147 readbacks,
148 next_readback: 0,
149 latest_ms_bits: Arc::new(AtomicU64::new((-1.0_f64).to_bits())),
150 latest_prepare_ms_bits: Arc::new(AtomicU64::new((-1.0_f64).to_bits())),
151 latest_frame_ms_bits: Arc::new(AtomicU64::new((-1.0_f64).to_bits())),
152 timestamp_period_ns: queue.get_timestamp_period() as f64,
153 }
154 }
155
156 fn reserve_readback(&mut self) -> Option<(wgpu::Buffer, Arc<AtomicBool>)> {
157 for _ in 0..self.readbacks.len() {
158 let index = self.next_readback;
159 self.next_readback = (self.next_readback + 1) % self.readbacks.len();
160 let readback = &self.readbacks[index];
161 if !readback.pending.load(AtomicOrdering::Relaxed) {
162 readback.pending.store(true, AtomicOrdering::Relaxed);
163 return Some((readback.buffer.clone(), readback.pending.clone()));
164 }
165 }
166 None
167 }
168}
169
170#[wasm_bindgen]
171pub struct GaussianRenderer {
172 surface: wgpu::Surface<'static>,
173 device: wgpu::Device,
174 queue: wgpu::Queue,
175 config: wgpu::SurfaceConfiguration,
176 pipeline: wgpu::RenderPipeline,
177 bind_group: wgpu::BindGroup,
178 camera_buffer: wgpu::Buffer,
179 index_buffer: wgpu::Buffer,
180 source_buffer: wgpu::Buffer,
181 gpu_preparation: Option<gpu_prepare::GpuPreparation>,
182 preparation: stg_prepare::StgPreparation,
183 last_time: Option<f32>,
184 last_presented: Option<(f32, Mat4, Mat4, [f32; 4], u32, u32)>,
185 render_on_demand: bool,
186 submitted_frames: u32,
187 inspection_records: Vec<[f32; FLOATS_PER_SPLAT]>,
188 inspection_frame: Option<inspection::InspectionFrame>,
189 sorted_indices: Vec<u32>,
190 telemetry: FrameTelemetry,
191 adapter_name: String,
192 gpu_timer: Option<GpuTimer>,
193}
194
195#[wasm_bindgen]
196impl GaussianRenderer {
197 #[wasm_bindgen(js_name = create)]
198 #[cfg(target_arch = "wasm32")]
199 pub async fn create(
200 canvas: web_sys::HtmlCanvasElement,
201 ply_data: js_sys::Uint8Array,
202 ) -> Result<GaussianRenderer, JsValue> {
203 console_error_panic_hook::set_once();
204 let ply_bytes = ply_data.to_vec();
205 let parsed = parse_stg_ply(&ply_bytes).map_err(js_error)?;
206 if parsed.source.len() > MAX_SPLATS {
207 return Err(JsValue::from_str(
208 "STG asset exceeds the renderer's declared splat budget",
209 ));
210 }
211
212 let width = canvas.width().max(1);
213 let height = canvas.height().max(1);
214 let instance = wgpu::Instance::new(wgpu::InstanceDescriptor::new_without_display_handle());
215 let surface = instance
216 .create_surface(wgpu::SurfaceTarget::Canvas(canvas))
217 .map_err(js_error)?;
218 let adapter = instance
219 .request_adapter(&wgpu::RequestAdapterOptions {
220 power_preference: wgpu::PowerPreference::HighPerformance,
221 compatible_surface: Some(&surface),
222 force_fallback_adapter: false,
223 apply_limit_buckets: false,
224 })
225 .await
226 .map_err(js_error)?;
227 let adapter_name = adapter.get_info().name;
228 let timestamp_query_supported =
229 adapter.features().contains(wgpu::Features::TIMESTAMP_QUERY);
230 let mut device_descriptor = wgpu::DeviceDescriptor::default();
231 if timestamp_query_supported {
232 device_descriptor.required_features = wgpu::Features::TIMESTAMP_QUERY;
233 }
234 let (device, queue) = adapter
235 .request_device(&device_descriptor)
236 .await
237 .map_err(js_error)?;
238 if (parsed.payload_end - parsed.payload_offset) as u64
239 > device.limits().max_storage_buffer_binding_size as u64
240 {
241 return Err(JsValue::from_str(
242 "Resident STG window exceeds GPU storage limit; request a smaller time range",
243 ));
244 }
245 let caps = surface.get_capabilities(&adapter);
246 let format = caps
247 .formats
248 .iter()
249 .copied()
250 .find(wgpu::TextureFormat::is_srgb)
251 .or_else(|| caps.formats.first().copied())
252 .ok_or_else(|| JsValue::from_str("WebGPU surface exposes no usable format"))?;
253 let config = wgpu::SurfaceConfiguration {
254 usage: wgpu::TextureUsages::RENDER_ATTACHMENT,
255 format,
256 color_space: wgpu::SurfaceColorSpace::Auto,
257 width,
258 height,
259 present_mode: wgpu::PresentMode::Fifo,
260 alpha_mode: caps.alpha_modes[0],
261 view_formats: vec![],
262 desired_maximum_frame_latency: 2,
263 };
264 surface.configure(&device, &config);
265
266 let camera_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
267 label: Some("stg-camera"),
268 contents: bytemuck::bytes_of(&CameraUniform::zeroed()),
269 usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
270 });
271 let source_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
272 label: Some("stg-source-records"),
273 contents: &ply_bytes[parsed.payload_offset..parsed.payload_end],
274 usage: wgpu::BufferUsages::STORAGE,
275 });
276 let index_buffer = device.create_buffer(&wgpu::BufferDescriptor {
277 label: Some("stg-depth-order"),
278 size: (parsed.source.len() * std::mem::size_of::<u32>()) as u64,
279 usage: wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
280 mapped_at_creation: false,
281 });
282
283 let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
284 label: Some("stg-render-layout"),
285 entries: &[
286 wgpu::BindGroupLayoutEntry {
287 binding: 0,
288 visibility: wgpu::ShaderStages::VERTEX,
289 ty: wgpu::BindingType::Buffer {
290 ty: wgpu::BufferBindingType::Uniform,
291 has_dynamic_offset: false,
292 min_binding_size: None,
293 },
294 count: None,
295 },
296 wgpu::BindGroupLayoutEntry {
297 binding: 1,
298 visibility: wgpu::ShaderStages::VERTEX,
299 ty: wgpu::BindingType::Buffer {
300 ty: wgpu::BufferBindingType::Storage { read_only: true },
301 has_dynamic_offset: false,
302 min_binding_size: None,
303 },
304 count: None,
305 },
306 ],
307 });
308 let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
309 label: Some("stg-render-bind-group"),
310 layout: &bind_group_layout,
311 entries: &[
312 wgpu::BindGroupEntry {
313 binding: 0,
314 resource: camera_buffer.as_entire_binding(),
315 },
316 wgpu::BindGroupEntry {
317 binding: 1,
318 resource: source_buffer.as_entire_binding(),
319 },
320 ],
321 });
322 let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
323 label: Some("stg-ewa-shader"),
324 source: wgpu::ShaderSource::Wgsl(include_str!("splat.wgsl").into()),
325 });
326 let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
327 label: Some("stg-pipeline-layout"),
328 bind_group_layouts: &[Some(&bind_group_layout)],
329 immediate_size: 0,
330 });
331 let vertex_layout = wgpu::VertexBufferLayout {
332 array_stride: std::mem::size_of::<u32>() as u64,
333 step_mode: wgpu::VertexStepMode::Instance,
334 attributes: &[wgpu::VertexAttribute {
335 format: wgpu::VertexFormat::Uint32,
336 offset: 0,
337 shader_location: 0,
338 }],
339 };
340 let pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
341 label: Some("stg-alpha-pipeline"),
342 layout: Some(&pipeline_layout),
343 vertex: wgpu::VertexState {
344 module: &shader,
345 entry_point: Some("vs_main"),
346 compilation_options: Default::default(),
347 buffers: &[Some(vertex_layout)],
348 },
349 fragment: Some(wgpu::FragmentState {
350 module: &shader,
351 entry_point: Some("fs_main"),
352 compilation_options: Default::default(),
353 targets: &[Some(wgpu::ColorTargetState {
354 format,
355 blend: Some(wgpu::BlendState::PREMULTIPLIED_ALPHA_BLENDING),
356 write_mask: wgpu::ColorWrites::ALL,
357 })],
358 }),
359 primitive: wgpu::PrimitiveState {
360 topology: wgpu::PrimitiveTopology::TriangleList,
361 cull_mode: None,
362 ..Default::default()
363 },
364 depth_stencil: None,
365 multisample: Default::default(),
366 multiview_mask: None,
367 cache: None,
368 });
369
370 let source_count = parsed.source.len();
371 let gpu_timer = timestamp_query_supported.then(|| GpuTimer::new(&device, &queue));
372 Ok(Self {
373 surface,
374 device,
375 queue,
376 config,
377 pipeline,
378 bind_group,
379 camera_buffer,
380 index_buffer,
381 source_buffer,
382 gpu_preparation: None,
383 preparation: stg_prepare::StgPreparation::new(parsed.source),
384 last_time: None,
385 last_presented: None,
386 render_on_demand: true,
387 submitted_frames: 0,
388 inspection_records: ply_bytes[parsed.payload_offset..parsed.payload_end]
389 .chunks_exact(BYTES_PER_SPLAT)
390 .map(|bytes| {
391 std::array::from_fn(|i| {
392 f32::from_le_bytes(bytes[i * 4..i * 4 + 4].try_into().unwrap())
393 })
394 })
395 .collect(),
396 inspection_frame: None,
397 sorted_indices: Vec::with_capacity(source_count),
398 telemetry: FrameTelemetry::default(),
399 adapter_name,
400 gpu_timer,
401 })
402 }
403
404 /// Atomically replace the resident temporal window on the existing device.
405 /// The caller must merge overlapping chunks by original ID before calling.
406 #[wasm_bindgen(js_name = replaceSource)]
407 pub fn replace_source(&mut self, ply_data: js_sys::Uint8Array) -> Result<(), JsValue> {
408 let bytes = ply_data.to_vec();
409 let parsed = parse_stg_ply(&bytes).map_err(js_error)?;
410 let payload = &bytes[parsed.payload_offset..parsed.payload_end];
411 if payload.len() as u64 > self.device.limits().max_storage_buffer_binding_size as u64 {
412 return Err(JsValue::from_str(
413 "Resident window exceeds GPU storage limit",
414 ));
415 }
416 let source = self
417 .device
418 .create_buffer_init(&wgpu::util::BufferInitDescriptor {
419 label: Some("stream-resident-records"),
420 contents: payload,
421 usage: wgpu::BufferUsages::STORAGE,
422 });
423 let indices = self.device.create_buffer(&wgpu::BufferDescriptor {
424 label: Some("stream-resident-order"),
425 size: (parsed.source.len() * 4) as u64,
426 usage: wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
427 mapped_at_creation: false,
428 });
429 let bindings = self.device.create_bind_group(&wgpu::BindGroupDescriptor {
430 label: Some("stream-resident-bindings"),
431 layout: &self.pipeline.get_bind_group_layout(0),
432 entries: &[
433 wgpu::BindGroupEntry {
434 binding: 0,
435 resource: self.camera_buffer.as_entire_binding(),
436 },
437 wgpu::BindGroupEntry {
438 binding: 1,
439 resource: source.as_entire_binding(),
440 },
441 ],
442 });
443 let records = payload
444 .chunks_exact(128)
445 .map(|b| {
446 std::array::from_fn(|i| f32::from_le_bytes(b[i * 4..i * 4 + 4].try_into().unwrap()))
447 })
448 .collect();
449 self.preparation = stg_prepare::StgPreparation::new(parsed.source);
450 if self.gpu_preparation.is_some() {
451 self.gpu_preparation = Some(gpu_prepare::GpuPreparation::new(
452 &self.device,
453 &source,
454 self.preparation.len() as u32,
455 ));
456 }
457 self.source_buffer = source;
458 self.inspection_records = records;
459 self.inspection_frame = None;
460 self.last_time = None;
461 self.last_presented = None;
462 self.sorted_indices.clear();
463 self.index_buffer = indices;
464 self.bind_group = bindings;
465 Ok(())
466 }
467
468 pub fn render(
469 &mut self,
470 time_seconds: f32,
471 yaw: f32,
472 pitch: f32,
473 distance: f32,
474 width: u32,
475 height: u32,
476 calibrated_camera: js_sys::Float32Array,
477 ) -> Result<(), JsValue> {
478 let prepare_started = Instant::now();
479 let width = width.max(1);
480 let height = height.max(1);
481 if self.config.width != width || self.config.height != height {
482 self.config.width = width;
483 self.config.height = height;
484 self.surface.configure(&self.device, &self.config);
485 }
486
487 let (view, projection, viewport) = if calibrated_camera.length() >= 12 {
488 let value = |index| calibrated_camera.get_index(index);
489 let eye = Vec3::new(value(0), value(1), value(2));
490 let mut forward = Vec3::new(value(3), value(4), value(5)).normalize_or_zero();
491 let mut up = Vec3::new(value(6), value(7), value(8)).normalize_or_zero();
492 if forward.length_squared() < 0.5 {
493 forward = Vec3::Z;
494 }
495 if up.length_squared() < 0.5 {
496 up = Vec3::NEG_Y;
497 }
498 let capture_aspect = (value(10) / value(11).max(1.0)).clamp(0.5, 3.0);
499 let surface_aspect = width as f32 / height as f32;
500 let viewport = if surface_aspect > capture_aspect {
501 let viewport_width = height as f32 * capture_aspect;
502 [
503 (width as f32 - viewport_width) * 0.5,
504 0.0,
505 viewport_width,
506 height as f32,
507 ]
508 } else {
509 let viewport_height = width as f32 / capture_aspect;
510 [
511 0.0,
512 (height as f32 - viewport_height) * 0.5,
513 width as f32,
514 viewport_height,
515 ]
516 };
517 (
518 Mat4::look_at_rh(eye, eye + forward, up),
519 Mat4::perspective_rh(value(9).clamp(0.2, 2.6), capture_aspect, 0.05, 180.0),
520 viewport,
521 )
522 } else {
523 let pitch = pitch.clamp(-0.55, 0.7);
524 let distance = distance.clamp(9.0, 34.0);
525 let orbit = Vec3::new(
526 yaw.sin() * pitch.cos(),
527 -pitch.sin(),
528 -yaw.cos() * pitch.cos(),
529 );
530 let eye = SCENE_TARGET + orbit * distance;
531 (
532 Mat4::look_at_rh(eye, SCENE_TARGET, Vec3::NEG_Y),
533 Mat4::perspective_rh(
534 69.5_f32.to_radians(),
535 width as f32 / height as f32,
536 0.05,
537 180.0,
538 ),
539 [0.0, 0.0, width as f32, height as f32],
540 )
541 };
542 let view_proj = projection * view;
543 let normalized_time = time_seconds.rem_euclid(LOOP_SECONDS) / LOOP_SECONDS;
544
545 let presented_key = (normalized_time, view, projection, viewport, width, height);
546 if self.render_on_demand && self.last_presented == Some(presented_key) {
547 self.telemetry.prepare_ms = prepare_started.elapsed().as_secs_f64() * 1000.0;
548 self.telemetry.sort_ms = 0.0;
549 self.telemetry.upload_bytes = 0;
550 return Ok(());
551 }
552 self.inspection_frame = None;
553 let variant = 9;
554 if self.gpu_preparation.is_none() {
555 self.preparation
556 .prepare(normalized_time, view, projection, variant);
557 }
558 self.last_time = Some(normalized_time);
559 self.telemetry.sort_ms = if self.gpu_preparation.is_some() {
560 0.0
561 } else {
562 self.preparation.sort_ms
563 };
564
565 let camera = CameraUniform {
566 view_proj: view_proj.to_cols_array_2d(),
567 viewport: [
568 viewport[2],
569 viewport[3],
570 1.0 / viewport[2],
571 1.0 / viewport[3],
572 ],
573 scene: [normalized_time, 0.0, 0.0, 0.0],
574 };
575 self.queue
576 .write_buffer(&self.camera_buffer, 0, bytemuck::bytes_of(&camera));
577 if self.gpu_preparation.is_none() && !self.preparation.cache_hit {
578 self.sorted_indices.clear();
579 self.sorted_indices
580 .extend_from_slice(self.preparation.indices());
581 self.queue.write_buffer(
582 &self.index_buffer,
583 0,
584 bytemuck::cast_slice(&self.sorted_indices),
585 );
586 }
587 self.telemetry.visible = self.sorted_indices.len() as u32;
588 self.telemetry.upload_bytes =
589 if self.gpu_preparation.is_some() || self.preparation.cache_hit {
590 0
591 } else {
592 (self.sorted_indices.len() * 4) as u32
593 };
594 self.telemetry.prepare_ms = prepare_started.elapsed().as_secs_f64() * 1000.0;
595
596 let inspection_frame = inspection::InspectionFrame {
597 time: normalized_time,
598 view,
599 view_proj,
600 viewport,
601 eye: view.inverse().transform_point3(Vec3::ZERO),
602 };
603 let frame = match self.surface.get_current_texture() {
604 wgpu::CurrentSurfaceTexture::Success(frame)
605 | wgpu::CurrentSurfaceTexture::Suboptimal(frame) => frame,
606 wgpu::CurrentSurfaceTexture::Outdated | wgpu::CurrentSurfaceTexture::Lost => {
607 self.surface.configure(&self.device, &self.config);
608 return Ok(());
609 }
610 wgpu::CurrentSurfaceTexture::Timeout | wgpu::CurrentSurfaceTexture::Occluded => {
611 return Ok(())
612 }
613 wgpu::CurrentSurfaceTexture::Validation => {
614 return Err(JsValue::from_str("WebGPU surface validation failed"))
615 }
616 };
617 let target_view = frame.texture.create_view(&Default::default());
618 let mut encoder = self
619 .device
620 .create_command_encoder(&wgpu::CommandEncoderDescriptor {
621 label: Some("stg-frame"),
622 });
623 let computed = if let Some(gpu) = self.gpu_preparation.as_mut() {
624 gpu.encode(
625 &self.queue,
626 &mut encoder,
627 normalized_time,
628 view,
629 projection,
630 self.gpu_timer.as_ref().map(|t| (&t.query_set, 0, 1)),
631 )
632 } else {
633 false
634 };
635 {
636 let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
637 label: Some("stg-splats"),
638 color_attachments: &[Some(wgpu::RenderPassColorAttachment {
639 view: &target_view,
640 resolve_target: None,
641 depth_slice: None,
642 ops: wgpu::Operations {
643 load: wgpu::LoadOp::Clear(wgpu::Color {
644 r: 0.0,
645 g: 0.0,
646 b: 0.0,
647 a: 1.0,
648 }),
649 store: wgpu::StoreOp::Store,
650 },
651 })],
652 depth_stencil_attachment: None,
653 timestamp_writes: self.gpu_timer.as_ref().map(|timer| {
654 wgpu::RenderPassTimestampWrites {
655 query_set: &timer.query_set,
656 beginning_of_pass_write_index: Some(2),
657 end_of_pass_write_index: Some(3),
658 }
659 }),
660 occlusion_query_set: None,
661 multiview_mask: None,
662 });
663 pass.set_pipeline(&self.pipeline);
664 pass.set_bind_group(0, &self.bind_group, &[]);
665 pass.set_viewport(viewport[0], viewport[1], viewport[2], viewport[3], 0.0, 1.0);
666 if let Some(gpu) = &self.gpu_preparation {
667 pass.set_vertex_buffer(0, gpu.indices.slice(..));
668 pass.draw_indirect(&gpu.indirect, 0);
669 } else {
670 pass.set_vertex_buffer(0, self.index_buffer.slice(..));
671 pass.draw(0..6, 0..self.telemetry.visible);
672 }
673 }
674 let mut timestamp_readback = None;
675 if let Some(timer) = self.gpu_timer.as_mut() {
676 if let Some((buffer, pending)) = timer.reserve_readback() {
677 encoder.resolve_query_set(
678 &timer.query_set,
679 if computed { 0..4 } else { 2..4 },
680 &timer.resolve_buffer,
681 0,
682 );
683 encoder.copy_buffer_to_buffer(
684 &timer.resolve_buffer,
685 0,
686 &buffer,
687 0,
688 GPU_TIMESTAMP_BYTES,
689 );
690 timestamp_readback = Some((
691 buffer,
692 pending,
693 timer.latest_ms_bits.clone(),
694 timer.latest_prepare_ms_bits.clone(),
695 timer.latest_frame_ms_bits.clone(),
696 timer.timestamp_period_ns,
697 ));
698 }
699 }
700 let count_copy = self
701 .gpu_preparation
702 .as_ref()
703 .is_some_and(|gpu| gpu.count_copy(&mut encoder));
704 self.telemetry.prepare_ms = prepare_started.elapsed().as_secs_f64() * 1000.0;
705 self.queue.submit([encoder.finish()]);
706 if count_copy {
707 self.gpu_preparation.as_ref().unwrap().map_count();
708 }
709 if let Some((
710 buffer,
711 pending,
712 latest_ms_bits,
713 latest_prepare_ms_bits,
714 latest_frame_ms_bits,
715 timestamp_period_ns,
716 )) = timestamp_readback
717 {
718 let callback_buffer = buffer.clone();
719 buffer
720 .slice(..GPU_TIMESTAMP_BYTES)
721 .map_async(wgpu::MapMode::Read, move |result| {
722 if result.is_ok() {
723 let mapped_result = callback_buffer
724 .slice(..GPU_TIMESTAMP_BYTES)
725 .get_mapped_range();
726 if let Ok(mapped) = mapped_result {
727 let offset = if computed { 16 } else { 0 };
728 let begin =
729 u64::from_le_bytes(mapped[offset..offset + 8].try_into().unwrap());
730 let end = u64::from_le_bytes(
731 mapped[offset + 8..offset + 16].try_into().unwrap(),
732 );
733 if end >= begin {
734 let elapsed_ms =
735 (end - begin) as f64 * timestamp_period_ns / 1_000_000.0;
736 latest_ms_bits.store(elapsed_ms.to_bits(), AtomicOrdering::Relaxed);
737 let first = if computed {
738 u64::from_le_bytes(mapped[0..8].try_into().unwrap())
739 } else {
740 begin
741 };
742 let compute_end = if computed {
743 u64::from_le_bytes(mapped[8..16].try_into().unwrap())
744 } else {
745 first
746 };
747 if compute_end >= first && end >= first {
748 latest_prepare_ms_bits.store(
749 ((compute_end - first) as f64 * timestamp_period_ns / 1e6)
750 .to_bits(),
751 AtomicOrdering::Relaxed,
752 );
753 latest_frame_ms_bits.store(
754 ((end - first) as f64 * timestamp_period_ns / 1e6)
755 .to_bits(),
756 AtomicOrdering::Relaxed,
757 );
758 }
759 }
760 drop(mapped);
761 }
762 callback_buffer.unmap();
763 }
764 pending.store(false, AtomicOrdering::Relaxed);
765 });
766 }
767 self.queue.present(frame);
768 self.last_presented = Some(presented_key);
769 self.submitted_frames = self.submitted_frames.wrapping_add(1);
770 self.inspection_frame = Some(inspection_frame);
771 Ok(())
772 }
773
774 /// Hosts with external canvas invalidation may request an explicit redraw.
775 #[wasm_bindgen(js_name = invalidate)]
776 pub fn invalidate(&mut self) {
777 self.last_presented = None;
778 }
779 #[wasm_bindgen(js_name = setRenderOnDemand)]
780 pub fn set_render_on_demand(&mut self, enabled: bool) {
781 self.render_on_demand = enabled;
782 self.invalidate();
783 }
784 #[wasm_bindgen(getter, js_name = submittedFrames)]
785 pub fn submitted_frames(&self) -> u32 {
786 self.submitted_frames
787 }
788
789 #[wasm_bindgen(getter, js_name = prepareMs)]
790 pub fn prepare_ms(&self) -> f64 {
791 self.telemetry.prepare_ms
792 }
793
794 #[wasm_bindgen(getter, js_name = sortMs)]
795 pub fn sort_ms(&self) -> f64 {
796 self.telemetry.sort_ms
797 }
798
799 #[wasm_bindgen(getter)]
800 pub fn visible(&self) -> u32 {
801 self.gpu_preparation
802 .as_ref()
803 .map_or(self.telemetry.visible, |g| g.visible())
804 }
805
806 #[wasm_bindgen(getter, js_name = uploadBytes)]
807 pub fn upload_bytes(&self) -> u32 {
808 self.telemetry.upload_bytes
809 }
810
811 #[wasm_bindgen(getter, js_name = gpuRenderMs)]
812 pub fn gpu_render_ms(&self) -> f64 {
813 self.gpu_timer
814 .as_ref()
815 .map(|timer| f64::from_bits(timer.latest_ms_bits.load(AtomicOrdering::Relaxed)))
816 .unwrap_or(-1.0)
817 }
818
819 #[wasm_bindgen(js_name = setGpuDriven)]
820 pub fn set_gpu_driven(&mut self, enabled: bool) {
821 if enabled == self.gpu_preparation.is_some() {
822 return;
823 }
824 self.gpu_preparation = enabled.then(|| {
825 gpu_prepare::GpuPreparation::new(
826 &self.device,
827 &self.source_buffer,
828 self.preparation.len() as u32,
829 )
830 });
831 self.preparation.invalidate();
832 self.invalidate();
833 }
834 #[wasm_bindgen(getter, js_name = gpuDriven)]
835 pub fn gpu_driven(&self) -> bool {
836 self.gpu_preparation.is_some()
837 }
838 #[wasm_bindgen(getter, js_name = gpuPrepareMs)]
839 pub fn gpu_prepare_ms(&self) -> f64 {
840 self.gpu_timer.as_ref().map_or(-1.0, |t| {
841 f64::from_bits(t.latest_prepare_ms_bits.load(AtomicOrdering::Relaxed))
842 })
843 }
844 #[wasm_bindgen(getter, js_name = gpuFrameMs)]
845 pub fn gpu_frame_ms(&self) -> f64 {
846 self.gpu_timer.as_ref().map_or(-1.0, |t| {
847 f64::from_bits(t.latest_frame_ms_bits.load(AtomicOrdering::Relaxed))
848 })
849 }
850
851 #[wasm_bindgen(getter, js_name = gpuTimingSupported)]
852 pub fn gpu_timing_supported(&self) -> bool {
853 self.gpu_timer.is_some()
854 }
855
856 #[wasm_bindgen(getter, js_name = adapterName)]
857 pub fn adapter_name(&self) -> String {
858 self.adapter_name.clone()
859 }
860
861 #[wasm_bindgen(getter, js_name = sourceCount)]
862 pub fn source_count(&self) -> u32 {
863 self.preparation.len() as u32
864 }
865
866 #[wasm_bindgen(js_name = splatRecord)]
867 pub fn splat_record(&self, id: u32) -> Vec<f32> {
868 self.inspection_records
869 .get(id as usize)
870 .map_or_else(Vec::new, |r| r.to_vec())
871 }
872
873 #[wasm_bindgen(js_name = inspectSplat)]
874 pub fn inspect_splat(&self, id: u32) -> Vec<f32> {
875 match (
876 self.inspection_records.get(id as usize),
877 self.inspection_frame.as_ref(),
878 ) {
879 (Some(row), Some(frame)) => {
880 inspection::project(row, frame).map_or_else(Vec::new, |s| s.pack(frame))
881 }
882 _ => vec![],
883 }
884 }
885
886 #[wasm_bindgen(js_name = pickSplats)]
887 pub fn pick_splats(&mut self, x: f32, y: f32) -> Vec<f32> {
888 if self.gpu_preparation.is_some() {
889 if let Some(frame) = self.inspection_frame {
890 // CPU picking is explicitly on-demand, never part of GPU playback.
891 self.preparation.prepare(
892 frame.time,
893 frame.view,
894 frame.view_proj * frame.view.inverse(),
895 9,
896 );
897 }
898 }
899 self.inspection_frame
900 .as_ref()
901 .map_or_else(Vec::new, |frame| {
902 inspection::pick(
903 &self.inspection_records,
904 self.preparation.ordered(),
905 frame,
906 glam::Vec2::new(x, y),
907 )
908 })
909 }
910}
911
912struct ParsedPly {
913 source: Vec<ResearchSplat>,
914 payload_offset: usize,
915 payload_end: usize,
916}
917
918fn parse_stg_ply(bytes: &[u8]) -> Result<ParsedPly, String> {
919 const END_HEADER: &[u8] = b"end_header\n";
920 let header_start = bytes
921 .windows(END_HEADER.len())
922 .position(|window| window == END_HEADER)
923 .ok_or_else(|| "PLY header is missing end_header".to_string())?;
924 let payload_offset = header_start + END_HEADER.len();
925 let header = std::str::from_utf8(&bytes[..payload_offset])
926 .map_err(|_| "PLY header is not UTF-8".to_string())?;
927 if !header.contains("format binary_little_endian 1.0") {
928 return Err("Expected a binary little-endian PLY".to_string());
929 }
930 let count = header
931 .lines()
932 .find_map(|line| line.strip_prefix("element vertex "))
933 .ok_or_else(|| "PLY vertex count is missing".to_string())?
934 .parse::<usize>()
935 .map_err(|_| "PLY vertex count is invalid".to_string())?;
936 let properties: Vec<_> = header
937 .lines()
938 .filter_map(|line| line.strip_prefix("property float "))
939 .collect();
940 if properties.len() != FLOATS_PER_SPLAT {
941 return Err(format!(
942 "Expected the 32-field STG-Lite layout, found {} fields",
943 properties.len()
944 ));
945 }
946 for (index, (actual, expected)) in properties.iter().zip(STG_PROPERTIES).enumerate() {
947 if *actual != expected {
948 return Err(format!(
949 "STG-Lite field {index} must be {expected}, found {actual}"
950 ));
951 }
952 }
953 if count == 0 || count > MAX_SPLATS {
954 return Err(format!(
955 "Expected 1..={MAX_SPLATS} STG-Lite splats, found {count}"
956 ));
957 }
958 let payload_end = count
959 .checked_mul(BYTES_PER_SPLAT)
960 .and_then(|length| payload_offset.checked_add(length))
961 .ok_or_else(|| "PLY payload length overflowed".to_string())?;
962 if payload_end > bytes.len() {
963 return Err("PLY payload is truncated".to_string());
964 }
965
966 let payload = &bytes[payload_offset..payload_end];
967 let mut source = Vec::with_capacity(count);
968 for (index, record) in payload.chunks_exact(BYTES_PER_SPLAT).enumerate() {
969 let value = |index: usize| {
970 let offset = index * 4;
971 f32::from_le_bytes(record[offset..offset + 4].try_into().unwrap())
972 };
973 if (0..FLOATS_PER_SPLAT).any(|field| !value(field).is_finite()) {
974 return Err(format!(
975 "STG-Lite record {index} contains a non-finite value"
976 ));
977 }
978 if [4, 21, 22, 23]
979 .iter()
980 .any(|&field| !value(field).exp().is_finite())
981 {
982 return Err(format!(
983 "STG-Lite record {index} has an overflowing log scale"
984 ));
985 }
986 source.push(ResearchSplat {
987 base: Vec3::new(value(0), value(1), value(2)),
988 temporal_center: value(3),
989 temporal_scale: value(4).exp().max(1e-6),
990 motion_linear: Vec3::new(value(8), value(9), value(10)),
991 motion_quadratic: Vec3::new(value(11), value(12), value(13)),
992 motion_cubic: Vec3::new(value(14), value(15), value(16)),
993 opacity_logit: value(20),
994 });
995 }
996 Ok(ParsedPly {
997 source,
998 payload_offset,
999 payload_end,
1000 })
1001}
1002
1003fn sigmoid(value: f32) -> f32 {
1004 1.0 / (1.0 + (-value).exp())
1005}
1006
1007fn js_error(error: impl std::fmt::Display) -> JsValue {
1008 JsValue::from_str(&error.to_string())
1009}
1010
1011#[cfg(test)]
1012mod tests {
1013 use super::*;
1014
1015 fn fixture_ply() -> Vec<u8> {
1016 let mut bytes = b"ply\nformat binary_little_endian 1.0\nelement vertex 1\n".to_vec();
1017 for property in STG_PROPERTIES {
1018 bytes.extend_from_slice(format!("property float {property}\n").as_bytes());
1019 }
1020 bytes.extend_from_slice(b"end_header\n");
1021 let mut values = [0.0_f32; FLOATS_PER_SPLAT];
1022 values[0] = 1.0;
1023 values[1] = 2.0;
1024 values[2] = 3.0;
1025 values[4] = 0.0;
1026 values[8] = 2.0;
1027 values[20] = 0.0;
1028 for value in values {
1029 bytes.extend_from_slice(&value.to_le_bytes());
1030 }
1031 bytes
1032 }
1033
1034 #[test]
1035 fn parses_the_official_stg_lite_record_shape() {
1036 let parsed = parse_stg_ply(&fixture_ply()).unwrap();
1037 assert_eq!(parsed.source.len(), 1);
1038 let (position, opacity) = parsed.source[0].sample(0.5);
1039 assert_eq!(position, Vec3::new(2.0, 2.0, 3.0));
1040 assert!(opacity > 0.38 && opacity < 0.40);
1041 }
1042
1043 #[test]
1044 fn rejects_an_unrelated_ply_layout() {
1045 let bytes = b"ply\nformat binary_little_endian 1.0\nelement vertex 0\nend_header\n";
1046 assert!(parse_stg_ply(bytes).is_err());
1047 }
1048
1049 #[test]
1050 fn rejects_a_reordered_stg_layout() {
1051 let mut bytes = fixture_ply();
1052 let x = bytes
1053 .windows(b"property float x\n".len())
1054 .position(|window| window == b"property float x\n")
1055 .unwrap();
1056 bytes[x + "property float ".len()] = b'q';
1057 assert!(parse_stg_ply(&bytes).is_err());
1058 }
1059
1060 #[test]
1061 fn rejects_non_finite_attributes_and_oversized_counts_before_allocation() {
1062 let mut bytes = fixture_ply();
1063 let offset = parse_stg_ply(&bytes).unwrap().payload_offset;
1064 bytes[offset..offset + 4].copy_from_slice(&f32::NAN.to_le_bytes());
1065 assert!(matches!(parse_stg_ply(&bytes), Err(message) if message.contains("non-finite")));
1066 let bytes = fixture_ply();
1067 let header = String::from_utf8_lossy(&bytes[..offset])
1068 .replace("element vertex 1", "element vertex 18446744073709551615");
1069 assert!(parse_stg_ply(header.as_bytes()).is_err());
1070 }
1071}
1072