Bevy game of life in fragment shader

main.rs

// License: Apache-2.0 / MIT

use bevy::{
    prelude::*,
    reflect::TypeUuid,
    render::{
        camera::{Camera, RenderTarget},
        render_resource::{
            AsBindGroup, Extent3d, ShaderRef, TextureDescriptor, TextureDimension, TextureFormat,
            TextureUsages,
        },
    },
    sprite::{Material2d, Material2dPlugin, MaterialMesh2dBundle},
};

fn main() {
    let mut app = App::new();
    app.add_plugins(DefaultPlugins)
        .insert_resource(ClearColor(Color::BLACK))
        .add_plugin(Material2dPlugin::<LifeMaterial>::default())
        .add_startup_system(setup)
        .add_system(swap_buffers)
        .add_system(update_handles)
        .add_system(un_init);

    app.run();
}

#[derive(Component)]
struct DoubleBuffer {
    buf_a: Handle<Image>,
    buf_b: Handle<Image>,
    phase: bool,
}

impl DoubleBuffer {
    fn swap(&mut self) {
        self.phase = !self.phase;
    }
    fn get(&self) -> Handle<Image> {
        if self.phase {
            self.buf_a.clone()
        } else {
            self.buf_b.clone()
        }
    }
}

fn swap_buffers(mut buffers: Query<&mut DoubleBuffer>) {
    for mut double_buffer in buffers.iter_mut() {
        double_buffer.swap()
    }
}

fn update_handles(
    mut buf_cameras: Query<(&mut Camera, &DoubleBuffer), Without<Handle<LifeMaterial>>>,
    mut buf_materials: Query<(&Handle<LifeMaterial>, &DoubleBuffer), Without<Camera>>,
    mut materials: ResMut<Assets<LifeMaterial>>,
) {
    for (mut camera, double_buffer) in buf_cameras.iter_mut() {
        camera.target = RenderTarget::Image(double_buffer.get())
    }
    for (material, double_buffer) in buf_materials.iter_mut() {
        let mut material = materials.get_mut(&material).unwrap();
        material.source_image = double_buffer.get()
    }
}

fn un_init(mut materials: ResMut<Assets<LifeMaterial>>) {
    for (_, mat) in materials.iter_mut() {
        mat.frame += 1.0;
    }
}

fn setup(
    mut commands: Commands,
    mut windows: ResMut<Windows>,
    mut meshes: ResMut<Assets<Mesh>>,
    mut post_processing_materials: ResMut<Assets<LifeMaterial>>,
    mut images: ResMut<Assets<Image>>,
    asset_server: Res<AssetServer>,
) {
    asset_server.watch_for_changes().unwrap();

    let window = windows.get_primary_mut().unwrap();
    let size = Extent3d {
        width: window.physical_width(),
        height: window.physical_height(),
        ..default()
    };

    let texture_descriptor = TextureDescriptor {
        label: None,
        size,
        dimension: TextureDimension::D2,
        format: TextureFormat::Bgra8UnormSrgb,
        mip_level_count: 1,
        sample_count: 1,
        usage: TextureUsages::TEXTURE_BINDING
            | TextureUsages::COPY_DST
            | TextureUsages::RENDER_ATTACHMENT,
    };

    let mut image_a = Image {
        texture_descriptor: texture_descriptor.clone(),
        ..default()
    };

    image_a.resize(size);
    let image_b = image_a.clone();

    let image_handle_a = images.add(image_a);
    let image_handle_b = images.add(image_b);

    // Main camera, first to render
    commands
        .spawn_bundle(Camera2dBundle {
            camera: Camera {
                target: RenderTarget::Image(image_handle_a.clone()),
                ..default()
            },
            ..default()
        })
        .insert(DoubleBuffer {
            buf_a: image_handle_a.clone(),
            buf_b: image_handle_b.clone(),
            phase: false,
        });

    let quad_handle = meshes.add(Mesh::from(shape::Quad::new(Vec2::new(
        size.width as f32,
        size.height as f32,
    ))));

    // This material has the texture that has been rendered.
    let material_handle = post_processing_materials.add(LifeMaterial {
        frame: 0.0,
        source_image: image_handle_b.clone(),
    });

    commands
        .spawn_bundle(MaterialMesh2dBundle {
            mesh: quad_handle.into(),
            material: material_handle,
            transform: Transform {
                translation: Vec3::new(0.0, 0.0, 1.5),
                ..default()
            },
            ..default()
        })
        .insert(DoubleBuffer {
            buf_a: image_handle_a.clone(),
            buf_b: image_handle_b.clone(),
            phase: true,
        });

    commands.spawn_bundle(Camera2dBundle {
        camera: Camera {
            // renders after the first main camera which has default value: 0.
            priority: 1,
            ..default()
        },
        ..Camera2dBundle::default()
    });
}

#[derive(AsBindGroup, TypeUuid, Clone)]
#[uuid = "bc2f08eb-a0fb-43f1-a908-54871ea597d5"]
struct LifeMaterial {
    #[uniform(0)]
    frame: f32,
    /// In this example, this image will be the result of the main camera.
    #[texture(1)]
    source_image: Handle<Image>,
}

impl Material2d for LifeMaterial {
    fn fragment_shader() -> ShaderRef {
        "shaders/game_of_life.wgsl".into()
    }
}

// put in shaders/game_of_life.wgsl:
/*
#import bevy_pbr::mesh_view_bindings

struct Material {
    frame: f32,
};

@group(1) @binding(0)
var<uniform> material: Material;
@group(1) @binding(1)
var texture: texture_2d<f32>;

fn hash(value: u32) -> u32 {
    var state = value;
    state = state ^ 2747636419u;
    state = state * 2654435769u;
    state = state ^ state >> 16u;
    state = state * 2654435769u;
    state = state ^ state >> 16u;
    state = state * 2654435769u;
    return state;
}

fn randomFloat(value: u32) -> f32 {
    return f32(hash(value)) / 4294967295.0;
}

fn is_alive(location: vec2<i32>, offset_x: i32, offset_y: i32) -> i32 {
    let value = textureLoad(texture, location + vec2<i32>(offset_x, offset_y), 0);
    return i32(value.x);
}

fn count_alive(location: vec2<i32>) -> i32 {
    return is_alive(location, -1, -1) +
           is_alive(location, -1,  0) +
           is_alive(location, -1,  1) +
           is_alive(location,  0, -1) +
           is_alive(location,  0,  1) +
           is_alive(location,  1, -1) +
           is_alive(location,  1,  0) +
           is_alive(location,  1,  1);
}

@fragment
fn fragment(
    @builtin(position) position: vec4<f32>,
    #import bevy_sprite::mesh2d_vertex_output
) -> @location(0) vec4<f32> {
    let ilocation = vec2<i32>(position.xy);
    var output_color = vec4<f32>(0.0);

    if material.frame < 30.0 {

        let uv = position.xy / vec2<f32>(view.width, view.height);
        let randomNumber = randomFloat(u32(position.x * position.y));
        let alive = randomNumber > 0.9;
        output_color = vec4<f32>(f32(alive));

    } else {

        let n_alive = count_alive(ilocation);

        var alive: bool;
        if (n_alive == 3) {
            alive = true;
        } else if (n_alive == 2) {
            let currently_alive = is_alive(ilocation, 0, 0);
            alive = bool(currently_alive);
        } else {
            alive = false;
        }
        output_color = vec4<f32>(f32(alive));
    }

    return output_color;
}
*/