Bevy egui 3d tile palette demo

egui_tile_palette.rs

#![allow(clippy::type_complexity)]
#![allow(clippy::too_many_arguments)]

use bevy::{
    core_pipeline::tonemapping::Tonemapping, prelude::*, render::view::RenderLayers,
    scene::SceneInstance,
};
use bevy_dolly::prelude::*;
use bevy_egui::{egui, EguiContexts, EguiPlugin, EguiUserTextures};
use bevy_inspector_egui::quick::WorldInspectorPlugin;
use bevy_polyline::prelude::*;
use leafwing_input_manager::prelude::*;
use std::collections::{HashMap, VecDeque};

fn main() {
    App::new()
        .add_plugins(DefaultPlugins.set(WindowPlugin {
            primary_window: Some(Window {
                title: "stack_scenes".to_string(),
                ..default()
            }),
            ..default()
        }))
        .add_plugin(EguiPlugin)
        .add_plugin(PolylinePlugin)
        .add_plugin(WorldInspectorPlugin::new())
        .add_plugin(InputManagerPlugin::<InputActions>::default())
        .add_event::<SetTile>()
        .register_type::<Tileset>()
        .register_type::<Tile>()
        .add_startup_system(setup)
        .add_system(Dolly::<MainCamera>::update_active)
        .add_system(handle_input)
        .add_system(load_tiles)
        .add_system(render_thumbnails)
        .add_system(tile_palette)
        .add_system(set_tile)
        .run();
}

