martian-chess/examples/shaders.rs

use bevy::core_pipeline::experimental::taa::TemporalAntiAliasSettings;
use bevy::core_pipeline::prepass::MotionVectorPrepass;
use bevy::core_pipeline::tonemapping::DebandDither;
use bevy::core_pipeline::Skybox;
use bevy::pbr::ExtendedMaterial;
use bevy::pbr::MaterialExtension;
use bevy::pbr::OpaqueRendererMethod;
use bevy::pbr::ScreenSpaceAmbientOcclusionBundle;
use bevy::prelude::*;
use bevy::render::render_resource::*;
use bevy::render::view::ColorGrading;

type MyMat = ExtendedMaterial<StandardMaterial, MatExt>;

fn main() {
    App::new()
        .add_plugins((
            DefaultPlugins.set(ImagePlugin::default_nearest()),
            MaterialPlugin::<MyMat>::default(),
        ))
        .insert_resource(Msaa::Off)
        .add_systems(Startup, setup)
        .add_systems(Update, apply_skybox)
        .add_systems(Update, set_scene)
        .add_systems(Update, rotate)
        .add_systems(
            Update,
            toggle_material
                .run_if(|keys: Res<Input<KeyCode>>| -> bool { keys.just_pressed(KeyCode::Space) }),
        )
        .add_systems(
            Update,
            init_materials
                .run_if(|events: Query<Entity, Added<Transform>>| -> bool { !events.is_empty() }),
        )
        .run();
}

#[derive(Component)]
struct Root;

fn setup(mut commands: Commands, assets: Res<AssetServer>) {
    commands.spawn((SceneBundle { ..default() }, Root));

    commands.spawn(PointLightBundle {
        point_light: PointLight {
            intensity: 10.0,
            ..default()
        },
        transform: Transform::from_xyz(-1.0, 4.5, 4.0).looking_at(Vec3::ZERO, Vec3::Y),
        ..default()
    });

    commands.spawn((Camera3dBundle {
        transform: Transform::from_xyz(-1.0, 1.5, 4.0).looking_at(Vec3::ZERO, Vec3::Y),
        ..default()
    },));
}

fn apply_skybox(
    cameras: Query<Entity, With<Camera>>,
    mut images: ResMut<Assets<Image>>,
    assets: Res<AssetServer>,
    mut commands: Commands,
    mut done: Local<bool>,
    mut handle: Local<Handle<Image>>,
) {
    if !*done {
        info!("Applying skybox...");
        *handle = assets.load("images/skybox.png");
        if let Some(image) = images.get_mut(handle.clone()) {
            info!("Loaded skybox image");
            // NOTE: PNGs do not have any metadata that could indicate they contain a cubemap texture,
            // so they appear as one texture. The following code reconfigures the texture as necessary.
            if image.texture_descriptor.array_layer_count() == 1 {
                image.reinterpret_stacked_2d_as_array(
                    image.texture_descriptor.size.height / image.texture_descriptor.size.width,
                );
                image.texture_view_descriptor = Some(TextureViewDescriptor {
                    dimension: Some(TextureViewDimension::Cube),
                    ..default()
                });

                cameras.iter().for_each(|e| {
                    commands.entity(e).insert(Skybox(handle.clone()));
                });

                *done = true;
            }
        } else {
            *done = false;
        }
    }
}

fn set_scene(
    mut events: Query<&mut Handle<Scene>, (With<Root>, Added<Root>)>,
    assets: Res<AssetServer>,
) {
    events.iter_mut().for_each(|mut scene| {
        *scene = assets.load("models/Martian Chess.glb#Scene0");
    });
}

#[derive(Asset, AsBindGroup, Reflect, Debug, Clone)]
struct MatExt {
    #[uniform(100)]
    cutoff: f32,
}

impl MaterialExtension for MatExt {
    fn fragment_shader() -> ShaderRef {
        "examples/shaders/dissolve.wgsl".into()
    }
}

#[derive(Debug, Component)]
struct Backup<T: Component>(T);

fn init_materials(
    events: Query<(Entity, &Handle<StandardMaterial>), Added<Handle<StandardMaterial>>>,
    standard_materials: Res<Assets<StandardMaterial>>,
    mut materials: ResMut<Assets<MyMat>>,
    mut commands: Commands,
) {
    info!("initializing materials");

    events.iter().for_each(|(entity, std_handle)| {
        // Extension we will add to existing gltf-sourced materials
        let extension = MatExt { cutoff: 1.0 };
        // Base material we will extend for the duration of the dissolve effect
        let mut base = standard_materials
            .get(std_handle)
            .expect("Resolve material data")
            .clone();

        base.opaque_render_method = OpaqueRendererMethod::Auto;

        let ext_handle = materials.add(ExtendedMaterial { base, extension });

        commands.entity(entity).insert(Backup(ext_handle.clone()));
    });
}

fn toggle_material(
    query: Query<Entity, Or<(With<Handle<StandardMaterial>>, With<Handle<MyMat>>)>>,
    std_mat: Query<(&Handle<StandardMaterial>, Option<&Backup<Handle<MyMat>>>)>,
    ext_mat: Query<(&Handle<MyMat>, Option<&Backup<Handle<StandardMaterial>>>)>,
    mut commands: Commands,
) {
    query.iter().for_each(|entity| {
        // Entity has standard material, and had extended material
        // Swap these materials
        if let Ok((std_handle, Some(Backup(ext_handle)))) = std_mat.get(entity) {
            info!("Swapping standard material for extended material");
            commands
                .entity(entity)
                .insert(ext_handle.clone())
                .insert(Backup(std_handle.clone()))
                .remove::<Backup<Handle<MyMat>>>()
                .remove::<Handle<StandardMaterial>>();

            return;
        }

        // In this branch we have the extended material assigned to the object
        if let Ok((ext_handle, Some(Backup(std_handle)))) = ext_mat.get(entity) {
            // Entity has standard material, and had extended material // Swap these materials
            info!("Swapping extended material for standard material");

            commands
                .entity(entity)
                .insert(std_handle.clone())
                .insert(Backup(ext_handle.clone()))
                .remove::<Backup<Handle<StandardMaterial>>>()
                .remove::<Handle<MyMat>>();

            return;
        }

        panic!("What is happening?")
    });
}

fn rotate(
    mut query: Query<&mut Transform, With<Handle<Mesh>>>,
    time: Res<Time>,
    mut materials: ResMut<Assets<MyMat>>,
) {
    query.iter_mut().for_each(|mut t| {
        t.rotate_local_y(time.delta_seconds() / 2.0);
    });
    materials.iter_mut().for_each(|(_id, m)| {
        m.extension.cutoff = time.elapsed_seconds().sin().abs();
    })
}