Rust

Voronoi shading the web

Introduction

Playing with GPU shaders has always been one of my main hobbies. It started out with OpenGL after a introductory graphics course during my bachelors. It peaked when Vulkan came out and scratched that part of my brain that yearns for well designed and documented APIs (Not to mention RTX support!). And now, frustrated with the fragility of C++, I came to Rust, eluded by the promise of Bevy. Whilst still in the honeymoon phase I discovered that it was not only possible, but near trivial to compile a bevy game to work in the browser! Now I can finally convey my shader-love directly in my posts! Important info: I used Bevy 0.8.1, it is still changing a lot so this post may not be a proper implementation reference in the future.

The goal: A Voronoi diagram

Some may remember the concept from high school maths. The idea is take a number of points on a 2d plane and then cut that plane up in cells such that each pixel on each cell is closest to its respective point. We can then give each cell its own color for stylistic effect. Obviously the cherry on the cake would be to animate the points in real-time.

Creating a 2d material and quad

I personally really dislike working with WGSL (the suggested replacement for GLSL when it comes to web), so I chose to implement the material using GLSL shaders. All we have to do is create a struct containing our set of points, let say a 100 of them, and them derive the Material2d trait. There is a crux however, because in order for the array to be mapped correctly to shader memory, each array element needs to be 16 byte aligned. We therefore allocate a Vec4 for each point, even though we only need 2.

// This is the struct that will be passed to your shader
#[derive(AsBindGroup, TypeUuid, Debug, Clone, ShaderType)]
#[uuid = "f690fdae-d598-45ab-8225-97e2a3f056e0"]
pub struct CustomMaterial {
    #[uniform(0)]
    points: [Vec4; 100],
}

full source

We can then inform bevy that we wish to override the shader pipeline with our custom shaders. Because GLSL is not the default we also need to override the entry point during specialisation.

impl Material2d for CustomMaterial {
    fn fragment_shader() -> ShaderRef {
        "shaders/custom_material.frag".into()
    }

    fn vertex_shader() -> ShaderRef {
        "shaders/custom_material.vert".into()
    }

    fn specialize(
        descriptor: &mut bevy::render::render_resource::RenderPipelineDescriptor,
        layout: &bevy::render::mesh::MeshVertexBufferLayout,
        key: bevy::sprite::Material2dKey<Self>,
    ) -> Result<(), bevy::render::render_resource::SpecializedMeshPipelineError> {
        descriptor.vertex.entry_point = "main".into();
        descriptor.fragment.as_mut().unwrap().entry_point = "main".into();
        Ok(())
    }
}

full source

The vertex shader is a simple one; we take the mesh inputs from bevy and simply set the gl_Position and extract the screenspace uv coordinates.

#version 450

layout(location = 0) in vec3 Vertex_Position;
layout(location = 1) in vec3 Vertex_Normal;
layout(location = 2) in vec2 Vertex_Uv;

layout(location = 0) out vec2 v_Uv;

layout(set = 0, binding = 0) uniform CameraViewProj {
    mat4 ViewProj;
    mat4 View;
    mat4 InverseView;
    mat4 Projection;
    vec3 WorldPosition;
    float width;
    float height;
};

layout(set = 2, binding = 0) uniform Mesh {
    mat4 Model;
    mat4 InverseTransposeModel;
    uint flags;
};

void main() {
    v_Uv = Vertex_Uv;
    gl_Position = ViewProj * Model * vec4(Vertex_Position, 1.0);
}

full source

Finally in the fragment shader we can implement our Voronoi logic. We simply iterate over all 100 points to find the point that has the smallest distance to our pixel (Note that we work in normalized screenspace for simplicity). We then sample a pleasant color in HSV color space and convert it to RGB. Finally we create a slight highlight at the location of the original point (where the distance is small) to make it look and feel a bit less flat.

#version 450
layout(location = 0) in vec2 v_Uv;

layout(location = 0) out vec4 o_Target;

layout(set = 1, binding = 0) uniform Points {
  vec4 data[100];
} pts;

uint rand_xorshift(uint seed)
{
    // Xorshift algorithm from George Marsaglia's paper
    seed ^= (seed << 13);
    seed ^= (seed >> 17);
    seed ^= (seed << 5);
    return seed;
}

// generate a random float
float rand(inout uint seed)
{
    seed = rand_xorshift(seed);
    return seed * 2.3283064365387e-10f;
}

vec3 hsv2rgb(vec3 c)
{
    vec4 K = vec4(1.0, 2.0 / 3.0, 1.0 / 3.0, 3.0);
    vec3 p = abs(fract(c.xxx + K.xyz) * 6.0 - K.www);
    return c.z * mix(K.xxx, clamp(p - K.xxx, 0.0, 1.0), c.y);
}

