shinespark/src/scripting/move_compiler.rs

use std::{cell::RefCell, collections::HashMap, rc::Rc};

use rhai::EvalAltResult;
use serde::{Deserialize, Serialize};

use crate::{
    fighter::attack::{self, AttackAngle, FighterDamage, Knockback, KnockbackType},
    math::{int::FGi32, vec::FGVec2, vec3::FGVec3},
    scripting::FighterAttackScript,
};

#[derive(Debug, Default)]
struct MoveCompiler {
    hitboxes: Vec<attack::AttackHitbox>,
    active_hitboxes: HashMap<u32, usize>,
    wont_autocancel_window: Option<(u32, u32)>,
    iasa_frame: Option<u32>,
    landing_duration: Option<u32>,
    cursor: u32,
    use_root_motion: bool,
}

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct SharedHitboxInfo {
    damage: FighterDamage,
    angle: AttackAngle,
    knockback_type: KnockbackType,
    knockback: Knockback,
    knockback_growth: FGi32,
}

impl SharedHitboxInfo {
    pub fn make(
        damage: FighterDamage,
        angle: AttackAngle,
        knockback_type: KnockbackType,
        knockback: Knockback,
        knockback_growth: FGi32,
    ) -> Self {
        Self {
            damage,
            angle,
            knockback_type,
            knockback,
            knockback_growth,
        }
    }
}

fn rhai_u32(value: i64, parameter: &str) -> Result<u32, Box<EvalAltResult>> {
    u32::try_from(value).map_err(|_| {
        format!(
            "{parameter} must be between 0 and {}, got {value}",
            u32::MAX
        )
        .into()
    })
}

impl MoveCompiler {
    fn frame(&mut self, new_frame: u32) -> Result<(), Box<EvalAltResult>> {
        if new_frame < self.cursor {
            return Err(format!(
                "Tried to go back from frame {} -> {}, rolling back is not supported",
                self.cursor, new_frame
            )
            .into());
        }

        self.cursor = new_frame;

        Ok(())
    }

    fn remove_hitbox(&mut self, id: u32) -> Result<(), Box<EvalAltResult>> {
        let hitbox_index = *self
            .active_hitboxes
            .get(&id)
            .ok_or(format!("Can't remove hitbox with {id}: Hitbox not found"))?;
        self.hitboxes[hitbox_index].end_frame = self.cursor;
        self.active_hitboxes.remove(&id);

        Ok(())
    }
    fn remove_all_hitboxes(&mut self) -> Result<(), Box<EvalAltResult>> {
        // We are gonna modify the array, so take a copy...
        for hitbox_id in self.active_hitboxes.keys().copied().collect::<Vec<_>>() {
            self.remove_hitbox(hitbox_id)?;
        }

        self.active_hitboxes.clear();

        Ok(())
    }

    fn set_wont_autocancel_window(
        &mut self,
        start: u32,
        end: u32,
    ) -> Result<(), Box<EvalAltResult>> {
        if end < start {
            return Err(format!(
                "Couldn't set won't autocancel window, end frame {} is smaller than start frame {}",
                end, start
            )
            .into());
        }

        self.wont_autocancel_window = Some((start, end));

        Ok(())
    }

    fn set_iasa_frame(&mut self, iasa_frame: u32) {
        self.iasa_frame = Some(iasa_frame);
    }

    fn enable_root_motion(&mut self) {
        self.use_root_motion = true;
    }

    fn into_attack_script(mut self) -> Result<FighterAttackScript, String> {
        self.remove_all_hitboxes().map_err(|err| err.to_string())?;
        let mut hitboxes = self.hitboxes;
        hitboxes.sort_unstable_by_key(|hitbox| (hitbox.id, hitbox.start_frame));

        Ok(FighterAttackScript {
            hitboxes,
            iasa_frame: self.iasa_frame,
            wont_autocancel_window: self.wont_autocancel_window,
            landing_duration: self.landing_duration,
            use_root_motion: self.use_root_motion,
        })
    }
}

// scripting functions

pub fn register_common_types(engine: &mut rhai::Engine) {
    engine
        .register_type_with_name::<KnockbackType>("KnockbackType")
        .register_fn("knockback_normal", KnockbackType::make_normal)
        .register_fn("knockback_fixed", KnockbackType::make_fixed)
        .register_type_with_name::<AttackAngle>("Angle")
        .register_fn("angle", |str: &str| {
            AttackAngle::make_normal(FGi32::lit(str))
        })
        .register_fn("angle_sakurai", AttackAngle::make_sakurai)
        .register_type_with_name::<Knockback>("Knockback")
        .register_fn("knockback", Knockback::lit)
        .register_type_with_name::<FGi32>("FGi32")
        .register_fn("fg32", FGi32::lit)
        .register_type_with_name::<FGVec2>("Vec2")
        .register_fn("vec2", FGVec2::lit)
        .register_type_with_name::<FGVec3>("Vec3")
        .register_fn("vec3", FGVec3::lit)
        .register_type_with_name::<FighterDamage>("FighterDamage")
        .register_type_with_name::<SharedHitboxInfo>("SharedHitboxInfo")
        .register_fn(
            "make_hitbox_info",
            |damage: FighterDamage,
             angle: AttackAngle,
             knockback_type: KnockbackType,
             knockback: Knockback,
             knockback_growth: FGi32| {
                SharedHitboxInfo::make(damage, angle, knockback_type, knockback, knockback_growth)
            },
        )
        .register_fn("damage", FighterDamage::lit);
}

