feat: depth/normal global switching

src/lib.rs

@@ -4,6 +4,7 @@ pub mod manifold;
 pub mod material;
 pub mod plucker;
 pub mod primitive;
+pub mod render;
 
 
 pub struct BevyZeroversePlugin;
@@ -16,6 +17,7 @@ impl Plugin for BevyZeroversePlugin {
             material::ZeroverseMaterialPlugin,
             plucker::PluckerPlugin,
             primitive::ZeroversePrimitivePlugin,
+            render::RenderPlugin,
         ));
     }
 }

src/render/depth.rs

@@ -0,0 +1,90 @@
+use bevy::{
+    prelude::*,
+    asset::load_internal_asset,
+    pbr::{
+        MaterialPipeline,
+        MaterialPipelineKey,
+    },
+    render::{
+        mesh::MeshVertexBufferLayout,
+        render_resource::*,
+    },
+};
+
+pub const DEPTH_SHADER_HANDLE: Handle<Shader> = Handle::weak_from_u128(63456234534534);
+
+#[derive(Component, Debug, Clone, Default, Reflect, Eq, PartialEq)]
+#[reflect(Component, Default)]
+pub struct Depth;
+
+
+#[derive(Debug, Default)]
+pub struct DepthPlugin;
+impl Plugin for DepthPlugin {
+    fn build(&self, app: &mut App) {
+        load_internal_asset!(
+            app,
+            DEPTH_SHADER_HANDLE,
+            "depth.wgsl",
+            Shader::from_wgsl
+        );
+
+        app.register_type::<Depth>();
+
+        app.add_plugins(MaterialPlugin::<DepthMaterial>::default());
+
+        app.add_systems(Startup, setup_global_depth_material);
+        app.add_systems(Update, apply_depth_material);
+    }
+}
+
+fn setup_global_depth_material(
+    mut commands: Commands,
+    mut materials: ResMut<Assets<DepthMaterial>>,
+) {
+    let depth_material = materials.add(DepthMaterial::default());
+    commands.insert_resource(GlobalDepthMaterial(depth_material));
+}
+
+fn apply_depth_material(
+    mut commands: Commands,
+    depths: Query<
+        Entity,
+        (With<Depth>, Without<Handle<DepthMaterial>>),
+    >,
+    mut removed_depths: RemovedComponents<Depth>,
+    global_material: Res<GlobalDepthMaterial>,
+) {
+    for e in removed_depths.read() {
+        if let Some(mut commands) = commands.get_entity(e) {
+            commands.remove::<Handle<DepthMaterial>>();
+        }
+    }
+
+    for e in &depths {
+        commands.entity(e).insert(global_material.0.clone());
+    }
+}
+
+
+#[derive(Default, AsBindGroup, TypePath, Debug, Clone, Asset)]
+pub struct DepthMaterial { }
+
+#[derive(Default, Resource, Debug, Clone)]
+pub struct GlobalDepthMaterial(pub Handle<DepthMaterial>);
+
+impl Material for DepthMaterial {
+    fn fragment_shader() -> ShaderRef {
+        DEPTH_SHADER_HANDLE.into()
+    }
+
+    fn specialize(
+        _pipeline: &MaterialPipeline<Self>,
+        descriptor: &mut RenderPipelineDescriptor,
+        _layout: &MeshVertexBufferLayout,
+        _key: MaterialPipelineKey<Self>,
+    ) -> Result<(), SpecializedMeshPipelineError> {
+        descriptor.primitive.cull_mode = None;
+        Ok(())
+    }
+}

src/render/depth.wgsl

@@ -0,0 +1,17 @@
+#import bevy_pbr::forward_io::VertexOutput
+
+
+fn depth_to_rgb(depth: f32) -> vec3<f32> {
+    let normalized_depth = clamp(depth, 0.0, 1.0);
+
+    let r = smoothstep(0.5, 1.0, normalized_depth);
+    let g = 1.0 - abs(normalized_depth - 0.5) * 2.0;
+    let b = 1.0 - smoothstep(0.0, 0.5, normalized_depth);
+
+    return vec3<f32>(r, g, b);
+}
+
+@fragment
+fn fragment(in: VertexOutput) -> @location(0) vec4<f32> {
+    return vec4<f32>(depth_to_rgb(in.position.w), 1.0);
+}

src/render/mod.rs

