Mandelbrot with bevy custom Material

* Left click + drag to pan around
* Mouse wheel to zoom
* Double left click to reset view
* Right click to change start value, distorts fractal
* Double right click to reset start value

Cargo.toml

@@ -0,0 +1,11 @@
+[package]
+name = "bevy_mandelbrot"
+version = "0.1.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+bevy = "0.7"
+#bevy_clicking = { path = "../clicking" }
+bevy_clicking = { git = "https://github.com/matelab/bevy_clicking" }

assets/shaders/mandelbrot.wgsl

@@ -0,0 +1,38 @@
+struct MandelbrotFS {
+    center: vec2<f32>;
+    start: vec2<f32>;
+    scale: f32;
+    aspect: f32;
+    iters: i32;
+};
+
+[[group(1), binding(0)]]
+var<uniform> fs: MandelbrotFS;
+
+[[stage(fragment)]]
+fn fragment([[location(2)]] uv: vec2<f32>) -> [[location(0)]] vec4<f32> {
+    var z: vec2<f32>;
+    var i: i32;
+    let iters = fs.iters;
+    z = fs.start;
+    var p: vec2<f32>;
+    p = vec2<f32>((fs.aspect * (uv.x - 0.5)) / fs.scale + fs.center.x, (uv.y - 0.5) / fs.scale + fs.center.y);
+    for (i = 0; i < iters; i = i + 1) {
+        let x = (z.x * z.x - z.y * z.y) + p.x;
+        let y = (2.0 * z.x * z.y) + p.y;
+
+        if (((x * x) + (y * y)) > 4.0) {
+            break;
+        }
+        z.x = x;
+        z.y = y;
+    }
+    var col: f32;
+    if (i == iters) {
+        col = 0.0;
+    } else {
+        col = f32(i) / f32(iters);
+    }
+    return vec4<f32>(col, col, col, 1.0);
+
+}

src/colormap.rs

