Initial Commitv0.1.0

4 files changed, 252 insertions, 0 deletions

diff --git a/src/components.rs b/src/components.rs
new file mode 100644
index 0000000..932aac3
--- /dev/null
+++ b/src/components.rs
@@ -0,0 +1,63 @@
+use bevy::{
+    prelude::*,
+    render::{
+        render_asset::RenderAssetUsages,
+        render_resource::{Extent3d, TextureDimension, TextureFormat, TextureUsages}, extract_component::ExtractComponent,
+    },
+};
+
+#[derive(Component, ExtractComponent, Clone)]
+pub struct DitherPostProcessSettings(Handle<Image>);
+
+impl DitherPostProcessSettings {
+    pub fn new(level: u32, mut images: ResMut<Assets<Image>>) -> Self {
+        let power = level + 1;
+        let map_size: u32 = 1 << power;
+        let mut buffer = Vec::<u8>::new();
+
+        for row in 0..map_size {
+            for col in 0..map_size {
+                let a = row ^ col;
+                // Interleave bits of `a` with bits of y coordinate in reverse order
+                let mut result: u64 = 0;
+                let mut bit = 0;
+                let mut mask = power as i32 - 1;
+                loop {
+                    if bit >= 2 * power {
+                        break;
+                    }
+                    result |= (((col >> mask) & 1) << bit) as u64;
+                    bit += 1;
+                    result |= (((a >> mask) & 1) << bit) as u64;
+                    bit += 1;
+                    mask -= 1;
+                }
+                let value = ((result as f32 / map_size.pow(2) as f32) * 255.0) as u8;
+                buffer.push(value);
+            }
+
+        }
+
+        let mut image = Image::new(
+            Extent3d {
+                width: map_size,
+                height: map_size,
+                depth_or_array_layers: 1,
+            },
+            TextureDimension::D2,
+            buffer,
+            TextureFormat::R8Unorm,
+            RenderAssetUsages::RENDER_WORLD,
+        );
+        image.texture_descriptor.usage =
+        TextureUsages::COPY_DST | TextureUsages::STORAGE_BINDING | TextureUsages::TEXTURE_BINDING;
+
+        let handle = images.add(image);
+
+        Self(handle)
+    }
+
+    pub fn handle(&self) -> Handle<Image> {
+        self.0.clone()
+    }
+}
diff --git a/src/lib.rs b/src/lib.rs
new file mode 100644
index 0000000..fdf3fa8
--- /dev/null
+++ b/src/lib.rs
@@ -0,0 +1,43 @@
+use bevy::{prelude::*, render::{RenderApp, render_graph::{RenderGraphApp, ViewNodeRunner}, extract_component::ExtractComponentPlugin}, asset::embedded_asset, core_pipeline::core_3d::graph::{Core3d, Node3d}};
+
+use crate::components::DitherPostProcessSettings;
+
+pub struct DitherPostProcessPlugin;
+
+pub mod components;
+mod resources;
+mod nodes;
+
+impl Plugin for DitherPostProcessPlugin {
+    fn build(&self, app: &mut App) {
+        embedded_asset!(app, "../assets/shaders/dither_post_process.wgsl");
+
+        app.add_plugins((
+            ExtractComponentPlugin::<DitherPostProcessSettings>::default(),
+        ));
+
+        let Ok(render_app) = app.get_sub_app_mut(RenderApp) else {
+            return;
+        };
+
+        render_app.add_render_graph_node::<ViewNodeRunner<nodes::DitherRenderNode>>(
+            Core3d,
+            nodes::DitherRenderLabel,
+        ).add_render_graph_edges(
+            Core3d, 
+            (
+                Node3d::Tonemapping,
+                nodes::DitherRenderLabel,
+                Node3d::EndMainPassPostProcessing,
+            ),
+        );
+    }
+    
+    fn finish(&self, app: &mut App) {
+        let Ok(render_app) = app.get_sub_app_mut(RenderApp) else {
+            return;
+        };
+
+        render_app.init_resource::<resources::DitherPostProcessPipeline>();
+    }
+}
diff --git a/src/nodes.rs b/src/nodes.rs
new file mode 100644
index 0000000..b33b2a0
--- /dev/null
+++ b/src/nodes.rs
@@ -0,0 +1,67 @@
+use bevy::{prelude::*, render::{render_graph::{ViewNode, NodeRunError, RenderGraphContext, RenderLabel}, view::ViewTarget, renderer::RenderContext, render_resource::{PipelineCache, BindGroupEntries, RenderPassDescriptor, RenderPassColorAttachment, Operations}, render_asset::RenderAssets}, ecs::query::QueryItem};
+
+use super::components;
+use super::resources;
+
+#[derive(RenderLabel, Clone, Eq, PartialEq, Hash, Debug)]
+pub struct DitherRenderLabel;
+
+#[derive(Default)]
+pub struct DitherRenderNode;
+
+impl ViewNode for DitherRenderNode {
+    type ViewQuery = (
+        &'static ViewTarget,
+        &'static components::DitherPostProcessSettings,
+    );
+
+    fn run(
+        &self,
+        _graph: &mut RenderGraphContext,
+        render_context: &mut RenderContext,
+        (view_target, dither_post_process_settings): QueryItem<Self::ViewQuery>,
+        world: &World,
+    ) -> Result<(), NodeRunError> {
+        let render_pipeline = world.resource::<resources::DitherPostProcessPipeline>();
+        let pipeline_cache = world.resource::<PipelineCache>();
+        let Some(pipeline) = pipeline_cache.get_render_pipeline(render_pipeline.pipeline_id) else {
+            warn!("Failed to get render pipeline from cache, skipping...");
+            return Ok(());
+        };
+
+        let post_process = view_target.post_process_write();
+
+        let Some(threshold_map) = world.resource::<RenderAssets<Image>>().get(dither_post_process_settings.handle()) else {
+            warn!("Failed to get threshold map, skipping...");
+            return Ok(());
+        };
+
+        let bind_group = render_context.render_device().create_bind_group(
+            "dither_post_process_bind_group",
+            &render_pipeline.layout,
+            &BindGroupEntries::sequential((
+                post_process.source,
+                &render_pipeline.screen_sampler,
+                &threshold_map.texture_view,
+                &render_pipeline.threshold_map_sampler,
+            )),
+        );
+
+        let mut render_pass = render_context.begin_tracked_render_pass(RenderPassDescriptor {
+            label: Some("dither_post_process_render_pass"),
+            color_attachments: &[Some(RenderPassColorAttachment {
+                view: post_process.destination,
+                ops: Operations::default(),
+                resolve_target: None,
+            })],
+            timestamp_writes: None,
+            depth_stencil_attachment: None,
+            occlusion_query_set: None,
+        });
+
+        render_pass.set_render_pipeline(pipeline);
+        render_pass.set_bind_group(0, &bind_group, &[]);
+        render_pass.draw(0..3, 0..1);
+        Ok(())
+    }
+}
diff --git a/src/resources.rs b/src/resources.rs
new file mode 100644
index 0000000..d900f18
--- /dev/null
+++ b/src/resources.rs
@@ -0,0 +1,79 @@
+use bevy::{
+    core_pipeline::fullscreen_vertex_shader::fullscreen_shader_vertex_state,
+    prelude::*,
+    render::{
+        render_resource::{
+            binding_types::{sampler, texture_2d},
+            BindGroupLayout, BindGroupLayoutEntries, CachedRenderPipelineId, ColorTargetState,
+            ColorWrites, FragmentState, MultisampleState, PipelineCache, PrimitiveState,
+            RenderPipelineDescriptor, Sampler, SamplerBindingType, SamplerDescriptor, ShaderStages,
+            TextureFormat, TextureSampleType,
+        },
+        renderer::RenderDevice,
+        texture::BevyDefault,
+    },
+};
+
+#[derive(Resource)]
+pub struct DitherPostProcessPipeline {
+    pub layout: BindGroupLayout,
+    pub screen_sampler: Sampler,
+    pub threshold_map_sampler: Sampler,
+    pub pipeline_id: CachedRenderPipelineId,
+}
+
+impl FromWorld for DitherPostProcessPipeline {
+    fn from_world(world: &mut World) -> Self {
+        let render_device = world.resource::<RenderDevice>();
+
+        let layout = render_device.create_bind_group_layout(
+            "dither_post_process_bind_group_layout",
+            &BindGroupLayoutEntries::sequential(
+                ShaderStages::FRAGMENT,
+                (
+                    texture_2d(TextureSampleType::Float { filterable: true }),
+                    sampler(SamplerBindingType::Filtering),
+                    texture_2d(TextureSampleType::Float { filterable: true }),
+                    sampler(SamplerBindingType::Filtering),
+                ),
+            ),
+        );
+
+        let screen_sampler = render_device.create_sampler(&SamplerDescriptor::default());
+        let threshold_map_sampler = render_device.create_sampler(&SamplerDescriptor::default());
+
+        let shader = world.resource::<AssetServer>().load::<Shader>(
+            "embedded://grex_dither_post_process/../assets/shaders/dither_post_process.wgsl",
+        );
+
+        let pipeline_id =
+            world
+                .resource_mut::<PipelineCache>()
+                .queue_render_pipeline(RenderPipelineDescriptor {
+                    label: Some("dither_post_process_render_pipeline".into()),
+                    layout: vec![layout.clone()],
+                    push_constant_ranges: vec![],
+                    vertex: fullscreen_shader_vertex_state(),
+                    primitive: PrimitiveState::default(),
+                    depth_stencil: None,
+                    multisample: MultisampleState::default(),
+                    fragment: Some(FragmentState {
+                        shader,
+                        shader_defs: vec![],
+                        entry_point: "fragment".into(),
+                        targets: vec![Some(ColorTargetState {
+                            format: TextureFormat::bevy_default(),
+                            blend: None,
+                            write_mask: ColorWrites::ALL,
+                        })],
+                    }),
+                });
+
+        Self {
+            layout,
+            screen_sampler,
+            threshold_map_sampler,
+            pipeline_id,
+        }
+    }
+}