void main() {
  float min_dis = 100000;
  uint min_id = 0;
  for(uint i=0; i<100; i++) {
    float dis = length(pts.data[i].xy - v_Uv);
    if (dis < min_dis) {
      min_dis = dis;
      min_id = i;
    }
  }

  uint seed = rand_xorshift(min_id);
  float hue = rand(seed);
  float saturation = 0.4 + 0.6 * rand(seed);
  float value = 0.2 + 0.8 * rand(seed);

  float intensity = pow(1.0 - min_dis, 10.5);
  vec3 color = intensity * hsv2rgb(vec3(hue, saturation, value));

  o_Target = vec4(color, 1.0);
}

full source

Creating a resource for animation

To animate the app we need to alter the points in the material, as well as keep track of the velocities of each point. We don't actually need the velocities in the shader so I kept them separate from the material and created a resource instead:

struct MatResource {
    handle: Handle<CustomMaterial>,
    velocities: [Vec2; 100],
}

full source

Startup system

Here we have the startup system 'setup', I will let the inline comments speak for themself.

fn setup(
    mut commands: Commands,
    mut windows: ResMut<Windows>,
    mut meshes: ResMut<Assets<Mesh>>,
    mut materials: ResMut<Assets<CustomMaterial>>,
    asset_server: Res<AssetServer>,
) {
    # Create a full screen quad
    let window = windows.get_primary_mut().unwrap();

    let quad = meshes.add(Mesh::from(shape::Quad::new(Vec2::new(
        window.physical_width() as f32,
        window.physical_height() as f32,
    ))));

    // Generate an initial set of points at random locations
    let mut points = [Vec4::ZERO; 100];
    for point in points.iter_mut() {
        point.x = rand::random::<f32>();
        point.y = rand::random::<f32>();
    }

    // Generate an initial set of velocities as random vectors
    // (not that they are not actually uniformly random but oh well)
    let mut velocities = [Vec2::ZERO; 100];
    for v  in velocities.iter_mut() {
        v.x = rand::random::<f32>() * 2.0 - 1.0;
        v.y = rand::random::<f32>() * 2.0 - 1.0;
        *v = v.normalize();
        v.x *= 0.002;
        v.y *= 0.002;
    }

    // register the material
    let material = materials.add(CustomMaterial {
        points,
    });

    // register the resource for updating 
    commands.insert_resource(MatResource {
        handle: material.clone(),
        velocities,
    });

    // Spawn the fullscreen quad with our material
    commands
        .spawn_bundle(MaterialMesh2dBundle {
            mesh: quad.into(),
            material,
            transform: Transform {
                translation: Vec3::new(0.0, 0.0, 1.5),
                ..default()
            },
            ..default()
        });

    // camera
    commands.spawn_bundle(Camera2dBundle::default());
}

full source

Updating the material

Each frame we want the points to move by their velocity. Should they hit the edge of the screen I will simply let them bounce.

fn tick(
    mut materials: ResMut<Assets<CustomMaterial>>,
    mut mat_res: ResMut<MatResource>) {
    if let Some(mat) = materials.get_mut(&mat_res.handle) {
        for (p, v) in mat.points.iter_mut().zip(mat_res.velocities.iter_mut()) {
            p.x = p.x + v.x;
            if p.x < 0.0 { p.x = 0.0; v.x *= -1.0; }
            if p.x > 1.0 { p.x = 1.0; v.x *= -1.0; }
            p.y = p.y + v.y;
            if p.y < 0.0 { p.y = 0.0; v.y *= -1.0; }
            if p.y > 1.0 { p.y = 1.0; v.y *= -1.0; }
        }
    }
}

full source

Jamming it together

Now all that is left is to combine our systems and materials in a bevy app that we can run.

fn main() {
    App::new()
        // We want our window to be a specific size
        // we also specify a html element where the app
        // should run in if compiled for the web
        .insert_resource(WindowDescriptor {
            width: 640.0,
            height: 480.0,
            canvas: Some("#shader_demo".into()),
            ..default()
        })
        .add_plugins(DefaultPlugins)
        .add_plugin(Material2dPlugin::<CustomMaterial>::default())
        .add_startup_system(setup)
        .add_system(tick)
        .run();
}

full source

Finally we can just Compile it to WASM. And we get a simple javascript file that we can include and execute. If you want to get a better look at the source, you can find it over on my Github. Feedback welcome, as always :)

last modified: May 29, 2023 at 13:07