@@ -0,0 +1,216 @@
+use bevy::{
+    core_pipeline::node::MAIN_PASS_DEPENDENCIES,
+    ecs::system::{lifetimeless::SRes, SystemParamItem},
+    prelude::*,
+    reflect::TypeUuid,
+    render::{
+        render_asset::{PrepareAssetError, RenderAsset, RenderAssetPlugin, RenderAssets},
+        render_component::ExtractComponentPlugin,
+        render_graph::{self, RenderGraph},
+        render_resource::std140::{AsStd140, Std140},
+        render_resource::*,
+        renderer::{RenderContext, RenderDevice},
+        RenderApp, RenderStage,
+    },
+};
+
+pub struct ColormapPlugin {
+    prev_node: &'static str,
+}
+
+impl ColormapPlugin {
+    pub fn with_previous(prev_node: &'static str) -> Self {
+        Self { prev_node }
+    }
+}
+
+impl Plugin for ColormapPlugin {
+    fn build(&self, app: &mut App) {
+        app.init_resource::<ColormapInputImage>()
+            .init_resource::<ColormapOutputImage>()
+            .init_resource::<ColormapMappingImage>();
+
+        let render_app = app.sub_app_mut(RenderApp);
+        render_app
+            .init_resource::<ColormapPipeline>()
+            .add_system_to_stage(RenderStage::Extract, extract_colormap)
+            .add_system_to_stage(RenderStage::Queue, queue_bind_group);
+
+        let mut render_graph = render_app.world.get_resource_mut::<RenderGraph>().unwrap();
+        render_graph.add_node("colormap", ColormapDispatch);
+        render_graph
+            .add_node_edge("colormap", MAIN_PASS_DEPENDENCIES)
+            .unwrap();
+
+        render_graph
+            .add_node_edge(self.prev_node, "colormap")
+            .unwrap();
+    }
+}
+
+#[derive(Default)]
+pub struct ColormapInputImage(pub Handle<Image>);
+#[derive(Default)]
+pub struct ColormapOutputImage(pub Handle<Image>);
+#[derive(Default)]
+pub struct ColormapMappingImage(pub Handle<Image>);
+struct ColormapBindGroup(BindGroup);
+
+struct ColormapSize(Size);
+
+struct ColormapPipeline {
+    pipeline: ComputePipeline,
+    bind_group_layout: BindGroupLayout,
+}
+struct ColormapDispatch;
+
+impl FromWorld for ColormapPipeline {
+    fn from_world(world: &mut World) -> Self {
+        let render_device = world.get_resource::<RenderDevice>().unwrap();
+
+        let shader_source = include_str!("../assets/shaders/colormap.wgsl");
+        let shader = render_device.create_shader_module(&ShaderModuleDescriptor {
+            label: Some("colormap_shader"),
+            source: ShaderSource::Wgsl(shader_source.into()),
+        });
+
+        let texture_bind_group_layout =
+            render_device.create_bind_group_layout(&BindGroupLayoutDescriptor {
+                label: Some("colormap_bind_group_layout"),
+                entries: &[
+                    BindGroupLayoutEntry {
+                        binding: 0,
+                        visibility: ShaderStages::COMPUTE,
+                        ty: BindingType::StorageTexture {
+                            access: StorageTextureAccess::ReadOnly,
+                            format: TextureFormat::R32Float,
+                            view_dimension: TextureViewDimension::D2,
+                        },
+                        count: None,
+                    },
+                    BindGroupLayoutEntry {
+                        binding: 1,
+                        visibility: ShaderStages::COMPUTE,
+                        ty: BindingType::StorageTexture {
+                            access: StorageTextureAccess::WriteOnly,
+                            format: TextureFormat::Rgba8Unorm,
+                            view_dimension: TextureViewDimension::D2,
+                        },
+                        count: None,
+                    },
+                    BindGroupLayoutEntry {
+                        binding: 2,
+                        visibility: ShaderStages::COMPUTE,
+                        ty: BindingType::StorageTexture {
+                            access: StorageTextureAccess::ReadOnly,
+                            format: TextureFormat::Rgba8Unorm,
+                            view_dimension: TextureViewDimension::D1,
+                        },
+                        count: None,
+                    },
+                ],
+            });
+
+        let pipeline_layout = render_device.create_pipeline_layout(&PipelineLayoutDescriptor {
+            label: Some("colormap_pipline_layout"),
+            bind_group_layouts: &[&texture_bind_group_layout],
+            push_constant_ranges: &[],
+        });
+
+        let pipeline = render_device.create_compute_pipeline(&ComputePipelineDescriptor {
+            label: Some("colormap_pipeline"),
+            layout: Some(&pipeline_layout),
+            module: &shader,
+            entry_point: "colormap",
+        });
+
+        ColormapPipeline {
+            pipeline,
+            bind_group_layout: texture_bind_group_layout,
+        }
+    }
+}
+
+impl render_graph::Node for ColormapDispatch {
+    fn update(&mut self, _world: &mut World) {}
+
+    fn run(
+        &self,
+        graph: &mut render_graph::RenderGraphContext,
+        render_context: &mut RenderContext,
+        world: &World,
+    ) -> Result<(), render_graph::NodeRunError> {
+        let pipeline = world.get_resource::<ColormapPipeline>().unwrap();
+        if let Some(texture_bind_group) = world.get_resource::<ColormapBindGroup>() {
+            let size = &world.get_resource::<ColormapSize>().unwrap();
+
+            let mut pass = render_context
+                .command_encoder
+                .begin_compute_pass(&ComputePassDescriptor::default());
+
+            pass.set_pipeline(&pipeline.pipeline);
+            pass.set_bind_group(0, &texture_bind_group.0, &[]);
+            pass.dispatch(
+                (size.0.width / 8.0).ceil() as u32,
+                (size.0.height / 8.0).ceil() as u32,
+                1,
+            );
+        }
+
+        Ok(())
+    }
+}
+
+fn extract_colormap(
+    mut commands: Commands,
+    input: Res<ColormapInputImage>,
+    output: Res<ColormapOutputImage>,
+    mapping: Res<ColormapMappingImage>,
+) {
+    commands.insert_resource(ColormapInputImage(input.0.clone()));
+    commands.insert_resource(ColormapOutputImage(output.0.clone()));
+    commands.insert_resource(ColormapMappingImage(mapping.0.clone()));
+}
+
+fn queue_bind_group(
+    mut commands: Commands,
+    pipeline: Res<ColormapPipeline>,
+    gpu_images: Res<RenderAssets<Image>>,
+    input: Res<ColormapInputImage>,
+    output: Res<ColormapOutputImage>,
+    mapping: Res<ColormapMappingImage>,
+    render_device: Res<RenderDevice>,
+) {
+    if let (Some(input), Some(output), Some(mapping)) = (
+        gpu_images.get(&input.0),
+        gpu_images.get(&output.0),
+        gpu_images.get(&mapping.0),
+    ) {
+        let ix = input.size.width.round() as i32;
+        let iy = input.size.height.round() as i32;
+        let ox = output.size.width.round() as i32;
+        let oy = output.size.height.round() as i32;
+        if (ix == ox) && (iy == oy) {
+            let bind_group = render_device.create_bind_group(&BindGroupDescriptor {
+                label: Some("colormap_bind_group"),
+                layout: &pipeline.bind_group_layout,
+                entries: &[
+                    BindGroupEntry {
+                        binding: 0,
+                        resource: BindingResource::TextureView(&input.texture_view),
+                    },
+                    BindGroupEntry {
+                        binding: 1,
+                        resource: BindingResource::TextureView(&output.texture_view),
+                    },
+                    BindGroupEntry {
+                        binding: 2,
+                        resource: BindingResource::TextureView(&mapping.texture_view),
+                    },
+                ],
+            });
+            commands.insert_resource(ColormapBindGroup(bind_group));
+            commands.insert_resource(ColormapSize(input.size));
+        }
+    }
+}