@@ -0,0 +1,116 @@
+use bevy::prelude::*;
+
+pub mod depth;
+pub mod normal;
+
+
+#[derive(Default, Debug, Resource, Reflect)]
+#[reflect(Resource)]
+pub enum RenderMode {
+    #[default]
+    Color,
+    Depth,
+    Normal,
+}
+
+
+#[derive(Debug, Default)]
+pub struct RenderPlugin;
+
+impl Plugin for RenderPlugin {
+    fn build(&self, app: &mut App) {
+        app.init_resource::<RenderMode>();
+        app.register_type::<RenderMode>();
+
+        app.add_plugins(depth::DepthPlugin);
+        app.add_plugins(normal::NormalPlugin);
+
+        // TODO: add wireframe pbr disable
+        app.add_systems(PreUpdate, auto_disable_pbr_material::<depth::Depth>);
+        app.add_systems(PreUpdate, auto_disable_pbr_material::<normal::Normal>);
+
+        app.add_systems(PreUpdate, apply_render_modes);
+        app.add_systems(PreUpdate, enable_pbr_material);
+    }
+}
+
+
+#[derive(Component, Default, Debug, Reflect)]
+pub struct DisabledPbrMaterial {
+    pub material: Handle<StandardMaterial>,
+}
+
+#[derive(Component, Default, Debug, Reflect)]
+pub struct EnablePbrMaterial;
+
+
+pub fn auto_disable_pbr_material<T: Component>(
+    mut commands: Commands,
+    mut disabled_materials: Query<
+        (
+            Entity,
+            &Handle<StandardMaterial>,
+        ),
+        (With<T>, Without<DisabledPbrMaterial>),
+    >,
+) {
+    for (
+        entity,
+        disabled_material,
+    ) in disabled_materials.iter_mut() {
+        commands.entity(entity)
+            .insert(DisabledPbrMaterial {
+                material: disabled_material.clone(),
+            })
+            .remove::<EnablePbrMaterial>()
+            .remove::<Handle<StandardMaterial>>();
+    }
+}
+
+fn enable_pbr_material(
+    mut commands: Commands,
+    mut enabled_materials: Query<
+        (
+            Entity,
+            &DisabledPbrMaterial,
+        ),
+        With<EnablePbrMaterial>,
+    >,
+) {
+    for (entity, disabled_material) in enabled_materials.iter_mut() {
+        commands.entity(entity)
+            .insert(disabled_material.material.clone())
+            .remove::<DisabledPbrMaterial>()
+            .remove::<EnablePbrMaterial>();
+    }
+}
+
+
+fn apply_render_modes(
+    mut commands: Commands,
+    render_mode: Res<RenderMode>,
+    meshes: Query<Entity, With<Handle<Mesh>>>,
+) {
+    if render_mode.is_changed() {
+        for entity in meshes.iter() {
+            commands.entity(entity)
+                .remove::<depth::Depth>()
+                .remove::<normal::Normal>();
+
+            match *render_mode {
+                RenderMode::Color => {
+                    commands.entity(entity)
+                        .insert(EnablePbrMaterial);
+                }
+                RenderMode::Depth => {
+                    commands.entity(entity)
+                        .insert(depth::Depth);
+                }
+                RenderMode::Normal => {
+                    commands.entity(entity)
+                        .insert(normal::Normal);
+                }
+            }
+        }
+    }
+}

src/render/normal.rs

@@ -0,0 +1,91 @@
+use bevy::{
+    prelude::*,
+    asset::load_internal_asset,
+    pbr::{
+        MaterialPipeline,
+        MaterialPipelineKey,
+    },
+    render::{
+        mesh::MeshVertexBufferLayout,
+        render_resource::*,
+    },
+};
+
+
+pub const NORMAL_SHADER_HANDLE: Handle<Shader> = Handle::weak_from_u128(234253434561);
+
+#[derive(Component, Debug, Clone, Default, Reflect, Eq, PartialEq)]
+#[reflect(Component, Default)]
+pub struct Normal;
+
+
+#[derive(Debug, Default)]
+pub struct NormalPlugin;
+impl Plugin for NormalPlugin {
+    fn build(&self, app: &mut App) {
+        load_internal_asset!(
+            app,
+            NORMAL_SHADER_HANDLE,
+            "normal.wgsl",
+            Shader::from_wgsl
+        );
+
+        app.register_type::<Normal>();
+
+        app.add_plugins(MaterialPlugin::<NormalMaterial>::default());
+
+        app.add_systems(Startup, setup_global_normal_material);
+        app.add_systems(Update, apply_normal_material);
+    }
+}
+
+fn setup_global_normal_material(
+    mut commands: Commands,
+    mut materials: ResMut<Assets<NormalMaterial>>,
+) {
+    let normal_material = materials.add(NormalMaterial::default());
+    commands.insert_resource(GlobalNormalMaterial(normal_material));
+}
+
+fn apply_normal_material(
+    mut commands: Commands,
+    normals: Query<
+        Entity,
+        (With<Normal>, Without<Handle<NormalMaterial>>),
+    >,
+    mut removed_normals: RemovedComponents<Normal>,
+    global_material: Res<GlobalNormalMaterial>,
+) {
+    for e in removed_normals.read() {
+        if let Some(mut commands) = commands.get_entity(e) {
+            commands.remove::<Handle<NormalMaterial>>();
+        }
+    }
+
+    for e in &normals {
+        commands.entity(e).insert(global_material.0.clone());
+    }
+}
+
+
+#[derive(Default, AsBindGroup, TypePath, Debug, Clone, Asset)]
+pub struct NormalMaterial { }
+
+#[derive(Default, Resource, Debug, Clone)]
+pub struct GlobalNormalMaterial(pub Handle<NormalMaterial>);
+
+impl Material for NormalMaterial {
+    fn fragment_shader() -> ShaderRef {
+        NORMAL_SHADER_HANDLE.into()
+    }
+
+    fn specialize(
+        _pipeline: &MaterialPipeline<Self>,
+        descriptor: &mut RenderPipelineDescriptor,
+        _layout: &MeshVertexBufferLayout,
+        _key: MaterialPipelineKey<Self>,
+    ) -> Result<(), SpecializedMeshPipelineError> {
+        descriptor.primitive.cull_mode = None;
+        Ok(())
+    }
+}

src/render/normal.wgsl

@@ -0,0 +1,8 @@
+#import bevy_pbr::forward_io::VertexOutput
+
+
+@fragment
+fn fragment(in: VertexOutput) -> @location(0) vec4<f32> {
+    let normal = (in.world_normal * 0.5) + vec3<f32>(0.5, 0.5, 0.5);
+    return vec4<f32>(normal, 1.0);
+}