Skip to main content

loam_render/raymarch/
mod.rs

1//! A user shader for [`RayMarchNode`] exports `vs_fullscreen(@builtin(vertex_index))` and
2//! `fs_main(@builtin(position))` and binds [`RayMarchUniforms`] at group 0, binding 0.
3
4mod hyperslice4d;
5mod polytope_data;
6pub mod scene;
7pub use hyperslice4d::{
8    BodyKind, BodyUniform, Hyperslice4DNode, Hyperslice4DUniforms, HYPERSLICE_KERNEL_WGSL,
9    SHAPE_120CELL, SHAPE_16CELL, SHAPE_24CELL, SHAPE_3SPHERE, SHAPE_600CELL, SHAPE_CLIFFORD_TORUS,
10    SHAPE_DUOCYLINDER, SHAPE_PENTATOPE, SHAPE_SPHERINDER, SHAPE_TESSERACT,
11};
12pub use polytope_data::{polytope_extended_sdfs_wgsl, polytope_stub_sdfs_wgsl};
13
14/// `None` for the smooth-surface SDFs.
15pub fn polytope4_from_shape_id(shape: u32) -> Option<loam_shape::polytope::Polytope4> {
16    use loam_shape::polytope::Polytope4;
17    Some(match shape {
18        s if s == SHAPE_PENTATOPE => Polytope4::Pentatope,
19        s if s == SHAPE_TESSERACT => Polytope4::Tesseract,
20        s if s == SHAPE_16CELL => Polytope4::Cell16,
21        s if s == SHAPE_24CELL => Polytope4::Cell24,
22        s if s == SHAPE_120CELL => Polytope4::Cell120,
23        s if s == SHAPE_600CELL => Polytope4::Cell600,
24        _ => return None,
25    })
26}
27
28#[derive(Copy, Clone, Debug, PartialEq, Eq)]
29pub enum RaymarchShape {
30    Polytope(loam_shape::polytope::Polytope4),
31    ThreeSphere,
32    Duocylinder,
33    CliffordTorus,
34    Spherinder,
35}
36
37impl RaymarchShape {
38    pub fn shape_id(&self) -> u32 {
39        use loam_shape::polytope::Polytope4;
40        match self {
41            RaymarchShape::Polytope(p) => match p {
42                Polytope4::Pentatope => SHAPE_PENTATOPE,
43                Polytope4::Tesseract => SHAPE_TESSERACT,
44                Polytope4::Cell16 => SHAPE_16CELL,
45                Polytope4::Cell24 => SHAPE_24CELL,
46                Polytope4::Cell120 => SHAPE_120CELL,
47                Polytope4::Cell600 => SHAPE_600CELL,
48            },
49            RaymarchShape::ThreeSphere => SHAPE_3SPHERE,
50            RaymarchShape::Duocylinder => SHAPE_DUOCYLINDER,
51            RaymarchShape::CliffordTorus => SHAPE_CLIFFORD_TORUS,
52            RaymarchShape::Spherinder => SHAPE_SPHERINDER,
53        }
54    }
55
56    pub fn polytope4(&self) -> Option<loam_shape::polytope::Polytope4> {
57        match self {
58            RaymarchShape::Polytope(p) => Some(*p),
59            _ => None,
60        }
61    }
62}
63
64impl From<loam_shape::polytope::Polytope4> for RaymarchShape {
65    fn from(p: loam_shape::polytope::Polytope4) -> Self {
66        RaymarchShape::Polytope(p)
67    }
68}
69
70impl From<RaymarchShape> for u32 {
71    fn from(s: RaymarchShape) -> u32 {
72        s.shape_id()
73    }
74}
75
76use bytemuck::{Pod, Zeroable};
77use wgpu::*;
78
79/// Bind group 0, binding 0, `std140`: every `vec3` pads to 16 bytes.
80#[repr(C)]
81#[derive(Copy, Clone, Debug, Pod, Zeroable)]
82pub struct RayMarchUniforms {
83    pub camera_pos: [f32; 3],
84    pub _pad0: f32,
85    pub camera_forward: [f32; 3],
86    pub _pad1: f32,
87    pub camera_right: [f32; 3],
88    pub _pad2: f32,
89    pub camera_up: [f32; 3],
90    pub fov_y_tan: f32,
91    pub resolution: [f32; 2],
92    pub time: f32,
93    pub tick: f32,
94    /// Four scalar knobs; semantics are up to the user shader.
95    pub params: [f32; 4],
96    /// Inverse of the eye's isometry on the domain's embedding; the view map composes it first.
97    pub eye_inverse: [[f32; 4]; 4],
98    /// Near distance of the root projection; see [`crate::view`].
99    pub near: f32,
100    pub _pad3: [f32; 3],
101}
102
103impl Default for RayMarchUniforms {
104    fn default() -> Self {
105        Self {
106            camera_pos: [0.0, 0.0, 3.0],
107            _pad0: 0.0,
108            camera_forward: [0.0, 0.0, -1.0],
109            _pad1: 0.0,
110            camera_right: [1.0, 0.0, 0.0],
111            _pad2: 0.0,
112            camera_up: [0.0, 1.0, 0.0],
113            fov_y_tan: (60.0_f32.to_radians() * 0.5).tan(),
114            resolution: [1.0, 1.0],
115            time: 0.0,
116            tick: 0.0,
117            params: [0.0; 4],
118            eye_inverse: glam::Mat4::IDENTITY.to_cols_array_2d(),
119            near: 0.05,
120            _pad3: [0.0; 3],
121        }
122    }
123}
124
125pub struct RayMarchNode {
126    pipeline: RenderPipeline,
127    uniforms: RayMarchUniforms,
128    uniform_buf: Buffer,
129    bind_group: BindGroup,
130    clear_color: Color,
131    has_depth: bool,
132}
133
134impl RayMarchNode {
135    pub fn new(device: &Device, surface_format: TextureFormat, shader: &ShaderModule) -> Self {
136        Self::with_depth(device, surface_format, shader, crate::DepthMode::Off)
137    }
138
139    /// The user's `fs_main` writes [`crate::view`]'s projective depth as `frag_depth`.
140    pub fn with_depth(
141        device: &Device,
142        surface_format: TextureFormat,
143        shader: &ShaderModule,
144        depth: crate::DepthMode,
145    ) -> Self {
146        let uniform_buf = device.create_buffer(&BufferDescriptor {
147            label: Some("raymarch uniforms"),
148            size: std::mem::size_of::<RayMarchUniforms>() as u64,
149            usage: BufferUsages::UNIFORM | BufferUsages::COPY_DST,
150            mapped_at_creation: false,
151        });
152
153        let bgl = device.create_bind_group_layout(&BindGroupLayoutDescriptor {
154            label: Some("raymarch bgl"),
155            entries: &[BindGroupLayoutEntry {
156                binding: 0,
157                visibility: ShaderStages::VERTEX_FRAGMENT,
158                ty: BindingType::Buffer {
159                    ty: BufferBindingType::Uniform,
160                    has_dynamic_offset: false,
161                    min_binding_size: None,
162                },
163                count: None,
164            }],
165        });
166
167        let bind_group = device.create_bind_group(&BindGroupDescriptor {
168            label: Some("raymarch bg"),
169            layout: &bgl,
170            entries: &[BindGroupEntry {
171                binding: 0,
172                resource: uniform_buf.as_entire_binding(),
173            }],
174        });
175
176        let pipeline_layout = device.create_pipeline_layout(&PipelineLayoutDescriptor {
177            label: Some("raymarch pipeline layout"),
178            bind_group_layouts: &[&bgl],
179            push_constant_ranges: &[],
180        });
181
182        let pipeline = device.create_render_pipeline(&RenderPipelineDescriptor {
183            label: Some("raymarch pipeline"),
184            layout: Some(&pipeline_layout),
185            vertex: VertexState {
186                module: shader,
187                entry_point: Some("vs_fullscreen"),
188                buffers: &[],
189                compilation_options: Default::default(),
190            },
191            fragment: Some(FragmentState {
192                module: shader,
193                entry_point: Some("fs_main"),
194                targets: &[Some(ColorTargetState {
195                    format: surface_format,
196                    blend: None,
197                    write_mask: ColorWrites::ALL,
198                })],
199                compilation_options: Default::default(),
200            }),
201            primitive: PrimitiveState {
202                topology: PrimitiveTopology::TriangleList,
203                ..Default::default()
204            },
205            depth_stencil: depth.format().map(|format| DepthStencilState {
206                format,
207                depth_write_enabled: depth.writes(),
208                depth_compare: crate::view::DEPTH_COMPARE,
209                stencil: StencilState::default(),
210                bias: DepthBiasState::default(),
211            }),
212            multisample: MultisampleState {
213                count: 1,
214                ..Default::default()
215            },
216            multiview: None,
217            cache: None,
218        });
219
220        Self {
221            pipeline,
222            uniforms: RayMarchUniforms::default(),
223            uniform_buf,
224            bind_group,
225            clear_color: Color::BLACK,
226            has_depth: depth.is_active(),
227        }
228    }
229
230    pub fn uniforms(&self) -> &RayMarchUniforms {
231        &self.uniforms
232    }
233
234    pub fn uniforms_mut(&mut self) -> &mut RayMarchUniforms {
235        &mut self.uniforms
236    }
237
238    pub fn set_uniforms(&mut self, queue: &Queue, uniforms: RayMarchUniforms) {
239        self.uniforms = uniforms;
240        queue.write_buffer(&self.uniform_buf, 0, bytemuck::bytes_of(&self.uniforms));
241    }
242
243    pub fn flush_uniforms(&self, queue: &Queue) {
244        queue.write_buffer(&self.uniform_buf, 0, bytemuck::bytes_of(&self.uniforms));
245    }
246}
247
248impl RayMarchNode {
249    /// Clears the whole attachment before shading the viewport.
250    pub fn record_in_viewport(
251        &mut self,
252        encoder: &mut wgpu::CommandEncoder,
253        view: &wgpu::TextureView,
254        viewport: crate::Viewport,
255    ) {
256        self.record(encoder, view, None, viewport);
257    }
258
259    /// Clears both attachments; depth clears to [`crate::view::DEPTH_CLEAR`].
260    pub fn record(
261        &self,
262        encoder: &mut wgpu::CommandEncoder,
263        view: &wgpu::TextureView,
264        depth_view: Option<&wgpu::TextureView>,
265        viewport: crate::Viewport,
266    ) {
267        match (self.has_depth, depth_view.is_some()) {
268            (true, false) => panic!(
269                "RayMarchNode::record: the pipeline has a depth format but no depth view was given"
270            ),
271            (false, true) => panic!(
272                "RayMarchNode::record: the pipeline has no depth format but a depth view was given"
273            ),
274            _ => {}
275        }
276        let depth_stencil_attachment = depth_view.map(|dv| RenderPassDepthStencilAttachment {
277            view: dv,
278            depth_ops: Some(Operations {
279                load: LoadOp::Clear(crate::view::DEPTH_CLEAR),
280                store: StoreOp::Store,
281            }),
282            stencil_ops: None,
283        });
284        let mut rp = encoder.begin_render_pass(&RenderPassDescriptor {
285            label: Some("raymarch pass"),
286            color_attachments: &[Some(RenderPassColorAttachment {
287                view,
288                depth_slice: None,
289                resolve_target: None,
290                ops: Operations {
291                    load: LoadOp::Clear(self.clear_color),
292                    store: StoreOp::Store,
293                },
294            })],
295            depth_stencil_attachment,
296            timestamp_writes: None,
297            occlusion_query_set: None,
298        });
299        viewport.apply(&mut rp);
300        rp.set_pipeline(&self.pipeline);
301        rp.set_bind_group(0, &self.bind_group, &[]);
302        rp.draw(0..3, 0..1);
303    }
304}