fn setup(mut commands: Commands) {
    // load the model
    let scenes = [
        "kenney_hexagon-kit/building_cabin.glb",
        "kenney_hexagon-kit/building_castle.glb",
        "kenney_hexagon-kit/building_dock.glb",
        "kenney_hexagon-kit/building_farm.glb",
        "kenney_hexagon-kit/building_house.glb",
        "kenney_hexagon-kit/building_market.glb",
        "kenney_hexagon-kit/building_mill.glb",
        "kenney_hexagon-kit/building_mine.glb",
        "kenney_hexagon-kit/building_sheep.glb",
        "kenney_hexagon-kit/building_sheep_recolor.glb",
        "kenney_hexagon-kit/building_smelter.glb",
        "kenney_hexagon-kit/building_tower.glb",
        "kenney_hexagon-kit/building_village.glb",
        "kenney_hexagon-kit/building_wall.glb",
        "kenney_hexagon-kit/building_water.glb",
        "kenney_hexagon-kit/dirt.glb",
        "kenney_hexagon-kit/dirt_lumber.glb",
        "kenney_hexagon-kit/grass.glb",
        "kenney_hexagon-kit/grass_forest.glb",
        "kenney_hexagon-kit/grass_hill.glb",
        "kenney_hexagon-kit/path_corner.glb",
        "kenney_hexagon-kit/path_cornerSharp.glb",
        "kenney_hexagon-kit/path_crossing.glb",
        "kenney_hexagon-kit/path_end.glb",
        "kenney_hexagon-kit/path_intersectionA.glb",
        "kenney_hexagon-kit/path_intersectionB.glb",
        "kenney_hexagon-kit/path_intersectionC.glb",
        "kenney_hexagon-kit/path_intersectionD.glb",
        "kenney_hexagon-kit/path_intersectionE.glb",
        "kenney_hexagon-kit/path_intersectionF.glb",
        "kenney_hexagon-kit/path_intersectionG.glb",
        "kenney_hexagon-kit/path_intersectionH.glb",
        "kenney_hexagon-kit/path_start.glb",
        "kenney_hexagon-kit/path_straight.glb",
        "kenney_hexagon-kit/river_corner.glb",
        "kenney_hexagon-kit/river_cornerSharp.glb",
        "kenney_hexagon-kit/river_crossing.glb",
        "kenney_hexagon-kit/river_end.glb",
        "kenney_hexagon-kit/river_intersectionA.glb",
        "kenney_hexagon-kit/river_intersectionB.glb",
        "kenney_hexagon-kit/river_intersectionC.glb",
        "kenney_hexagon-kit/river_intersectionD.glb",
        "kenney_hexagon-kit/river_intersectionE.glb",
        "kenney_hexagon-kit/river_intersectionF.glb",
        "kenney_hexagon-kit/river_intersectionG.glb",
        "kenney_hexagon-kit/river_intersectionH.glb",
        "kenney_hexagon-kit/river_start.glb",
        "kenney_hexagon-kit/river_straight.glb",
        "kenney_hexagon-kit/sand.glb",
        "kenney_hexagon-kit/sand_rocks.glb",
        "kenney_hexagon-kit/stone.glb",
        "kenney_hexagon-kit/stone_hill.glb",
        "kenney_hexagon-kit/stone_mountain.glb",
        "kenney_hexagon-kit/stone_rocks.glb",
        "kenney_hexagon-kit/unit_boat.glb",
        "kenney_hexagon-kit/unit_house.glb",
        "kenney_hexagon-kit/unit_houseLarge.glb",
        "kenney_hexagon-kit/unit_mill.glb",
        "kenney_hexagon-kit/unit_tower.glb",
        "kenney_hexagon-kit/unit_tree.glb",
        "kenney_hexagon-kit/unit_wallTower.glb",
        "kenney_hexagon-kit/water.glb",
        "kenney_hexagon-kit/water_island.glb",
        "kenney_hexagon-kit/water_rocks.glb",
    ];

    let mut tileset = Tileset::new();
    for path in scenes {
        tileset.add_tile(path.into());
    }
    commands.spawn((Name::new("Tileset"), tileset));
    commands.insert_resource(ThumbnailRenderQueue::default());

    // Add the world camera
    commands.spawn((
        MainCamera,
        bevy::render::view::RenderLayers::layer(0),
        Camera3dBundle {
            tonemapping: Tonemapping::None,
            projection: OrthographicProjection {
                near: -100.0,
                far: 100.0,
                scaling_mode: bevy::render::camera::ScalingMode::WindowSize(48.0),
                ..default()
            }
            .into(),
            ..default()
        },
        InputManagerBundle::<InputActions> {
            action_state: ActionState::default(),
            input_map: input_map(),
        },
        Rig::builder()
            .with(Position::new(Vec3::new(0.0, 0.0, 0.0)))
            .with(YawPitch::new().pitch_degrees(-30.0).yaw_degrees(45.0))
            .with(Smooth::new_position(0.3))
            .with(Smooth::new_rotation(0.3))
            .with(Arm::new(Vec3::Z * 5.0))
            .build(),
    ));

    // Add a directional light (the sun)
    commands.spawn((DirectionalLightBundle {
        directional_light: DirectionalLight {
            illuminance: 18000.0,
            ..default()
        },
        transform: Transform::from_rotation(Quat::from_rotation_x(-0.5)),
        ..default()
    },));

    // fun colors
    commands.insert_resource(ClearColor(Color::rgb(1.0, 210.0 / 255.0, 202.0 / 255.0)));

    // Also ambient light
    commands.insert_resource(AmbientLight {
        color: Color::WHITE,
        brightness: 0.15,
    });

    // add a thumbnail rendering camera
    commands.spawn((
        Name::new("thumbnail_render_camera"),
        ThumbnailCamera,
        bevy::render::view::RenderLayers::layer(1),
        Camera3dBundle {
            camera_3d: Camera3d {
                clear_color: bevy::core_pipeline::clear_color::ClearColorConfig::Custom(
                    Color::rgba(0.3, 0.3, 0.3, 1.0),
                ),
                ..default()
            },
            camera: Camera {
                // render before the "main pass" camera
                order: -1,
                is_active: false,
                ..default()
            },
            transform: Transform::from_translation(Vec3::new(3.0, 2.5, 3.0))
                .looking_at(Vec3::new(0.0, 0.25, 0.0), Vec3::Y),
            tonemapping: Tonemapping::None,
            projection: OrthographicProjection {
                near: -100.0,
                far: 100.0,
                scaling_mode: bevy::render::camera::ScalingMode::Fixed {
                    width: 1.3,
                    height: 1.3,
                },
                scale: 1.0,
                ..default()
            }
            .into(),
            ..default()
        },
    ));
}

