Skip to main content

orbits/
orbits.rs

1use cuneus::prelude::ComputeShader;
2use cuneus::{Core, ExportManager, RenderKit, ShaderApp, ShaderControls, ShaderManager, UniformProvider};
3use cuneus::WindowEvent;
4use winit::event::{ElementState, MouseButton, MouseScrollDelta};
5#[repr(C)]
6#[derive(Copy, Clone, Debug, bytemuck::Pod, bytemuck::Zeroable)]
7pub struct ShaderParams {
8    base_color: [f32; 3],
9    x: f32,
10    rim_color: [f32; 3],
11    y: f32,
12    accent_color: [f32; 3],
13    gamma_correction: f32,
14    travel_speed: f32,
15    iteration: i32,
16    col_ext: f32,
17    zoom: f32,
18    trap_pow: f32,
19    trap_x: f32,
20    trap_y: f32,
21    trap_c1: f32,
22    aa: i32,
23    trap_s1: f32,
24    wave_speed: f32,
25    fold_intensity: f32,
26}
27
28impl UniformProvider for ShaderParams {
29    fn as_bytes(&self) -> &[u8] {
30        bytemuck::bytes_of(self)
31    }
32}
33
34struct Shader {
35    base: RenderKit,
36    compute_shader: ComputeShader,
37    mouse_dragging: bool,
38    drag_start: [f32; 2],
39    drag_start_pos: [f32; 2],
40    zoom_level: f32,
41    current_params: ShaderParams,
42}
43fn main() -> Result<(), Box<dyn std::error::Error>> {
44    env_logger::init();
45    let (app, event_loop) = ShaderApp::new("orbits", 800, 600);
46    app.run(event_loop, Shader::init)
47}
48impl ShaderManager for Shader {
49    fn init(core: &Core) -> Self {
50        let initial_zoom = 0.0004;
51        let initial_x = 2.14278;
52        let initial_y = 2.14278;
53
54        // Create texture display layout
55        let base = RenderKit::new(core);
56
57        let config = ComputeShader::builder()
58            .with_entry_point("main")
59            .with_custom_uniforms::<ShaderParams>()
60            .with_mouse()
61            .build();
62
63        let compute_shader = cuneus::compute_shader!(core, "shaders/orbits.wgsl", config);
64
65        let initial_params = ShaderParams {
66            base_color: [0.0, 0.5, 1.0],
67            x: initial_x,
68            rim_color: [0.0, 0.5, 1.0],
69            y: initial_y,
70            accent_color: [0.018, 0.018, 0.018],
71            gamma_correction: 0.4,
72            travel_speed: 1.0,
73            iteration: 355,
74            col_ext: 2.0,
75            zoom: initial_zoom,
76            trap_pow: 1.0,
77            trap_x: -0.5,
78            trap_y: 2.0,
79            trap_c1: 0.2,
80            aa: 1,
81            trap_s1: 0.8,
82            wave_speed: 0.1,
83            fold_intensity: 1.0,
84        };
85
86        compute_shader.set_custom_params(initial_params, &core.queue);
87
88        Self {
89            base,
90            compute_shader,
91            mouse_dragging: false,
92            drag_start: [0.0, 0.0],
93            drag_start_pos: [initial_x, initial_y],
94            zoom_level: initial_zoom,
95            current_params: initial_params,
96        }
97    }
98
99    fn update(&mut self, core: &Core) {
100        // Handle export
101        self.compute_shader.handle_export(core, &mut self.base);
102    }
103
104    fn render(&mut self, core: &Core) -> Result<(), cuneus::SurfaceError> {
105        let mut frame = self.base.begin_frame(core)?;
106
107        let mut params = self.current_params;
108        let mut changed = false;
109        let mut should_start_export = false;
110        let mut export_request = self.base.export_manager.get_ui_request();
111
112        let mut controls_request = self
113            .base
114            .controls
115            .get_ui_request(&self.base.start_time, &core.size, self.base.fps_tracker.fps());
116
117        let full_output = if self.base.key_handler.show_ui {
118            self.base.render_ui(core, |ctx| {
119                RenderKit::apply_default_style(ctx);
120                egui::Window::new("Orbits")
121                    .collapsible(true)
122                    .resizable(true)
123                    .default_width(280.0)
124                    .show(ctx, |ui| {
125                        egui::CollapsingHeader::new("Colors")
126                            .default_open(false)
127                            .show(ui, |ui| {
128                                ui.horizontal(|ui| {
129                                    ui.label("Base:");
130                                    changed |=
131                                        ui.color_edit_button_rgb(&mut params.base_color).changed();
132                                });
133                                ui.horizontal(|ui| {
134                                    ui.label("Orbit:");
135                                    changed |=
136                                        ui.color_edit_button_rgb(&mut params.rim_color).changed();
137                                });
138                                ui.horizontal(|ui| {
139                                    ui.label("Exterior:");
140                                    changed |= ui
141                                        .color_edit_button_rgb(&mut params.accent_color)
142                                        .changed();
143                                });
144                            });
145
146                        egui::CollapsingHeader::new("Rendering")
147                            .default_open(false)
148                            .show(ui, |ui| {
149                                changed |= ui
150                                    .add(
151                                        egui::Slider::new(&mut params.iteration, 50..=500)
152                                            .text("Iterations"),
153                                    )
154                                    .changed();
155                                changed |= ui
156                                    .add(
157                                        egui::Slider::new(&mut params.aa, 1..=4)
158                                            .text("Anti-aliasing"),
159                                    )
160                                    .changed();
161                                changed |= ui
162                                    .add(
163                                        egui::Slider::new(&mut params.gamma_correction, 0.1..=2.0)
164                                            .text("Gamma"),
165                                    )
166                                    .changed();
167                            });
168
169                        egui::CollapsingHeader::new("Traps")
170                            .default_open(false)
171                            .show(ui, |ui| {
172                                changed |= ui
173                                    .add(
174                                        egui::Slider::new(&mut params.trap_x, -5.0..=5.0)
175                                            .text("Trap X"),
176                                    )
177                                    .changed();
178                                changed |= ui
179                                    .add(
180                                        egui::Slider::new(&mut params.trap_y, -5.0..=5.0)
181                                            .text("Trap Y"),
182                                    )
183                                    .changed();
184                                changed |= ui
185                                    .add(
186                                        egui::Slider::new(&mut params.trap_pow, 0.0..=3.0)
187                                            .text("Trap Power"),
188                                    )
189                                    .changed();
190                                changed |= ui
191                                    .add(
192                                        egui::Slider::new(&mut params.trap_c1, 0.0..=1.0)
193                                            .text("Trap Mix"),
194                                    )
195                                    .changed();
196                                changed |= ui
197                                    .add(
198                                        egui::Slider::new(&mut params.trap_s1, 0.0..=2.0)
199                                            .text("Trap Blend"),
200                                    )
201                                    .changed();
202                            });
203
204                        egui::CollapsingHeader::new("Animation")
205                            .default_open(false)
206                            .show(ui, |ui| {
207                                changed |= ui
208                                    .add(
209                                        egui::Slider::new(&mut params.travel_speed, 0.0..=2.0)
210                                            .text("Travel Speed"),
211                                    )
212                                    .changed();
213                                changed |= ui
214                                    .add(
215                                        egui::Slider::new(&mut params.wave_speed, 0.0..=2.0)
216                                            .text("Wave Speed"),
217                                    )
218                                    .changed();
219                                changed |= ui
220                                    .add(
221                                        egui::Slider::new(&mut params.fold_intensity, 0.0..=3.0)
222                                            .text("Fold Intensity"),
223                                    )
224                                    .changed();
225                                changed |= ui
226                                    .add(
227                                        egui::Slider::new(&mut params.col_ext, 0.0..=10.0)
228                                            .text("Color Extension"),
229                                    )
230                                    .changed();
231                            });
232
233                        egui::CollapsingHeader::new("Navigation")
234                            .default_open(false)
235                            .show(ui, |ui| {
236                                ui.label("Left-click + drag: Pan view");
237                                ui.label("Mouse wheel: Zoom");
238                                ui.separator();
239                                let old_zoom = params.zoom;
240                                changed |= ui
241                                    .add(
242                                        egui::Slider::new(&mut params.zoom, 0.0001..=1.0)
243                                            .text("Zoom")
244                                            .logarithmic(true),
245                                    )
246                                    .changed();
247                                if old_zoom != params.zoom {
248                                    self.zoom_level = params.zoom;
249                                }
250                                changed |= ui
251                                    .add(
252                                        egui::Slider::new(&mut params.x, 0.0..=3.0)
253                                            .text("X Position"),
254                                    )
255                                    .changed();
256                                changed |= ui
257                                    .add(
258                                        egui::Slider::new(&mut params.y, 0.0..=6.0)
259                                            .text("Y Position"),
260                                    )
261                                    .changed();
262                            });
263
264                        ui.separator();
265                        ShaderControls::render_controls_widget(ui, &mut controls_request);
266                        ui.separator();
267                        should_start_export =
268                            ExportManager::render_export_ui_widget(ui, &mut export_request);
269                    });
270            })
271        } else {
272            self.base.render_ui(core, |_ctx| {})
273        };
274
275        self.base.export_manager.apply_ui_request(export_request);
276        self.base.apply_control_request(controls_request);
277
278        if changed {
279            self.current_params = params;
280            self.compute_shader.set_custom_params(params, &core.queue);
281        }
282
283        if should_start_export {
284            self.base.export_manager.start_export();
285        }
286
287        // Create command encoder
288
289        // Update time uniform
290        let current_time = self.base.controls.get_time(&self.base.start_time);
291        let delta_time = 1.0 / 60.0;
292        self.compute_shader
293            .set_time(current_time, delta_time, &core.queue);
294
295        // Update mouse uniform
296        self.compute_shader
297            .update_mouse_uniform(&self.base.mouse_tracker.uniform, &core.queue);
298
299        // Dispatch compute shader
300        self.compute_shader.dispatch(&mut frame.encoder, core);
301
302        self.base.renderer.render_to_view(&mut frame.encoder, &frame.view, &self.compute_shader.get_output_texture().bind_group);
303
304        self.base.end_frame(core, frame, full_output);
305
306        Ok(())
307    }
308    fn resize(&mut self, core: &Core) {
309        self.base.default_resize(core, &mut self.compute_shader);
310    }
311    fn handle_input(&mut self, core: &Core, event: &WindowEvent) -> bool {
312        if self.base.default_handle_input(core, event) {
313            return true;
314        }
315        match event {
316            WindowEvent::MouseInput { state, button, .. } => {
317                if button == &MouseButton::Left {
318                    match state {
319                        ElementState::Pressed => {
320                            let mouse_pos = self.base.mouse_tracker.uniform.position;
321                            self.mouse_dragging = true;
322                            self.drag_start = mouse_pos;
323                            self.drag_start_pos = [self.current_params.x, self.current_params.y];
324                            return true;
325                        }
326                        ElementState::Released => {
327                            self.mouse_dragging = false;
328                            return true;
329                        }
330                    }
331                }
332                false
333            }
334            WindowEvent::CursorMoved { .. } => {
335                if self.mouse_dragging {
336                    let current_pos = self.base.mouse_tracker.uniform.position;
337                    let dx = (current_pos[0] - self.drag_start[0]) * 3.0 * self.zoom_level;
338                    let dy = (current_pos[1] - self.drag_start[1]) * 6.0 * self.zoom_level;
339                    let mut new_x = self.drag_start_pos[0] + dx;
340                    let mut new_y = self.drag_start_pos[1] + dy;
341                    new_x = new_x.clamp(0.0, 3.0);
342                    new_y = new_y.clamp(0.0, 6.0);
343                    self.current_params.x = new_x;
344                    self.current_params.y = new_y;
345                    self.compute_shader
346                        .set_custom_params(self.current_params, &core.queue);
347                }
348                self.base.handle_mouse_input(core, event, false)
349            }
350            WindowEvent::MouseWheel { delta, .. } => {
351                let zoom_delta = match delta {
352                    MouseScrollDelta::LineDelta(_, y) => *y * 0.1,
353                    MouseScrollDelta::PixelDelta(pos) => (pos.y as f32) * 0.001,
354                };
355
356                if zoom_delta != 0.0 {
357                    let mouse_pos = self.base.mouse_tracker.uniform.position;
358                    let center_x = self.current_params.x;
359                    let center_y = self.current_params.y;
360
361                    let rel_x = mouse_pos[0] - 0.5;
362                    let rel_y = mouse_pos[1] - 0.5;
363
364                    let zoom_factor = if zoom_delta > 0.0 { 0.9 } else { 1.1 };
365                    self.zoom_level = (self.zoom_level * zoom_factor).clamp(0.0001, 1.5);
366
367                    let scale_change = 1.0 - zoom_factor;
368                    let dx = rel_x * scale_change * 3.0 * self.zoom_level;
369                    let dy = rel_y * scale_change * 6.0 * self.zoom_level;
370                    self.current_params.zoom = self.zoom_level;
371                    self.current_params.x = (center_x + dx).clamp(0.0, 3.0);
372                    self.current_params.y = (center_y + dy).clamp(0.0, 6.0);
373                    self.compute_shader
374                        .set_custom_params(self.current_params, &core.queue);
375                }
376                self.base.handle_mouse_input(core, event, false)
377            }
378
379            _ => self.base.handle_mouse_input(core, event, false),
380        }
381    }
382}