src/main.rs

@@ -0,0 +1,210 @@
+mod mandelbrot;
+
+use bevy_clicking::{ClickEvent, ClickingPlugin, DoubleclickEvent};
+use mandelbrot::{MandelbrotMaterial, MandelbrotMesh2dBundle, MandelbrotPlugin};
+
+use bevy::{
+    input::mouse::{MouseMotion, MouseWheel},
+    prelude::*,
+    window::WindowResized,
+};
+
+struct Screen {
+    width: f32,
+    height: f32,
+    aspect: f32,
+}
+
+type MousePos = Vec2;
+
+fn main() {
+    App::new()
+        .add_plugins(DefaultPlugins)
+        .add_plugin(MandelbrotPlugin::default())
+        .add_plugin(ClickingPlugin)
+        .insert_resource(ClearColor(Color::rgb(0.0, 0.0, 0.0)))
+        .insert_resource(Screen {
+            width: 1.0,
+            height: 1.0,
+            aspect: 1.0,
+        })
+        .insert_resource(MousePos::new(0.0, 0.0))
+        .add_startup_system(setup)
+        .add_system(bevy::input::system::exit_on_esc_system)
+        .add_system(bevy::input::mouse::mouse_button_input_system)
+        .add_system(fractal_drag)
+        .add_system(fractal_zoom)
+        .add_system(fractal_start)
+        .add_system(cursor_moved)
+        .add_system(change_iters)
+        .add_system(window_size)
+        .add_system(reset_start)
+        .add_system(reset_view)
+        .run();
+}
+
+fn setup(
+    mut commands: Commands,
+    mut meshes: ResMut<Assets<Mesh>>,
+    mut materials: ResMut<Assets<MandelbrotMaterial>>,
+    asset_server: ResMut<AssetServer>,
+) {
+    asset_server.watch_for_changes().unwrap();
+    commands.spawn_bundle(OrthographicCameraBundle::new_2d());
+    commands.spawn_bundle(MandelbrotMesh2dBundle {
+        mesh: meshes.add(Mesh::from(shape::Quad::default())).into(),
+        transform: Transform::default(),
+        material: materials.add(MandelbrotMaterial {
+            center: Vec2::new(-0.4, 0.0),
+            start: Vec2::new(0.0, 0.0),
+            scale: 0.4,
+            aspect: 1.0,
+            iters: 64,
+        }),
+        ..Default::default()
+    });
+}
+
+fn window_size(
+    mut size_event: EventReader<WindowResized>,
+    mut query: Query<(&mut Transform, &Handle<MandelbrotMaterial>)>,
+    mut mat: ResMut<Assets<MandelbrotMaterial>>,
+    mut screen: ResMut<Screen>,
+) {
+    for wse in size_event.iter() {
+        for (mut tr, handle) in query.iter_mut() {
+            tr.scale = Vec3::new(wse.width as f32, wse.height as f32, 1.0);
+            mat.get_mut(handle).unwrap().aspect = wse.width as f32 / wse.height as f32;
+            screen.width = wse.width as f32;
+            screen.height = wse.height as f32;
+            screen.aspect = screen.width / screen.height;
+        }
+    }
+}
+
+fn fractal_drag(
+    mut mouse_event: EventReader<MouseMotion>,
+    mut query: Query<&Handle<MandelbrotMaterial>>,
+    screen: Res<Screen>,
+    lmb: Res<Input<MouseButton>>,
+    mut mats: ResMut<Assets<MandelbrotMaterial>>,
+) {
+    for ev in mouse_event.iter() {
+        let dx = ev.delta.x / screen.height; // No typo, height is reference
+        let dy = ev.delta.y / screen.height;
+        if lmb.pressed(MouseButton::Left) {
+            for handle in query.iter_mut() {
+                let fractal = mats.get_mut(handle).unwrap();
+                fractal.center.x -= dx / fractal.scale;
+                fractal.center.y -= dy / fractal.scale;
+            }
+        }
+    }
+}
+
+fn fractal_start(
+    mut mouse_event: EventReader<MouseMotion>,
+    mut query: Query<&Handle<MandelbrotMaterial>>,
+    screen: Res<Screen>,
+    rmb: Res<Input<MouseButton>>,
+    mut mats: ResMut<Assets<MandelbrotMaterial>>,
+) {
+    if rmb.pressed(MouseButton::Right) {
+        for ev in mouse_event.iter() {
+            for handle in query.iter_mut() {
+                let fractal = mats.get_mut(handle).unwrap();
+                let dx = ev.delta.x / screen.height;
+                let dy = ev.delta.y / screen.height;
+                fractal.start.x -= dx / 4.;
+                fractal.start.y -= dy / 4.;
+                println!("start = {}, {}", fractal.start.x, fractal.start.y);
+            }
+        }
+    }
+}
+
+fn fractal_zoom(
+    mut query: Query<&Handle<MandelbrotMaterial>>,
+    mut scroll_event: EventReader<MouseWheel>,
+    mut mats: ResMut<Assets<MandelbrotMaterial>>,
+    mouse_pos: Res<MousePos>,
+) {
+    for ev in scroll_event.iter() {
+        let amt = ev.y * 0.05;
+        let factor = 1.0 + amt;
+        for handle in query.iter_mut() {
+            let fractal = mats.get_mut(handle).unwrap();
+
+            // Correct center position to zoom towards mouse position
+            fractal.center.x += mouse_pos.x / fractal.scale * amt / 2.0;
+            fractal.center.y -= mouse_pos.y / fractal.scale * amt / 2.0;
+            fractal.scale *= factor;
+        }
+    }
+}
+
+fn cursor_moved(
+    mut cursor_moved: EventReader<CursorMoved>,
+    mut mouse_pos: ResMut<MousePos>,
+    screen: Res<Screen>,
+) {
+    for vm in cursor_moved.iter() {
+        mouse_pos.x = ((vm.position.x / screen.width) - 0.5) * 2.0 * screen.aspect;
+        mouse_pos.y = ((vm.position.y / screen.height) - 0.5) * 2.0;
+    }
+}
+
+fn change_iters(
+    mut scroll_event: EventReader<MouseWheel>,
+    mut mats: ResMut<Assets<MandelbrotMaterial>>,
+    mut query: Query<&Handle<MandelbrotMaterial>>,
+) {
+    for ev in scroll_event.iter() {
+        if (ev.x.abs() < 0.05) {
+            continue;
+        }
+
+        let dir = ev.x.signum() as i32;
+        for handle in query.iter_mut() {
+            let fractal = mats.get_mut(handle).unwrap();
+
+            fractal.iters += dir;
+            fractal.iters = fractal.iters.max(2);
+            println!("Iterations: {}", fractal.iters);
+        }
+    }
+}
+
+fn reset_start(
+    mut cl: EventReader<DoubleclickEvent>,
+    mut query: Query<&Handle<MandelbrotMaterial>>,
+    mut mats: ResMut<Assets<MandelbrotMaterial>>,
+) {
+    for ev in cl.iter() {
+        if ev.button == MouseButton::Right {
+            for handle in query.iter_mut() {
+                let fractal = mats.get_mut(handle).unwrap();
+                fractal.start.x = 0.0;
+                fractal.start.y = 0.0;
+            }
+        }
+    }
+}
+
+fn reset_view(
+    mut cl: EventReader<DoubleclickEvent>,
+    mut query: Query<&Handle<MandelbrotMaterial>>,
+    mut mats: ResMut<Assets<MandelbrotMaterial>>,
+) {
+    for ev in cl.iter() {
+        if ev.button == MouseButton::Left {
+            for handle in query.iter_mut() {
+                let fractal = mats.get_mut(handle).unwrap();
+                fractal.center.x = -0.4;
+                fractal.center.y = 0.0;
+                fractal.scale = 0.4;
+                fractal.iters = 64;
+            }
+        }
+    }
+}

