Add anti-alising of jump-flood outline when using MSAA.

Author: komadori

src/flood/compose_output.rs

@@ -1,3 +1,4 @@
+use bevy::shader::ShaderDefVal;
 use bevy::{
     core_pipeline::FullscreenShader,
     platform::collections::HashMap,
@@ -12,7 +13,7 @@ use bevy::{
             SamplerDescriptor, ShaderType, StoreOp,
         },
         renderer::{RenderContext, RenderDevice, RenderQueue},
-        sync_world::{MainEntity, MainEntityHashMap},
+        sync_world::MainEntity,
         texture::CachedTexture,
         view::{ExtractedView, ViewDepthTexture, ViewTarget},
     },
@@ -23,38 +24,61 @@ use wgpu_types::{
     TextureFormat, TextureSampleType,
 };
 
-use crate::{pipeline_key::ViewPipelineKey, uniforms::RenderOutlineInstances};
+use crate::{
+    culling::RenderExtractedOutlineEntities, pipeline_key::ViewPipelineKey,
+    uniforms::RenderOutlineInstances,
+};
 
 use super::{DrawMode, OutlineViewUniform, COMPOSE_OUTPUT_SHADER_HANDLE};
 
 #[derive(Clone, ShaderType)]
 pub(crate) struct ComposeOutputUniform {
+    pub flat_depth: f32,
     pub volume_offset: f32,
     pub volume_colour: Vec4,
 }
 
 #[derive(Resource, Default)]
 pub(crate) struct ComposeOutputUniforms {
     pub buffer: DynamicUniformBuffer<ComposeOutputUniform>,
-    pub offsets: MainEntityHashMap<u32>,
+    pub offsets: HashMap<(Entity, MainEntity), u32>,
 }
 
 pub(crate) fn prepare_compose_output_uniform(
     render_outlines: Res<RenderOutlineInstances>,
+    render_extracted: Res<RenderExtractedOutlineEntities>,
+    views: Query<(Entity, &ExtractedView, &OutlineViewUniform)>,
     mut uniforms: ResMut<ComposeOutputUniforms>,
     render_device: Res<RenderDevice>,
     render_queue: Res<RenderQueue>,
 ) {
     let uniforms = uniforms.as_mut();
     uniforms.buffer.clear();
     uniforms.offsets.clear();
-    for (main_entity, outline) in render_outlines.iter() {
-        if outline.draw_mode == DrawMode::JumpFlood {
+    for (view_entity, view_extracted, view) in views.iter() {
+        let Some(render_view_extracted) = render_extracted
+            .views
+            .get(&view_extracted.retained_view_entity)
+        else {
+            continue;
+        };
+        let model_eye = view_extracted.world_from_view.forward();
+        for (_, main_entity) in render_view_extracted.visible_entities.entities.iter() {
+            let Some(outline) = render_outlines.get(main_entity) else {
+                continue;
+            };
+            if outline.draw_mode != DrawMode::JumpFlood {
+                continue;
+            }
+            let world_plane = outline.instance_data.world_plane_origin
+                + *model_eye * outline.instance_data.world_plane_offset;
+            let clip = view.clip_from_world * world_plane.extend(1.0);
             let offset = uniforms.buffer.push(&ComposeOutputUniform {
+                flat_depth: clip.z / clip.w,
                 volume_offset: outline.instance_data.volume_offset,
                 volume_colour: outline.instance_data.volume_colour,
             });
-            uniforms.offsets.insert(*main_entity, offset);
+            uniforms.offsets.insert((view_entity, *main_entity), offset);
         }
     }
     uniforms.buffer.write_buffer(&render_device, &render_queue);
@@ -101,13 +125,17 @@ impl ComposeOutputPipeline {
         key: ViewPipelineKey,
     ) -> CachedRenderPipelineId {
         *self.pipeline_cache.entry(key).or_insert_with(|| {
+            let mut shader_defs = vec![];
+            if key.msaa().samples() > 1 {
+                shader_defs.push(ShaderDefVal::from("MSAA"));
+            }
             pipeline_cache.queue_render_pipeline(RenderPipelineDescriptor {
                 label: Some("outline_flood_compose_output_pipeline".into()),
                 layout: vec![self.layout.clone()],
                 vertex: fullscreen_shader.to_vertex_state(),
                 fragment: Some(FragmentState {
                     shader: COMPOSE_OUTPUT_SHADER_HANDLE,
-                    shader_defs: vec![],
+                    shader_defs,
                     entry_point: None,
                     targets: vec![Some(ColorTargetState {
                         format: key.target_format(),
@@ -214,7 +242,7 @@ impl<'w> ComposeOutputPass<'w> {
         let dynamic_index = *self
             .compose_output_uniforms
             .offsets
-            .get(&main_entity)
+            .get(&(view_entity, main_entity))
             .unwrap();
 
         let bind_group = render_context.render_device().create_bind_group(

src/flood/compose_output.wgsl

@@ -2,6 +2,7 @@
 #import bevy_mod_outline::common::OutlineViewUniform
 
 struct ComposeOutputUniform {
+    flat_depth: f32,
     volume_offset: f32,
     volume_colour: vec4<f32>,
 }
@@ -19,13 +20,24 @@ struct FragmentOutput {
 @fragment
 fn fragment(in: FullscreenVertexOutput) -> FragmentOutput {
     let tex = textureSample(screen_texture, texture_sampler, in.uv);
+    let threshold = view.scale_physical_from_logical * instance.volume_offset;
+    let dist = distance(in.position.xy, tex.xy);
     var out: FragmentOutput;
-    if distance(in.position.xy, tex.xy) <= view.scale_physical_from_logical * instance.volume_offset {
-        out.colour = instance.volume_colour;
-        out.frag_depth = tex.z;
+#ifdef MSAA
+    let inner = max(threshold - 1.0, 0.0);
+    let coverage = 1.0 - smoothstep(inner, threshold, dist);
+    if coverage <= 0.0 {
+        discard;
     }
-    else {
+    out.colour = vec4<f32>(instance.volume_colour.rgb, instance.volume_colour.a * coverage);
+    out.frag_depth = instance.flat_depth;
+#else
+    if dist <= threshold {
+        out.colour = instance.volume_colour;
+        out.frag_depth = instance.flat_depth;
+    } else {
         discard;
     }
+#endif
     return out;
 }

src/flood/flood_init.rs

@@ -158,7 +158,7 @@ impl<'w> FloodInitPass<'w> {
         }
     }
 
-    pub fn execute(
+    pub fn execute_direct(
         &mut self,
         render_context: &mut RenderContext<'_, '_>,
         range: Range<usize>,
@@ -172,13 +172,42 @@ impl<'w> FloodInitPass<'w> {
                 load: LoadOp::Clear(wgpu_types::Color {
                     r: -1.0,
                     g: -1.0,
-                    b: -1.0,
+                    b: 0.0,
                     a: 0.0,
                 }),
                 store: StoreOp::Store,
             },
         };
 
+        self.run(render_context, range, color_attachment);
+    }
+
+    pub fn execute_coverage(
+        &mut self,
+        render_context: &mut RenderContext<'_, '_>,
+        range: Range<usize>,
+        coverage_msaa: &CachedTexture,
+        resolve_target: &CachedTexture,
+    ) {
+        let color_attachment = RenderPassColorAttachment {
+            view: &coverage_msaa.default_view,
+            depth_slice: None,
+            resolve_target: Some(&resolve_target.default_view),
+            ops: Operations {
+                load: LoadOp::Clear(wgpu_types::Color::TRANSPARENT),
+                store: StoreOp::Store,
+            },
+        };
+
+        self.run(render_context, range, color_attachment);
+    }
+
+    fn run(
+        &mut self,
+        render_context: &mut RenderContext<'_, '_>,
+        range: Range<usize>,
+        color_attachment: RenderPassColorAttachment<'_>,
+    ) {
         let mut init_pass = render_context.begin_tracked_render_pass(RenderPassDescriptor {
             label: Some("outline_flood_init"),
             color_attachments: &[Some(color_attachment)],

src/flood/jump_flood.rs

@@ -65,7 +65,7 @@ pub(crate) fn init_jump_flood_pipeline(
             shader_defs: vec![],
             entry_point: None,
             targets: vec![Some(ColorTargetState {
-                format: TextureFormat::Rgba16Float,
+                format: TextureFormat::Rg16Float,
                 blend: None,
                 write_mask: ColorWrites::ALL,
             })],

src/flood/mod.rs

@@ -25,6 +25,7 @@ use compose_output::{
 use flood_init::{prepare_flood_phases, queue_flood_meshes};
 use jump_flood::init_jump_flood_pipeline;
 use node::{flood_render_pass, FloodOutline};
+use sobel_init::init_sobel_init_pipeline;
 
 use crate::add_dummy_phase_buffer;
 use crate::node::outline_render_pass;
@@ -37,17 +38,27 @@ mod compose_output;
 mod flood_init;
 mod jump_flood;
 mod node;
+mod sobel_init;
 
 const JUMP_FLOOD_SHADER_HANDLE: Handle<Shader> =
     uuid_handle!("66f5981f-0cc2-4e62-8221-cd495062f3ac");
 const COMPOSE_OUTPUT_SHADER_HANDLE: Handle<Shader> =
     uuid_handle!("3c0c1990-4202-48ef-8aa4-bbbb3a334471");
+const SOBEL_INIT_SHADER_HANDLE: Handle<Shader> =
+    uuid_handle!("e011500d-544c-4a0a-85ee-e7de0b1fda3f");
+
+#[derive(Clone)]
+pub(crate) struct FloodCoverageTextures {
+    pub msaa_tex: CachedTexture,
+    pub resolved: CachedTexture,
+}
 
 #[derive(Clone, Component)]
 pub(crate) struct FloodTextures {
     pub flip: bool,
     pub texture_a: CachedTexture,
     pub texture_b: CachedTexture,
+    pub coverage: Option<FloodCoverageTextures>,
 }
 
 impl FloodTextures {
@@ -76,9 +87,9 @@ pub fn prepare_flood_textures(
     mut commands: Commands,
     mut texture_cache: ResMut<TextureCache>,
     render_device: Res<RenderDevice>,
-    cameras: Query<(Entity, &ExtractedCamera)>,
+    cameras: Query<(Entity, &ExtractedCamera, &Msaa)>,
 ) {
-    for (entity, camera) in cameras.iter() {
+    for (entity, camera, msaa) in cameras.iter() {
         let Some(target_size) = camera.physical_target_size else {
             continue;
         };
@@ -95,15 +106,42 @@ pub fn prepare_flood_textures(
             mip_level_count: 1,
             sample_count: 1,
             dimension: TextureDimension::D2,
-            format: TextureFormat::Rgba16Float,
+            format: TextureFormat::Rg16Float,
             usage: TextureUsages::TEXTURE_BINDING | TextureUsages::RENDER_ATTACHMENT,
             view_formats: &[],
         };
 
+        let coverage = if msaa.samples() > 1 {
+            let coverage_descriptor = TextureDescriptor {
+                format: TextureFormat::R8Unorm,
+                ..texture_descriptor.clone()
+            };
+            let msaa_tex = texture_cache.get(
+                &render_device,
+                TextureDescriptor {
+                    label: Some("outline_flood_coverage_msaa"),
+                    sample_count: msaa.samples(),
+                    usage: TextureUsages::RENDER_ATTACHMENT,
+                    ..coverage_descriptor.clone()
+                },
+            );
+            let resolved = texture_cache.get(
+                &render_device,
+                TextureDescriptor {
+                    label: Some("outline_flood_coverage_resolved"),
+                    ..coverage_descriptor
+                },
+            );
+            Some(FloodCoverageTextures { msaa_tex, resolved })
+        } else {
+            None
+        };
+
         commands.entity(entity).insert(FloodTextures {
             flip: false,
             texture_a: texture_cache.get(&render_device, texture_descriptor.clone()),
             texture_b: texture_cache.get(&render_device, texture_descriptor),
+            coverage,
         });
     }
 }
@@ -131,6 +169,13 @@ impl Plugin for FloodPlugin {
             "compose_output.wgsl",
             Shader::from_wgsl
         );
+        load_internal_asset!(
+            app,
+            SOBEL_INIT_SHADER_HANDLE,
+            "sobel_init.wgsl",
+            Shader::from_wgsl
+        );
+
         app.add_plugins(SortedRenderPhasePlugin::<FloodOutline, OutlinePipeline>::new(
             RenderDebugFlags::empty(),
         ))
@@ -165,7 +210,11 @@ impl Plugin for FloodPlugin {
         let render_app = app.sub_app_mut(RenderApp);
         render_app.add_systems(
             RenderStartup,
-            (init_jump_flood_pipeline, init_compose_output_pipeline),
+            (
+                init_sobel_init_pipeline,
+                init_jump_flood_pipeline,
+                init_compose_output_pipeline,
+            ),
         );
 
         let gpu_preprocessing_support = render_app.world().resource::<GpuPreprocessingSupport>();