Implement hitbox viewer (#81)

Reviewed-on: #81
Co-authored-by: Álex Román Núñez <eirexe@eirteam.moe>

src/hitbox_viewer/mod.rs

@@ -0,0 +1,694 @@
+//! Standalone, read-only attack-script previewer used by `hitbox-editor`.
+
+use bevy::{gltf::Gltf, prelude::*, world_serialization::WorldAssetRoot};
+use bevy_egui::{EguiContexts, EguiPrimaryContextPass, egui};
+use clap::{Parser, ValueEnum};
+
+use crate::{
+    GameAssetPlugin,
+    fighter::{
+        animation::FighterAnimationFrame,
+        attack::{AttackHitbox, AttackKind},
+        baked_animation::{
+            BakedAnimationAssetProcessingPlugin, BakedAnimationSamplingSet,
+            BakedAnimationVisualPlugin, BakedFighterAnimations, FighterBoneIndexMap,
+            FighterBoneMatrices, FixedMat4,
+        },
+        manifest::{FighterHurtbox, FighterManifest},
+        visual::FighterAnimations,
+    },
+    filled_gizmos::{FilledGizmos, FilledGizmosPlugin},
+    math::{
+        int::{FGi32, FGi32Ext},
+        vec3::FGVec3,
+    },
+    scripting::{FighterAttackScript, importer::FighterScriptImportPlugin},
+};
+
+#[derive(Clone, Copy, Debug, ValueEnum)]
+pub enum ViewerAttack {
+    NeutralAir,
+    ForwardAir,
+    BackAir,
+    DownAir,
+    UpAir,
+    Jab,
+    UpTilt,
+    DownTilt,
+    ForwardTilt,
+    UpSmash,
+    DownSmash,
+    ForwardSmash,
+}
+
+impl From<ViewerAttack> for AttackKind {
+    fn from(value: ViewerAttack) -> Self {
+        match value {
+            ViewerAttack::NeutralAir => Self::NeutralAir,
+            ViewerAttack::ForwardAir => Self::ForwardAir,
+            ViewerAttack::BackAir => Self::BackAir,
+            ViewerAttack::DownAir => Self::DownAir,
+            ViewerAttack::UpAir => Self::UpAir,
+            ViewerAttack::Jab => Self::Jab,
+            ViewerAttack::UpTilt => Self::UpTilt,
+            ViewerAttack::DownTilt => Self::DownTilt,
+            ViewerAttack::ForwardTilt => Self::ForwardTilt,
+            ViewerAttack::UpSmash => Self::UpSmash,
+            ViewerAttack::DownSmash => Self::DownSmash,
+            ViewerAttack::ForwardSmash => Self::ForwardSmash,
+        }
+    }
+}
+
+#[derive(Parser, Resource, Debug)]
+#[command(version, about = "Live preview for Shinespark attack hitboxes")]
+pub struct HitboxViewerArgs {
+    /// Fighter manifest path, relative to the assets directory.
+    #[arg(long, default_value = "fighters/test/test.fighter.ron")]
+    pub manifest: String,
+    /// Attack script and matching animation to preview.
+    #[arg(long, value_enum, default_value_t = ViewerAttack::Jab)]
+    pub attack: ViewerAttack,
+    /// Start with fighter hurtboxes hidden.
+    #[arg(long)]
+    pub hide_hurtboxes: bool,
+}
+
+#[derive(Resource)]
+struct ViewerAssets {
+    manifest: Handle<FighterManifest>,
+    model: Option<Handle<Gltf>>,
+    baked: Option<Handle<BakedFighterAnimations>>,
+    script: Option<Handle<FighterAttackScript>>,
+    fighter: Option<Entity>,
+    error: Option<String>,
+}
+
+#[derive(Resource)]
+struct ViewerState {
+    attack: AttackKind,
+    frame: u32,
+    playing: bool,
+    show_hurtboxes: bool,
+    elapsed: f32,
+}
+
+pub struct HitboxViewerPlugin;
+
+impl Plugin for HitboxViewerPlugin {
+    fn build(&self, app: &mut App) {
+        let args = app.world().resource::<HitboxViewerArgs>();
+        let attack = args.attack.into();
+        let show_hurtboxes = !args.hide_hurtboxes;
+        app.add_plugins((
+            GameAssetPlugin,
+            FighterScriptImportPlugin,
+            BakedAnimationAssetProcessingPlugin,
+            BakedAnimationVisualPlugin,
+            FilledGizmosPlugin,
+        ))
+        .insert_resource(ViewerState {
+            attack,
+            frame: 0,
+            playing: true,
+            show_hurtboxes,
+            elapsed: 0.0,
+        })
+        .add_systems(Startup, setup)
+        .add_systems(
+            FixedUpdate,
+            (
+                request_fighter_dependencies,
+                spawn_preview_fighter,
+                handle_keyboard,
+                advance_preview,
+                sync_preview_frame,
+                sample_preview_pose.in_set(BakedAnimationSamplingSet),
+                draw_preview,
+            )
+                .chain(),
+        )
+        .add_systems(EguiPrimaryContextPass, draw_ui);
+    }
+}
+
+fn setup(mut commands: Commands, asset_server: Res<AssetServer>, args: Res<HitboxViewerArgs>) {
+    commands.insert_resource(ViewerAssets {
+        manifest: asset_server.load(args.manifest.clone()),
+        model: None,
+        baked: None,
+        script: None,
+        fighter: None,
+        error: None,
+    });
+    commands.spawn((
+        Name::new("Hitbox viewer camera"),
+        Camera3d::default(),
+        Transform::from_xyz(0.0, 9.0, 38.0).looking_at(Vec3::new(0.0, 8.0, 0.0), Vec3::Y),
+    ));
+    commands.spawn((
+        Name::new("Hitbox viewer light"),
+        DirectionalLight {
+            illuminance: 8_000.0,
+            shadow_maps_enabled: true,
+            ..default()
+        },
+        Transform::from_rotation(Quat::from_euler(EulerRot::XYZ, -0.7, -0.5, 0.0)),
+    ));
+    commands.insert_resource(GlobalAmbientLight {
+        color: Color::WHITE,
+        brightness: 250.0,
+        affects_lightmapped_meshes: true,
+    });
+}
+
+fn request_fighter_dependencies(
+    asset_server: Res<AssetServer>,
+    manifests: Res<Assets<FighterManifest>>,
+    mut viewer: ResMut<ViewerAssets>,
+    state: Res<ViewerState>,
+) {
+    if viewer.model.is_some() {
+        return;
+    }
+    let Some(manifest) = manifests.get(&viewer.manifest) else {
+        return;
+    };
+    let Some(script_path) = manifest.attack_scripts.get(&state.attack) else {
+        viewer.error = Some(format!(
+            "{:?} has no {:?} attack script",
+            manifest.id, state.attack
+        ));
+        return;
+    };
+    viewer.model = Some(asset_server.load(manifest.model_path.clone()));
+    viewer.baked = Some(asset_server.load(manifest.baked_animation_path.clone()));
+    viewer.script = Some(asset_server.load(script_path.clone()));
+    viewer.error = None;
+}
+
+fn spawn_preview_fighter(
+    mut commands: Commands,
+    gltfs: Res<Assets<Gltf>>,
+    baked_assets: Res<Assets<BakedFighterAnimations>>,
+    scripts: Res<Assets<FighterAttackScript>>,
+    manifests: Res<Assets<FighterManifest>>,
+    mut viewer: ResMut<ViewerAssets>,
+    state: Res<ViewerState>,
+) {
+    if viewer.fighter.is_some() {
+        return;
+    }
+    let (Some(model_handle), Some(baked_handle), Some(script_handle)) =
+        (&viewer.model, &viewer.baked, &viewer.script)
+    else {
+        return;
+    };
+    let (Some(gltf), Some(_), Some(_), Some(manifest)) = (
+        gltfs.get(model_handle),
+        baked_assets.get(baked_handle),
+        scripts.get(script_handle),
+        manifests.get(&viewer.manifest),
+    ) else {
+        return;
+    };
+    let Some(scene) = gltf.scenes.first().cloned() else {
+        viewer.error = Some("fighter model contains no scene".into());
+        return;
+    };
+    let entity = commands
+        .spawn((
+            Name::new("Hitbox preview fighter"),
+            FighterAnimations {
+                baked: baked_handle.clone(),
+            },
+            FighterAnimationFrame {
+                kind: state.attack.get_animation(),
+                frame: state.frame,
+                repeat: false,
+            },
+            FighterBoneMatrices::default(),
+            FighterBoneIndexMap::default(),
+            WorldAssetRoot(scene),
+            Transform::from_scale(Vec3::splat(manifest.model_scale.to_num()))
+                .with_rotation(Quat::from_rotation_y(std::f32::consts::FRAC_PI_2)),
+        ))
+        .id();
+    viewer.fighter = Some(entity);
+}
+
+fn timeline_end(
+    state: &ViewerState,
+    viewer: &ViewerAssets,
+    scripts: &Assets<FighterAttackScript>,
+    baked_assets: &Assets<BakedFighterAnimations>,
+) -> u32 {
+    let animation_end = viewer
+        .baked
+        .as_ref()
+        .and_then(|handle| baked_assets.get(handle))
+        .and_then(|baked| baked.frame_count(state.attack.get_animation()))
+        .unwrap_or(1);
+    let hitbox_end = viewer
+        .script
+        .as_ref()
+        .and_then(|handle| scripts.get(handle))
+        .and_then(|script| script.hitboxes.iter().map(|hitbox| hitbox.end_frame).max())
+        .unwrap_or(1);
+    animation_end.max(hitbox_end).max(1)
+}
+
+fn handle_keyboard(
+    keys: Res<ButtonInput<KeyCode>>,
+    mut state: ResMut<ViewerState>,
+    viewer: Res<ViewerAssets>,
+    scripts: Res<Assets<FighterAttackScript>>,
+    baked_assets: Res<Assets<BakedFighterAnimations>>,
+) {
+    let end = timeline_end(&state, &viewer, &scripts, &baked_assets);
+    if keys.just_pressed(KeyCode::Space) {
+        state.playing = !state.playing;
+    }
+    if keys.just_pressed(KeyCode::ArrowLeft) {
+        state.playing = false;
+        state.frame = state.frame.saturating_sub(1);
+    }
+    if keys.just_pressed(KeyCode::ArrowRight) {
+        state.playing = false;
+        state.frame = (state.frame + 1).min(end.saturating_sub(1));
+    }
+    if keys.just_pressed(KeyCode::Home) {
+        state.playing = false;
+        state.frame = 0;
+    }
+}
+
+fn advance_preview(
+    time: Res<Time>,
+    mut state: ResMut<ViewerState>,
+    viewer: Res<ViewerAssets>,
+    scripts: Res<Assets<FighterAttackScript>>,
+    baked_assets: Res<Assets<BakedFighterAnimations>>,
+) {
+    let end = timeline_end(&state, &viewer, &scripts, &baked_assets);
+    state.frame = state.frame.min(end.saturating_sub(1));
+    if !state.playing {
+        state.elapsed = 0.0;
+        return;
+    }
+    state.elapsed += time.delta_secs();
+    while state.elapsed >= 1.0 / 60.0 {
+        state.elapsed -= 1.0 / 60.0;
+        state.frame = (state.frame + 1) % end;
+    }
+}
+
+fn sync_preview_frame(state: Res<ViewerState>, mut frames: Query<&mut FighterAnimationFrame>) {
+    for mut frame in &mut frames {
+        frame.kind = state.attack.get_animation();
+        frame.frame = state.frame;
+        frame.repeat = false;
+    }
+}
+
+fn sample_preview_pose(
+    viewer: Res<ViewerAssets>,
+    manifests: Res<Assets<FighterManifest>>,
+    baked_assets: Res<Assets<BakedFighterAnimations>>,
+    mut preview: Query<(
+        &FighterAnimationFrame,
+        &FighterAnimations,
+        &mut FighterBoneMatrices,
+        &mut FighterBoneIndexMap,
+    )>,
+) {
+    let Some(manifest) = manifests.get(&viewer.manifest) else {
+        return;
+    };
+    for (frame, animations, mut matrices, mut bone_indices) in &mut preview {
+        let Some(baked) = baked_assets.get(&animations.baked) else {
+            continue;
+        };
+        if matrices.sample_baked_pose(baked, frame, manifest.model_scale)
+            && bone_indices.0.is_empty()
+        {
+            bone_indices.0.extend(
+                baked
+                    .bone_names
+                    .iter()
+                    .enumerate()
+                    .map(|(index, name)| (name.clone(), index)),
+            );
+        }
+    }
+}
+
+fn draw_preview(
+    viewer: Res<ViewerAssets>,
+    state: Res<ViewerState>,
+    scripts: Res<Assets<FighterAttackScript>>,
+    baked_assets: Res<Assets<BakedFighterAnimations>>,
+    manifests: Res<Assets<FighterManifest>>,
+    mut gizmos: Gizmos,
+    mut filled: FilledGizmos,
+) {
+    let (Some(script_handle), Some(baked_handle)) = (&viewer.script, &viewer.baked) else {
+        return;
+    };
+    let (Some(script), Some(baked), Some(manifest)) = (
+        scripts.get(script_handle),
+        baked_assets.get(baked_handle),
+        manifests.get(&viewer.manifest),
+    ) else {
+        return;
+    };
+    let animation = FighterAnimationFrame {
+        kind: state.attack.get_animation(),
+        frame: state.frame,
+        repeat: false,
+    };
+    let previous_animation = FighterAnimationFrame {
+        frame: state.frame.saturating_sub(1),
+        ..animation
+    };
+    let fighter_transform = FixedMat4::IDENTITY.rotate_y(1);
+
+    if state.show_hurtboxes {
+        for hurtbox in &manifest.hurtboxes {
+            draw_hurtbox(
+                hurtbox,
+                baked,
+                &animation,
+                manifest.model_scale,
+                fighter_transform,
+                &mut filled,
+            );
+        }
+    }
+
+    for hitbox in script
+        .hitboxes
+        .iter()
+        .filter(|hitbox| hitbox.start_frame <= state.frame && state.frame < hitbox.end_frame)
+    {
+        draw_hitbox(
+            hitbox,
+            baked,
+            &animation,
+            &previous_animation,
+            manifest.model_scale,
+            fighter_transform,
+            &mut gizmos,
+            &mut filled,
+        );
+    }
+}
+
+fn draw_hitbox(
+    hitbox: &AttackHitbox,
+    baked: &BakedFighterAnimations,
+    animation: &FighterAnimationFrame,
+    previous_animation: &FighterAnimationFrame,
+    model_scale: FGi32,
+    fighter_transform: FixedMat4,
+    gizmos: &mut Gizmos,
+    filled: &mut FilledGizmos,
+) {
+    let Some(current_bone) =
+        baked.sample_bone_model_transform(animation, &hitbox.bone, model_scale)
+    else {
+        return;
+    };
+    let current = fighter_transform
+        .mul(current_bone)
+        .transform_point(hitbox.offset)
+        .to_vec3();
+    let previous = if animation.frame == hitbox.start_frame {
+        current
+    } else {
+        baked
+            .sample_bone_model_transform(previous_animation, &hitbox.bone, model_scale)
+            .map(|bone| {
+                fighter_transform
+                    .mul(bone)
+                    .transform_point(hitbox.offset)
+                    .to_vec3()
+            })
+            .unwrap_or(current)
+    };
+    let center = (previous + current) * 0.5;
+    let length = previous.distance(current);
+    let direction = (current - previous).normalize_or(Vec3::Y);
+    let mut transform = Transform::from_translation(center);
+    transform.align(
+        Dir3::Y,
+        direction,
+        Dir3::X,
+        direction.any_orthonormal_vector(),
+    );
+    let radius = hitbox.radius.to_num();
+    filled.capsule(
+        transform.to_isometry(),
+        radius,
+        length * 0.5,
+        Color::srgba(1.0, 0.1, 0.05, 0.32),
+    );
+    gizmos.primitive_3d(
+        &Capsule3d::new(radius, length),
+        transform.to_isometry(),
+        Color::srgb(1.0, 0.15, 0.1),
+    );
+}
+
+fn draw_hurtbox(
+    hurtbox: &FighterHurtbox,
+    baked: &BakedFighterAnimations,
+    animation: &FighterAnimationFrame,
+    model_scale: FGi32,
+    fighter_transform: FixedMat4,
+    filled: &mut FilledGizmos,
+) {
+    let Some(bone) = baked.sample_bone_model_transform(animation, &hurtbox.bone, model_scale)
+    else {
+        return;
+    };
+    let local = FixedMat4::IDENTITY
+        .with_rotation(FGVec3::new(
+            hurtbox.rotation.x.to_radians(),
+            hurtbox.rotation.y.to_radians(),
+            hurtbox.rotation.z.to_radians(),
+        ))
+        .with_translation(hurtbox.offset);
+    filled.affine_capsule(
+        fighter_transform.mul(bone).mul(local).to_mat4(),
+        hurtbox.radius.to_num(),
+        hurtbox.half_length.to_num(),
+        Color::srgba(0.15, 0.55, 1.0, 0.22),
+    );
+}
+
+fn draw_ui(
+    mut contexts: EguiContexts,
+    mut state: ResMut<ViewerState>,
+    viewer: Res<ViewerAssets>,
+    scripts: Res<Assets<FighterAttackScript>>,
+    baked_assets: Res<Assets<BakedFighterAnimations>>,
+) -> Result {
+    let end = timeline_end(&state, &viewer, &scripts, &baked_assets);
+    let script = viewer
+        .script
+        .as_ref()
+        .and_then(|handle| scripts.get(handle));
+    egui::Window::new("Timeline")
+        .anchor(egui::Align2::CENTER_BOTTOM, egui::vec2(0.0, -10.0))
+        .resizable(true)
+        .default_size(egui::vec2(780.0, 190.0))
+        .show(contexts.ctx_mut()?, |ui| {
+            ui.horizontal(|ui| {
+                if ui
+                    .button(if state.playing { "Pause" } else { "Play" })
+                    .clicked()
+                {
+                    state.playing = !state.playing;
+                }
+                if ui.button("<").clicked() {
+                    state.playing = false;
+                    state.frame = state.frame.saturating_sub(1);
+                }
+                if ui.button(">").clicked() {
+                    state.playing = false;
+                    state.frame = (state.frame + 1).min(end.saturating_sub(1));
+                }
+                ui.label(format!("Frame {} / {}", state.frame, end.saturating_sub(1)));
+                ui.checkbox(&mut state.show_hurtboxes, "Hurtboxes");
+                ui.separator();
+                ui.label(format!("{:?}", state.attack));
+            });
+            if let Some(error) = &viewer.error {
+                ui.colored_label(egui::Color32::RED, error);
+            }
+            let Some(script) = script else {
+                ui.spinner();
+                ui.label("Waiting for the processed attack script…");
+                return;
+            };
+            draw_summary_timeline(ui, &mut state, script, end);
+            draw_timeline(ui, &mut state, script, end);
+        });
+    Ok(())
+}
+
+fn draw_summary_timeline(
+    ui: &mut egui::Ui,
+    state: &mut ViewerState,
+    script: &FighterAttackScript,
+    end: u32,
+) {
+    let row_height = 20.0;
+    let label_width = 55.0;
+    let height = (row_height + 25.0f32).min(ui.available_height());
+    let (response, painter) = ui.allocate_painter(
+        egui::vec2(ui.available_width(), height),
+        egui::Sense::click_and_drag(),
+    );
+    let rect = response.rect;
+    let timeline = egui::Rect::from_min_max(
+        egui::pos2(rect.left() + label_width, rect.top()),
+        rect.right_bottom(),
+    );
+    if (response.clicked() || response.dragged())
+        && let Some(pointer) = response.interact_pointer_pos()
+        && timeline.contains(pointer)
+    {
+        let fraction = ((pointer.x - timeline.left()) / timeline.width()).clamp(0.0, 1.0);
+        state.frame = ((end.saturating_sub(1)) as f32 * fraction).round() as u32;
+        state.playing = false;
+    }
+    for frame in 0..end {
+        if frame % 5 != 0 {
+            continue;
+        }
+        let x = frame_x(timeline, frame, end);
+        painter.line_segment(
+            [egui::pos2(x, rect.top()), egui::pos2(x, rect.bottom())],
+            egui::Stroke::new(1.0_f32, egui::Color32::from_gray(55)),
+        );
+        painter.text(
+            egui::pos2(x + 2.0, rect.top()),
+            egui::Align2::LEFT_TOP,
+            frame.to_string(),
+            egui::FontId::monospace(10.0),
+            egui::Color32::GRAY,
+        );
+    }
+    for frame in 0..end {
+        let top = rect.top() + 20.0 as f32;
+        painter.text(
+            egui::pos2(rect.left(), top + 2.0),
+            egui::Align2::LEFT_TOP,
+            format!("HB {}", 2),
+            egui::FontId::monospace(12.0),
+            egui::Color32::LIGHT_GRAY,
+        );
+        let active_hitbox = script
+            .hitboxes
+            .iter()
+            .filter(|hitbox| hitbox.start_frame <= frame && hitbox.end_frame > frame)
+            .max_by(|a, b| a.id.cmp(&b.id));
+        let bar = egui::Rect::from_min_max(
+            egui::pos2(frame_x(timeline, frame, end), top + 2.0),
+            egui::pos2(frame_x(timeline, frame + 1, end), top + row_height - 2.0),
+        );
+
+        if let Some(_active_hitbox) = active_hitbox {
+            painter.rect_filled(bar, 2.0, egui::Color32::from_rgb(245, 75, 55));
+        } else {
+            painter.rect_stroke(bar, 2.0, egui::Stroke::new(2.0f32, egui::Color32::from_rgb(75, 75, 75)), egui::StrokeKind::Inside);
+        }
+    }
+}
+
+fn draw_timeline(
+    ui: &mut egui::Ui,
+    state: &mut ViewerState,
+    script: &FighterAttackScript,
+    end: u32,
+) {
+    let row_height = 20.0;
+    let label_width = 55.0;
+    let height =
+        (script.hitboxes.len().max(1) as f32 * row_height + 25.0).min(ui.available_height());
+    let (response, painter) = ui.allocate_painter(
+        egui::vec2(ui.available_width(), height),
+        egui::Sense::click_and_drag(),
+    );
+    let rect = response.rect;
+    let timeline = egui::Rect::from_min_max(
+        egui::pos2(rect.left() + label_width, rect.top()),
+        rect.right_bottom(),
+    );
+    if (response.clicked() || response.dragged())
+        && let Some(pointer) = response.interact_pointer_pos()
+        && timeline.contains(pointer)
+    {
+        let fraction = ((pointer.x - timeline.left()) / timeline.width()).clamp(0.0, 1.0);
+        state.frame = ((end.saturating_sub(1)) as f32 * fraction).round() as u32;
+        state.playing = false;
+    }
+    for frame in 0..end {
+        if frame % 5 != 0 {
+            continue;
+        }
+        let x = frame_x(timeline, frame, end);
+        painter.line_segment(
+            [egui::pos2(x, rect.top()), egui::pos2(x, rect.bottom())],
+            egui::Stroke::new(1.0_f32, egui::Color32::from_gray(55)),
+        );
+        painter.text(
+            egui::pos2(x + 2.0, rect.top()),
+            egui::Align2::LEFT_TOP,
+            frame.to_string(),
+            egui::FontId::monospace(10.0),
+            egui::Color32::GRAY,
+        );
+    }
+    for (row, hitbox) in script.hitboxes.iter().enumerate() {
+        let top = rect.top() + 20.0 + row as f32 * row_height;
+        painter.text(
+            egui::pos2(rect.left(), top + 2.0),
+            egui::Align2::LEFT_TOP,
+            format!("HB {}", hitbox.id),
+            egui::FontId::monospace(12.0),
+            egui::Color32::LIGHT_GRAY,
+        );
+        let active = hitbox.start_frame <= state.frame && state.frame < hitbox.end_frame;
+        let bar = egui::Rect::from_min_max(
+            egui::pos2(frame_x(timeline, hitbox.start_frame, end), top + 2.0),
+            egui::pos2(
+                frame_x(timeline, hitbox.end_frame, end),
+                top + row_height - 2.0,
+            ),
+        );
+        painter.rect_filled(
+            bar,
+            2.0,
+            if active {
+                egui::Color32::from_rgb(245, 75, 55)
+            } else {
+                egui::Color32::from_rgb(125, 48, 42)
+            },
+        );
+    }
+    let playhead = frame_x(timeline, state.frame, end);
+    painter.line_segment(
+        [
+            egui::pos2(playhead, rect.top()),
+            egui::pos2(playhead, rect.bottom()),
+        ],
+        egui::Stroke::new(2.0_f32, egui::Color32::WHITE),
+    );
+}
+
+fn frame_x(rect: egui::Rect, frame: u32, end: u32) -> f32 {
+    rect.left() + rect.width() * frame.min(end) as f32 / end.max(1) as f32
+}