type TileId = usize;

#[derive(Debug, Reflect, FromReflect, Component)]
struct Tile {
    id: TileId,
    name: String,
    path: std::path::PathBuf,
    model: Option<Handle<Scene>>,
    #[reflect(ignore)]
    egui_texture_id: Option<egui::TextureId>,
}

#[derive(Component, Reflect)]
struct Tileset {
    tiles: HashMap<TileId, Tile>,
    tile_order: Vec<TileId>,
    tile_id_max: TileId,
}

impl Tileset {
    fn new() -> Self {
        Self {
            tiles: HashMap::new(),
            tile_order: Vec::new(),
            tile_id_max: 0,
        }
    }

    fn add_tile(&mut self, path: std::path::PathBuf) {
        let tile = Tile {
            id: self.tile_id_max,
            name: format!("{:?}", path.file_stem().unwrap()),
            path,
            model: None,
            egui_texture_id: None,
        };
        self.tile_order.push(tile.id);
        self.tiles.insert(tile.id, tile);
        self.tile_id_max += 1;
    }
}

#[derive(Resource, Default, Debug)]
struct ThumbnailRenderQueue {
    queue: VecDeque<(Handle<Image>, Handle<Scene>)>,
    scene: Option<Entity>,
}

#[derive(Component)]
struct ThumbnailCamera;

#[derive(Component)]
struct ThumbnailScene;

#[derive(Component)]
struct MainCamera;

#[derive(Component)]
struct SceneRenderLayersPropagated;

struct SetTile(Entity, TileId);

#[derive(Component)]
struct ActiveTile(Entity, TileId);

#[derive(Actionlike, PartialEq, Eq, Clone, Copy, Hash, Debug)]
pub enum InputActions {
    Click,
    Rotate,
    Scale,
    ResetCamera,
    ZeroCamera,
    RedoAabb,
    Run,
}

#[rustfmt::skip]
fn input_map() -> InputMap<InputActions> {
    InputMap::default()
        .insert(MouseButton::Left, InputActions::Click)
        .insert(DualAxis::mouse_motion(), InputActions::Rotate)
        .insert(SingleAxis::mouse_wheel_y(), InputActions::Scale)
        .insert(KeyCode::Z, InputActions::ResetCamera)
        .insert(KeyCode::Key0, InputActions::ZeroCamera)
        .insert(KeyCode::R, InputActions::RedoAabb)
        .insert(KeyCode::Space, InputActions::Run)
        .build()
}

fn handle_input(
    mut camera: Query<(&mut Rig, &mut Projection, &ActionState<InputActions>), With<MainCamera>>,
) {
    let (mut rig, mut projection, actions) = camera.single_mut();
    let camera_yp = rig.driver_mut::<YawPitch>();
    let Projection::Orthographic(projection) = projection.as_mut() else { panic!("wrong scaling mode") };

    if actions.just_pressed(InputActions::ResetCamera) {
        camera_yp.yaw_degrees = 45.0;
        camera_yp.pitch_degrees = -30.0;
        projection.scale = 1.0;
    }

    if actions.just_pressed(InputActions::ZeroCamera) {
        camera_yp.yaw_degrees = 0.0;
        camera_yp.pitch_degrees = 0.0;
        projection.scale = 1.0;
    }

    if actions.pressed(InputActions::Click) {
        let vector = actions.axis_pair(InputActions::Rotate).unwrap().xy();
        camera_yp.rotate_yaw_pitch(-0.1 * vector.x * 15.0, -0.1 * vector.y * 15.0);
    }

    let scale = actions.value(InputActions::Scale);
    if scale != 0.0 {
        projection.scale = (projection.scale * (1.0 - scale * 0.005)).clamp(0.001, 15.0);
    }
}