src/mandelbrot.rs

@@ -0,0 +1,118 @@
+use bevy::{
+    ecs::system::{lifetimeless::SRes, SystemParamItem},
+    prelude::*,
+    reflect::TypeUuid,
+    render::{
+        render_asset::{PrepareAssetError, RenderAsset},
+        render_resource::{
+            std140::{AsStd140, Std140},
+            *,
+        },
+        renderer::RenderDevice,
+    },
+    sprite::{Material2d, Material2dPipeline, Material2dPlugin, MaterialMesh2dBundle},
+};
+
+#[derive(Debug, Clone, TypeUuid, Component)]
+#[uuid = "d29793f4-c24d-43f0-97c7-4d417a99188a"]
+pub struct MandelbrotMaterial {
+    pub center: Vec2,
+    pub start: Vec2,
+    pub scale: f32,
+    pub aspect: f32,
+    pub iters: i32,
+}
+
+#[derive(Clone, Default, AsStd140)]
+pub struct MandelbrotFSUniformData {
+    pub center: Vec2,
+    pub start: Vec2,
+    pub scale: f32,
+    pub aspect: f32,
+    pub iters: i32,
+}
+
+#[derive(Debug, Clone)]
+pub struct GpuMandelbrotMaterial {
+    pub fs_buffer: Buffer,
+    pub bind_group: BindGroup,
+}
+
+impl RenderAsset for MandelbrotMaterial {
+    type ExtractedAsset = MandelbrotMaterial;
+    type PreparedAsset = GpuMandelbrotMaterial;
+    type Param = (
+        SRes<RenderDevice>,
+        SRes<Material2dPipeline<MandelbrotMaterial>>,
+    );
+
+    fn extract_asset(&self) -> Self::ExtractedAsset {
+        self.clone()
+    }
+
+    fn prepare_asset(
+        material: Self::ExtractedAsset,
+        (render_device, mandelbrot_pipeline): &mut SystemParamItem<Self::Param>,
+    ) -> Result<Self::PreparedAsset, PrepareAssetError<Self::ExtractedAsset>> {
+        let fs_value = MandelbrotFSUniformData {
+            center: material.center,
+            start: material.start,
+            scale: material.scale,
+            aspect: material.aspect,
+            iters: material.iters,
+        };
+        let fs_value_std140 = fs_value.as_std140();
+
+        let fs_buffer = render_device.create_buffer_with_data(&BufferInitDescriptor {
+            label: Some("mandelbrot_material_uniform_fs_buffer"),
+            usage: BufferUsages::UNIFORM | BufferUsages::COPY_DST,
+            contents: fs_value_std140.as_bytes(),
+        });
+
+        let bind_group = render_device.create_bind_group(&BindGroupDescriptor {
+            entries: &[BindGroupEntry {
+                binding: 0,
+                resource: fs_buffer.as_entire_binding(),
+            }],
+            label: Some("mandelbrot_material_bind_group"),
+            layout: &mandelbrot_pipeline.material2d_layout,
+        });
+
+        Ok(GpuMandelbrotMaterial {
+            fs_buffer,
+            bind_group,
+        })
+    }
+}
+
+impl Material2d for MandelbrotMaterial {
+    fn fragment_shader(asset_server: &AssetServer) -> Option<Handle<Shader>> {
+        Some(asset_server.load("shaders/mandelbrot.wgsl"))
+    }
+
+    #[inline]
+    fn bind_group(render_asset: &<Self as RenderAsset>::PreparedAsset) -> &BindGroup {
+        &render_asset.bind_group
+    }
+
+    fn bind_group_layout(render_device: &RenderDevice) -> BindGroupLayout {
+        render_device.create_bind_group_layout(&BindGroupLayoutDescriptor {
+            entries: &[BindGroupLayoutEntry {
+                binding: 0,
+                visibility: ShaderStages::FRAGMENT,
+                ty: BindingType::Buffer {
+                    ty: BufferBindingType::Uniform,
+                    has_dynamic_offset: false,
+                    min_binding_size: BufferSize::new(
+                        MandelbrotFSUniformData::std140_size_static() as u64,
+                    ),
+                },
+                count: None,
+            }],
+            label: Some("mandelbrot_material_layout"),
+        })
+    }
+}
+
+pub type MandelbrotMesh2dBundle = MaterialMesh2dBundle<MandelbrotMaterial>;
+pub type MandelbrotPlugin = Material2dPlugin<MandelbrotMaterial>;