pub fn compile_script(text: &str) -> Result<FighterAttackScript, String> {
    let mc = Rc::new(RefCell::new(MoveCompiler {
        cursor: 0,
        ..Default::default()
    }));
    let result = {
        let mut engine = rhai::Engine::new();

        register_common_types(&mut engine);

        {
            let mc = mc.clone();
            engine.register_fn(
                "frame",
                move |frame: i64| -> Result<(), Box<EvalAltResult>> {
                    let mut b = mc.borrow_mut();
                    b.frame(rhai_u32(frame, "frame")?)
                },
            );
        }

        {
            let mc = mc.clone();
            engine.register_fn(
                "remove_all_hitboxes",
                move || -> Result<(), Box<EvalAltResult>> {
                    let mut b = mc.borrow_mut();
                    b.remove_all_hitboxes()
                },
            );
        }

        {
            let mc = mc.clone();
            engine.register_fn(
                "remove_hitbox",
                move |id: i64| -> Result<(), Box<EvalAltResult>> {
                    let mut b = mc.borrow_mut();
                    b.remove_hitbox(rhai_u32(id, "hitbox id")?)
                },
            );
        }

        {
            let mc = mc.clone();
            engine.register_fn(
                "set_wont_autocancel_window",
                move |start: i64, end: i64| -> Result<(), Box<EvalAltResult>> {
                    let mut b = mc.borrow_mut();
                    b.set_wont_autocancel_window(
                        rhai_u32(start, "autocancel start")?,
                        rhai_u32(end, "autocancel end")?,
                    )
                },
            );
        }

        {
            let mc = mc.clone();
            engine.register_fn(
                "set_iasa_frame",
                move |frame: i64| -> Result<(), Box<EvalAltResult>> {
                    let mut b = mc.borrow_mut();
                    b.set_iasa_frame(rhai_u32(frame, "IASA frame")?);
                    Ok(())
                },
            );
        }

        {
            let mc = mc.clone();
            engine.register_fn(
                "enable_root_motion",
                move || -> Result<(), Box<EvalAltResult>> {
                    let mut b = mc.borrow_mut();
                    b.enable_root_motion();
                    Ok(())
                },
            );
        }

        {
            let mc = mc.clone();
            engine.register_fn(
                "hitbox",
                move |id: i64,
                      group: i64,
                      bone: &str,
                      offset: FGVec3,
                      radius: FGi32,
                      shared: SharedHitboxInfo|
                      -> Result<(), Box<EvalAltResult>> {
                    add_hitbox(
                        &mc,
                        id,
                        group,
                        bone,
                        shared.damage,
                        offset,
                        radius,
                        shared.angle.clone(),
                        shared.knockback_type,
                        shared.knockback,
                        shared.knockback_growth,
                    )
                },
            );
        }

        // Keep the original script syntax working. Melee's default collision
        // group is shared by all hitbox slots in an attack.
        {
            let mc = mc.clone();
            engine.register_fn(
                "hitbox",
                move |id: i64,
                      bone: &str,
                      damage: FighterDamage,
                      offset: FGVec3,
                      radius: FGi32,
                      angle: AttackAngle,
                      knockback_type: KnockbackType,
                      knockback: Knockback,
                      knockback_growth: FGi32|
                      -> Result<(), Box<EvalAltResult>> {
                    add_hitbox(
                        &mc,
                        id,
                        0,
                        bone,
                        damage,
                        offset,
                        radius,
                        angle,
                        knockback_type,
                        knockback,
                        knockback_growth,
                    )
                },
            );
        }

        {
            let mc = mc.clone();
            engine.register_fn(
                "hitbox",
                move |id: i64,
                      group: i64,
                      bone: &str,
                      damage: FighterDamage,
                      offset: FGVec3,
                      radius: FGi32,
                      angle: AttackAngle,
                      knockback_type: KnockbackType,
                      knockback: Knockback,
                      knockback_growth: FGi32|
                      -> Result<(), Box<EvalAltResult>> {
                    add_hitbox(
                        &mc,
                        id,
                        group,
                        bone,
                        damage,
                        offset,
                        radius,
                        angle,
                        knockback_type,
                        knockback,
                        knockback_growth,
                    )
                },
            );
        }
        engine.run(&text)
    };
    let script = result
        .map(|_| {
            Rc::into_inner(mc)
                .unwrap()
                .into_inner()
                .into_attack_script()
        })
        .map_err(|err| err.to_string())??;
    script.validate().map(|()| script)
}