fn alloc_render_image(width: u32, height: u32) -> Image {
    use bevy::render::render_resource::{
        Extent3d, TextureDescriptor, TextureDimension, TextureFormat, TextureUsages,
    };
    let size = Extent3d {
        width,
        height,
        ..default()
    };

    let mut image = Image {
        texture_descriptor: TextureDescriptor {
            label: None,
            size,
            dimension: TextureDimension::D2,
            format: TextureFormat::Bgra8UnormSrgb,
            mip_level_count: 1,
            sample_count: 1,
            usage: TextureUsages::TEXTURE_BINDING
                | TextureUsages::COPY_DST
                | TextureUsages::RENDER_ATTACHMENT,
            view_formats: &[],
        },
        ..default()
    };

    // fill image.data with zeroes
    image.resize(size);
    image
}

fn load_tiles(
    asset_server: Res<AssetServer>,
    mut tilesets: Query<&mut Tileset, Changed<Tileset>>,
    mut images: ResMut<Assets<Image>>,
    mut render_queue: ResMut<ThumbnailRenderQueue>,
    mut egui_user_textures: ResMut<EguiUserTextures>,
) {
    for mut tileset in &mut tilesets {
        debug!("loading tileset");
        for mut tile in tileset.tiles.values_mut() {
            match tile.model {
                Some(_) => continue,
                None => {
                    tile.model =
                        Some(asset_server.load(format!("{}#Scene0", tile.path.to_string_lossy())))
                }
            }
            match tile.egui_texture_id {
                Some(_) => continue,
                None => {
                    let image = alloc_render_image(48 * 2, 48 * 2);
                    let handle = images.add(image);
                    tile.egui_texture_id = Some(egui_user_textures.add_image(handle.clone()));
                    render_queue
                        .queue
                        .push_back((handle, tile.model.as_ref().unwrap().clone()));
                }
            }
        }
    }
}

fn render_thumbnails(
    mut commands: Commands,
    mut render_queue: ResMut<ThumbnailRenderQueue>,
    mut camera: Query<(&mut Camera, &RenderLayers), With<ThumbnailCamera>>,
    scene_instances: Query<&SceneInstance, With<ThumbnailScene>>,
    scene_manager: Res<SceneSpawner>,
) {
    use bevy::render::camera::RenderTarget;

    let (mut camera, render_layers) = camera
        .get_single_mut()
        .expect("a single ThumbnailCamera to exist");

    // if we're working on an existing scene, see if it's loaded
    if let Some(scene) = render_queue.scene {
        if let Ok(instance) = scene_instances.get(scene) {
            // check if the scene has been loaded
            if !scene_manager.instance_is_ready(**instance) {
                debug!("scene not loaded {:?}", scene);
                return;
            }

            // scene is loaded, update all the child entities to be in the
            // proper render layer
            for entity in scene_manager.iter_instance_entities(**instance) {
                commands.entity(entity).insert(*render_layers);
            }

            // enable the camera, and clear the tag; we'll render the scene to
            // the image, then despawn the scene entity on the next call of
            // this system.
            debug!("render thumbnail {:?}", scene);
            camera.is_active = true;
            commands
                .entity(scene)
                .remove::<ThumbnailScene>()
                .insert(Visibility::Visible);
            return;
        } else {
            debug!("despawn thumbnail {:?}", scene);
            camera.is_active = false;
            commands.entity(scene).despawn_recursive();
            render_queue.scene = None;
        }
    }

    // scene has been loaded, so let's pop the request off the queue
    let Some((image, scene)) = render_queue.queue.pop_front() else { return };

    // update camera to write to the new image
    camera.target = RenderTarget::Image(image);

    // spawn the new model
    let entity = commands
        .spawn((
            ThumbnailScene,
            SceneBundle {
                scene,
                visibility: Visibility::Hidden,
                ..default()
            },
            *render_layers,
        ))
        .id();
    render_queue.scene = Some(entity);
    debug!("spawn thumbnail {:?}", entity);
}

