From 3eb0117e18a24e0bd8bfc0f264dd69c7f814417c Mon Sep 17 00:00:00 2001 From: polygon Date: Mon, 18 Apr 2022 20:24:49 +0200 Subject: [PATCH] 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 | 11 ++ assets/shaders/mandelbrot.wgsl | 38 ++++++ src/colormap.rs | 216 +++++++++++++++++++++++++++++++++ src/main.rs | 210 ++++++++++++++++++++++++++++++++ src/mandelbrot.rs | 118 ++++++++++++++++++ 5 files changed, 593 insertions(+) create mode 100644 Cargo.toml create mode 100644 assets/shaders/mandelbrot.wgsl create mode 100644 src/colormap.rs create mode 100644 src/main.rs create mode 100644 src/mandelbrot.rs diff --git a/Cargo.toml b/Cargo.toml new file mode 100644 index 0000000..f467f75 --- /dev/null +++ b/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" } diff --git a/assets/shaders/mandelbrot.wgsl b/assets/shaders/mandelbrot.wgsl new file mode 100644 index 0000000..c0c7690 --- /dev/null +++ b/assets/shaders/mandelbrot.wgsl @@ -0,0 +1,38 @@ +struct MandelbrotFS { + center: vec2; + start: vec2; + scale: f32; + aspect: f32; + iters: i32; +}; + +[[group(1), binding(0)]] +var fs: MandelbrotFS; + +[[stage(fragment)]] +fn fragment([[location(2)]] uv: vec2) -> [[location(0)]] vec4 { + var z: vec2; + var i: i32; + let iters = fs.iters; + z = fs.start; + var p: vec2; + p = vec2((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(col, col, col, 1.0); + +} diff --git a/src/colormap.rs b/src/colormap.rs new file mode 100644 index 0000000..7b37ec5 --- /dev/null +++ b/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::() + .init_resource::() + .init_resource::(); + + let render_app = app.sub_app_mut(RenderApp); + render_app + .init_resource::() + .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::().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); +#[derive(Default)] +pub struct ColormapOutputImage(pub Handle); +#[derive(Default)] +pub struct ColormapMappingImage(pub Handle); +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::().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::().unwrap(); + if let Some(texture_bind_group) = world.get_resource::() { + let size = &world.get_resource::().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, + output: Res, + mapping: Res, +) { + 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, + gpu_images: Res>, + input: Res, + output: Res, + mapping: Res, + render_device: Res, +) { + 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)); + } + } +} diff --git a/src/main.rs b/src/main.rs new file mode 100644 index 0000000..a860a0a --- /dev/null +++ b/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>, + mut materials: ResMut>, + asset_server: ResMut, +) { + 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, + mut query: Query<(&mut Transform, &Handle)>, + mut mat: ResMut>, + mut screen: ResMut, +) { + 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, + mut query: Query<&Handle>, + screen: Res, + lmb: Res>, + mut mats: ResMut>, +) { + 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, + mut query: Query<&Handle>, + screen: Res, + rmb: Res>, + mut mats: ResMut>, +) { + 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>, + mut scroll_event: EventReader, + mut mats: ResMut>, + mouse_pos: Res, +) { + 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, + mut mouse_pos: ResMut, + screen: Res, +) { + 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, + mut mats: ResMut>, + mut query: Query<&Handle>, +) { + 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, + mut query: Query<&Handle>, + mut mats: ResMut>, +) { + 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, + mut query: Query<&Handle>, + mut mats: ResMut>, +) { + 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; + } + } + } +} diff --git a/src/mandelbrot.rs b/src/mandelbrot.rs new file mode 100644 index 0000000..d8eed06 --- /dev/null +++ b/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, + SRes>, + ); + + fn extract_asset(&self) -> Self::ExtractedAsset { + self.clone() + } + + fn prepare_asset( + material: Self::ExtractedAsset, + (render_device, mandelbrot_pipeline): &mut SystemParamItem, + ) -> Result> { + 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> { + Some(asset_server.load("shaders/mandelbrot.wgsl")) + } + + #[inline] + fn bind_group(render_asset: &::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; +pub type MandelbrotPlugin = Material2dPlugin;