Skip to main content

compute/
compute.rs

1//! A compute pass: rain on a pond. A wave simulation runs on the GPU each
2//! frame in a render target (height now in red, the frame before in
3//! green), and the water's shader reads it to bend its normals.
4//!
5//!   `cargo run --release --example compute --features 3d`
6//!
7//! Click the water to drop a stone in. `SHOT=path.png` saves frame 120
8//! and quits.
9
10use ponko::*;
11
12/// The wave equation on a grid: each texel moves toward the average of
13/// its neighbours, overshooting like water does, and calms slowly.
14/// material.custom[0]: a drop (uv x, uv y, strength; 0 for none).
15const RIPPLES: &str = "
16fn compute(c: Cell) -> vec4<f32> {
17    let around = cell_at(c.pixel + vec2<i32>(1, 0)).x + cell_at(c.pixel - vec2<i32>(1, 0)).x
18        + cell_at(c.pixel + vec2<i32>(0, 1)).x + cell_at(c.pixel - vec2<i32>(0, 1)).x;
19    var height = (around * 0.5 - c.value.y) * 0.985;
20    let drop = material.custom[0];
21    let r = distance(c.uv, drop.xy) * f32(c.size.x);
22    if drop.z > 0.0 && r < 3.0 {
23        height -= drop.z * (1.0 - r / 3.0);
24    }
25    return vec4<f32>(height, c.value.x, 0.0, 0.0);
26}";
27
28/// The water: its normals follow the simulated heights (custom_texture0).
29const WATER: &str = "
30fn surface(s: Surface) -> Surface {
31    var out = s;
32    let texel = 1.0 / vec2<f32>(textureDimensions(custom_texture0));
33    let x = textureSample(custom_texture0, material_sampler, s.uv + vec2<f32>(texel.x, 0.0)).x
34        - textureSample(custom_texture0, material_sampler, s.uv - vec2<f32>(texel.x, 0.0)).x;
35    let z = textureSample(custom_texture0, material_sampler, s.uv + vec2<f32>(0.0, texel.y)).x
36        - textureSample(custom_texture0, material_sampler, s.uv - vec2<f32>(0.0, texel.y)).x;
37    out.normal = normalize(vec3<f32>(-x * 8.0, 1.0, -z * 8.0));
38    out.roughness = 0.06;
39    return out;
40}";
41
42const SIZE: f32 = 12.0;
43
44fn main() {
45    let mut seed = 11u32;
46    let mut random = move || {
47        seed = seed.wrapping_mul(1_664_525).wrapping_add(1_013_904_223);
48        (seed >> 8) as f32 / (1 << 24) as f32
49    };
50    let mut next_drop = 0.0;
51    Game::new("Crabby Engine: compute", 1280, 720)
52        .system(move |ctx: &mut Ctx| {
53            if ctx.key_pressed(Key::Escape) {
54                ctx.quit();
55            }
56            let Some(pass) = ctx.compute_passes.first().copied() else {
57                return;
58            };
59            let Some(drops) = pass.material else { return };
60            // Rain every so often, or a stone where the mouse clicks.
61            let mut drop = Vec4::ZERO;
62            next_drop -= ctx.dt();
63            if next_drop <= 0.0 {
64                next_drop = 0.05 + random() * 0.2;
65                drop = vec4(random(), random(), 0.3 + random() * 0.4, 0.0);
66            }
67            if ctx.mouse_pressed(MouseButton::Left) {
68                let (from, dir) = ctx.camera3d.screen_ray(ctx.mouse_pos());
69                if dir.y < 0.0 {
70                    let at = from + dir * (-from.y / dir.y);
71                    drop = vec4(at.x / SIZE + 0.5, at.z / SIZE + 0.5, 3.0, 0.0);
72                }
73            }
74            ctx.update_material(
75                drops,
76                MaterialDesc {
77                    custom: [drop, Vec4::ZERO, Vec4::ZERO, Vec4::ZERO],
78                    ..Default::default()
79                },
80            );
81            if ctx.frame() == 120
82                && let Ok(path) = std::env::var("SHOT")
83            {
84                ctx.save_screenshot(path);
85            }
86            if ctx.frame() == 122 && std::env::var("SHOT").is_ok() {
87                ctx.quit();
88            }
89        })
90        .run(|ctx| {
91            ctx.camera3d.position = vec3(0.0, 7.0, 10.0);
92            ctx.camera3d.look_at(vec3(0.0, 0.0, -1.0));
93            ctx.lighting.sun_direction = Sky::sun_direction_at(15.0);
94            ctx.post.bloom = 0.15;
95
96            let heights = ctx.create_render_target(512, 512);
97            let ripples = ctx.create_compute("ripples", RIPPLES);
98            let drops = ctx.create_material(MaterialDesc::default());
99            ctx.compute_passes
100                .push(ComputePass::new(ripples, heights).with_material(drops));
101
102            let shader = ctx.create_shader("water", WATER);
103            let water = ctx.create_material(MaterialDesc {
104                shader: Some(shader),
105                custom_textures: [Some(heights), None],
106                ..MaterialDesc::color(Color::hex(0x2A5A75))
107            });
108            let surface = ctx.create_mesh(&MeshData::plane(SIZE), water);
109            ctx.world.spawn((Transform3D::default(), Model::new(surface)));
110
111            let stone = ctx.create_material(MaterialDesc {
112                roughness: 0.8,
113                ..MaterialDesc::color(Color::hex(0x8A8578))
114            });
115            let rock = ctx.create_mesh(&MeshData::sphere(0.6, 16), stone);
116            for i in 0..40 {
117                let a = i as f32 / 40.0 * std::f32::consts::TAU;
118                let edge = vec3(a.cos(), 0.0, a.sin()) * SIZE * 0.55;
119                let t = Transform3D {
120                    scale: Vec3::splat(0.8 + (i % 3) as f32 * 0.3),
121                    ..Transform3D::at(edge)
122                };
123                ctx.world.spawn((t, Model::new(rock)));
124            }
125            let grass = ctx.create_material(MaterialDesc {
126                roughness: 0.95,
127                ..MaterialDesc::color(Color::hex(0x4E6B33))
128            });
129            let ground = ctx.create_mesh(&MeshData::plane(200.0), grass);
130            ctx.world
131                .spawn((Transform3D::at(vec3(0.0, -0.3, 0.0)), Model::new(ground)));
132        });
133}