fn tile_palette(
    mut contexts: EguiContexts,
    mut tilesets: Query<(Entity, &mut Tileset)>,
    mut events: EventWriter<SetTile>,
) {
    use egui::*;

    let Ok((entity, mut tileset)) = tilesets.get_single_mut() else { return };

    let context = contexts.ctx_mut();
    let palette_id = Id::new("tile_palette").with(entity);
    let mut swap_tiles = None;

    // we use our own temp variable in egui Memory because something tweaks the
    // memory.interaction.drag_id is getting stomped in some way that prevented
    // us from being able to handle both click and drag & drop.

    egui::Window::new("Tileset")
        .default_width(200.0)
        .vscroll(true)
        .show(context, |ui| {
            ui.columns(4, |cols| {
                let mut column_index = (0..=3).cycle();

                for (index, tile_id) in tileset.tile_order.iter().enumerate() {
                    let Some(tile) = tileset.tiles.get(tile_id)  else { continue } ;
                    let Some(texture_id) = tile.egui_texture_id else { continue };

                    let ui = &mut cols[column_index.next().unwrap()];
                    let button = ImageButton::new(texture_id, [48.0, 48.0])
                        .frame(false)
                        .sense(Sense::click_and_drag());
                    let drag_id = ui.memory_mut(|mem| mem.data.get_temp::<usize>(palette_id));

                    let drag_id = match drag_id {
                        // nothing being dragged
                        None => {
                            let res = ui.add(button);
                            if res.clicked() {
                                debug!("clicked tile {}", tile.name);
                                events.send(SetTile(entity, tile.id));
                            } else if res.drag_delta().length() > 4.0 {
                                // set the temp value for our tile index
                                ui.memory_mut(|mem| mem.data.insert_temp(palette_id, index));
                            }
                            continue;
                        }
                        Some(v) => v,
                    };

                    // dragging this button
                    if drag_id == index {
                        ui.ctx().set_cursor_icon(CursorIcon::Grabbing);

                        let layer_id = LayerId::new(Order::Tooltip, palette_id.with(index));
                        let response = ui.with_layer_id(layer_id, |ui| ui.add(button)).response;
                        if let Some(pos) = ui.ctx().pointer_interact_pos() {
                            let delta = pos - response.rect.center();
                            ui.ctx().translate_layer(layer_id, delta);
                        }

                    // dragging, but not this button
                    } else {
                        let res = ui.add(button);

                        // if we're hovering over this button, and the mouse
                        // button was just released, the drag has ended.  We
                        // need to do it this way because drag_release()
                        // doesn't register on this one.
                        if res.hovered() && ui.input(|i| i.pointer.any_released()) {
                            swap_tiles = Some((drag_id, index));

                            // clear our dragged value
                            ui.memory_mut(|mem| mem.data.remove::<usize>(palette_id));
                        }
                    }
                }
            });
        });

    if let Some((old, new)) = swap_tiles {
        let ele = tileset.tile_order.remove(old);
        tileset.tile_order.insert(new, ele);
    }
}

fn set_tile(
    mut commands: Commands,
    spawned_tiles: Query<Entity, With<ActiveTile>>,
    mut events: EventReader<SetTile>,
    tilesets: Query<&Tileset>,
) {
    let Some(SetTile(entity, tile_id)) = events.iter().last() else { return };
    let Ok(tileset) = tilesets.get(*entity) else { return };

    for tile in &spawned_tiles {
        commands.entity(tile).despawn_recursive();
    }

    let Some(tile) = tileset.tiles.get(tile_id) else { return };
    let Some(scene) = &tile.model else { return };

    commands.spawn((
        ActiveTile(*entity, *tile_id),
        SceneBundle {
            scene: scene.clone(),
            ..default()
        },
    ));

    events.clear();
}