fn add_hitbox(
    mc: &Rc<RefCell<MoveCompiler>>,
    raw_id: i64,
    raw_group: i64,
    bone: &str,
    damage: FighterDamage,
    offset: FGVec3,
    radius: FGi32,
    angle: AttackAngle,
    knockback_type: KnockbackType,
    knockback: Knockback,
    knockback_growth: FGi32,
) -> Result<(), Box<EvalAltResult>> {
    let mut b = mc.borrow_mut();
    let id = rhai_u32(raw_id, "hitbox id")?;
    let group = rhai_u32(raw_group, "hitbox group")?;
    if b.active_hitboxes.contains_key(&id) {
        return Err(format!("Hitbox with ID {} is already active!", id).into());
    }
    let hitbox_index = b.hitboxes.len();
    let frame = b.cursor;
    b.hitboxes.push(attack::AttackHitbox {
        id,
        group,
        bone: bone.to_owned(),
        damage,
        offset,
        radius,
        angle,
        knockback_type,
        knockback,
        knockback_growth,
        start_frame: frame,
        end_frame: frame,
    });
    b.active_hitboxes.insert(id, hitbox_index);
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    const HITBOX: &str = "hitbox(0, \"hand_l\", damage(\"1\"), vec3(\"0\", \"0\", \"0\"), fg32(\"1\"), angle(\"0\"), knockback_normal(), knockback(\"1\"), fg32(\"1\"));";

    #[test]
    fn compiled_hitboxes_are_sorted_by_authored_id() {
        let script = compile_script(
            r#"
                frame(1);
                hitbox(20, "hand_l", damage("1"), vec3("0", "0", "0"), fg32("1"), angle("0"), knockback_normal(), knockback("1"), fg32("1"));
                frame(2);
                hitbox(3, "hand_l", damage("2"), vec3("0", "0", "0"), fg32("1"), angle("0"), knockback_normal(), knockback("1"), fg32("1"));
                frame(3);
                hitbox(10, "hand_l", damage("3"), vec3("0", "0", "0"), fg32("1"), angle("0"), knockback_normal(), knockback("1"), fg32("1"));
                frame(4);
            "#,
        )
        .unwrap();

        assert_eq!(
            script
                .hitboxes
                .iter()
                .map(|hitbox| hitbox.id)
                .collect::<Vec<_>>(),
            vec![3, 10, 20]
        );
        assert_eq!(
            script
                .hitboxes
                .iter()
                .map(|hitbox| hitbox.start_frame)
                .collect::<Vec<_>>(),
            vec![2, 3, 1]
        );
    }

    #[test]
    fn negative_and_overflowing_frames_are_rejected() {
        for source in ["frame(-1);", "frame(4294967296);"] {
            let error = match compile_script(source) {
                Ok(_) => panic!("invalid frame was accepted"),
                Err(error) => error,
            };
            assert!(error.contains("frame"));
        }
    }

    #[test]
    fn remove_hitbox_ends_only_the_named_hitbox() {
        let script =
            compile_script(&format!("frame(1); {HITBOX} frame(2); remove_hitbox(0);")).unwrap();

        assert_eq!(script.hitboxes.len(), 1);
        assert_eq!(script.hitboxes[0].start_frame, 1);
        assert_eq!(script.hitboxes[0].end_frame, 2);
    }

    #[test]
    fn hitbox_ids_can_be_reused_after_removal() {
        let script = compile_script(&format!(
            "frame(1); {HITBOX} frame(2); remove_hitbox(0); frame(4); {HITBOX} frame(5);"
        ))
        .unwrap();

        assert_eq!(script.hitboxes.len(), 2);
        assert_eq!(script.hitboxes[0].start_frame, 1);
        assert_eq!(script.hitboxes[0].end_frame, 2);
        assert_eq!(script.hitboxes[1].start_frame, 4);
        assert_eq!(script.hitboxes[1].end_frame, 5);
    }

    #[test]
    fn explicit_hitbox_group_is_compiled() {
        let script = compile_script(
            r#"frame(1); hitbox(7, 3, "hand_l", damage("1"), vec3("0", "0", "0"), fg32("1"), angle("0"), knockback_normal(), knockback("1"), fg32("1")); frame(2);"#,
        )
        .unwrap();

        assert_eq!(script.hitboxes[0].id, 7);
        assert_eq!(script.hitboxes[0].group, 3);
    }

    #[test]
    fn legacy_hitboxes_share_the_default_collision_group() {
        let script = compile_script(
            r#"
                frame(1);
                hitbox(0, "hand_l", damage("1"), vec3("0", "0", "0"), fg32("1"), angle("0"), knockback_normal(), knockback("1"), fg32("1"));
                hitbox(1, "hand_l", damage("1"), vec3("0", "0", "0"), fg32("1"), angle("0"), knockback_normal(), knockback("1"), fg32("1"));
                frame(2);
            "#,
        )
        .unwrap();

        assert_eq!(script.hitboxes.iter().map(|hitbox| hitbox.group).collect::<Vec<_>>(), vec![0, 0]);
    }
}