bevy-expedition-demo/src/input.rs

use std::{any::{Any, TypeId}, collections::HashMap, hash::Hash, marker::PhantomData};

use bevy::prelude::*;
use bevy_common_assets::toml::TomlAssetPlugin;
use leafwing_input_manager::{Actionlike, prelude::{Axislike, Buttonlike, DualAxislike, InputMap}};
use serde::{Deserialize, Serialize, de::DeserializeOwned};

const INPUT_ASSET_EXTENSIONS: [&'static str; 1] = ["input.toml"];

#[derive(Clone, Debug, Serialize, Deserialize, Default, PartialEq, Eq)]
pub struct MultiInput {
    pub keyboard: Option<Vec<KeyCode>>,
    pub mouse: Option<Vec<MouseButton>>,
    pub gamepad: Option<Vec<GamepadButton>>,
}

impl From<MultiInput> for InputKind {
    fn from(value: MultiInput) -> Self {
        Self::Button(value)
    }
}

#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
#[serde(untagged)]
pub enum InputKind {
    Button(MultiInput),
    Axis(Vec<Box<dyn Axislike>>),
    DualAxis(Vec<Box<dyn DualAxislike>>),
}

#[derive(Default, Deref, DerefMut, Debug, Asset, Reflect, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct InputAsset<Name> 
    where Name: Sized + Hash + Eq + Reflect + TypePath + Actionlike {
    #[serde(flatten)]
    events: HashMap<Name, InputKind>,
}

fn copy_keys<Name, Button>(input_map: &mut InputMap<Name>, name: &Name, buttons: &Vec<Button>)
    where Name: Sized + Hash + Eq + Reflect + TypePath + Actionlike,
    Button: Buttonlike + Clone {
    let bindings: Vec<(Name, Button)> = buttons.iter()
        .map(|k| (name.clone(), k.to_owned())).collect();
    input_map.insert_multiple(bindings);
}

impl<Name> From<InputAsset<Name>> for InputMap<Name>
    where Name: Sized + Hash + Eq + Reflect + TypePath + Actionlike {
    fn from(asset: InputAsset<Name>) -> Self {
        let new_asset = (*asset).clone();
        let mut input_map = InputMap::default();
        for (name, input_kind) in new_asset.into_iter() {
            match input_kind {
                InputKind::Button(input) => {
                    if let Some(buttons) = input.keyboard {
                        copy_keys(&mut input_map, &name, &buttons);
                    }
                    if let Some(buttons) = input.mouse {
                        copy_keys(&mut input_map, &name, &buttons);
                    }
                    if let Some(buttons) = input.gamepad {
                        copy_keys(&mut input_map, &name, &buttons);
                    }
                },
                InputKind::Axis(axises) => {
                    for axis in axises {
                        input_map.insert_axis_boxed(name.clone(), axis);
                    }
                },
                InputKind::DualAxis(axises) => {
                    for axis in axises {
                        input_map.insert_dual_axis_boxed(name.clone(), axis);
                    }
                },
            }
        }
        input_map
    }
}

impl<Name> InputAsset<Name>
    where Name: Sized + Hash + Eq + Reflect + TypePath + Actionlike {
    /// This method does several things:
    /// - Replace all actions which both are contained in the [`InputAsset`] and in the [`InputMap`]
    /// and have the same "type" (e.g. [`Buttonlike`]).
    /// - Insert all actions which are not contained in the map
    /// Common usecase for this is to update default values with this asset
    pub fn replace_input_map_actions(&self, map: &mut InputMap<Name>) {
        for (name, input_kind) in self.iter() {
            match input_kind {
                InputKind::Button(input) => {
                    if map.iter_buttonlike().find(|(n, _)| n == &name).is_some() {
                        map.clear_action(name);
                    }
                    if let Some(buttons) = &input.keyboard {
                        copy_keys(map, name, &buttons);
                    }
                    if let Some(buttons) = &input.mouse {
                        copy_keys(map, name, &buttons);
                    }
                    if let Some(buttons) = &input.gamepad {
                        copy_keys(map, name, &buttons);
                    }
                },
                InputKind::Axis(axises) => {
                    if map.iter_axislike().find(|(n, _)| n == &name).is_some() {
                        map.clear_action(name);
                    }
                    for axis in axises.iter() {
                        map.insert_axis_boxed(name.to_owned(), axis.to_owned());
                    }
                },
                InputKind::DualAxis(axises) => {
                    if map.iter_dual_axislike().find(|(n, _)| n == &name).is_some() {
                        map.clear_action(name);
                    }
                    for axis in axises.iter() {
                        map.insert_dual_axis_boxed(name.to_owned(), axis.to_owned());
                    }
                }
            }
        }
    }
}

fn try_multi_input_insert<T: Clone>(input_type: &mut Option<Vec<T>>, new_key: T) {
    match input_type {
        Some(key) => key.push(new_key),
        None => {
            let wrapped_new_key = Some(vec![new_key]);
            *input_type = wrapped_new_key;
        },
    }
}

impl<Name> From<InputMap<Name>> for InputAsset<Name>
    where Name: Sized + Hash + Eq + Reflect + TypePath + Actionlike + Default {
    fn from(map: InputMap<Name>) -> Self {
        let mut asset: InputAsset<Name> = InputAsset::default();
        const KC_TYPE: TypeId = TypeId::of::<KeyCode>();
        const MB_TYPE: TypeId = TypeId::of::<MouseButton>();
        const GP_TYPE: TypeId = TypeId::of::<GamepadButton>();
        for (name, buttonlikes) in map.iter_buttonlike() {
            let multi_input = {
                let input_kind = match asset.get_mut(name) {
                    Some(input_kind) => input_kind,
                    None => {
                        asset.insert(name.to_owned(), MultiInput::default().into());
                        asset.get_mut(name).unwrap()
                    }
                };
                let InputKind::Button(btn) = input_kind else {
                    continue;
                };
                btn
            };
            for buttonlike in buttonlikes.iter() {
                let buttonlike = &(**buttonlike) as &dyn Any;
                let checked_type = buttonlike.type_id();
                if checked_type == KC_TYPE {
                    let new_key = (*buttonlike).downcast_ref::<KeyCode>().unwrap().to_owned();
                    try_multi_input_insert(&mut multi_input.keyboard, new_key);
                } else if checked_type == MB_TYPE {
                    let new_key = (*buttonlike).downcast_ref::<MouseButton>().unwrap().to_owned();
                    try_multi_input_insert(&mut multi_input.mouse, new_key);
                } else if checked_type == GP_TYPE {
                    let new_key = (*buttonlike).downcast_ref::<GamepadButton>().unwrap().to_owned();
                    try_multi_input_insert(&mut multi_input.gamepad, new_key);
                }
            }
        }

        for (name, axislikes) in map.iter_axislike() {
            match asset.get_mut(name) {
                Some(input_kind) => {
                    let InputKind::Axis(axises) = input_kind else {
                        continue;
                    };
                    for axislike in axislikes {
                        axises.push(axislike.clone());
                    }
                },
                None => {
                    asset.insert(name.clone(), InputKind::Axis(axislikes.clone()));
                }
            }
        }

        for (name, axislikes) in map.iter_dual_axislike() {
            match asset.get_mut(name) {
                Some(input_kind) => {
                    let InputKind::DualAxis(axises) = input_kind else {
                        continue;
                    };
                    for axislike in axislikes {
                        axises.push(axislike.clone());
                    }
                },
                None => {
                    asset.insert(name.clone(), InputKind::DualAxis(axislikes.clone()));
                }
            }
        }

        asset
    }
}

#[derive(Debug, Default)]
pub struct InputAssetPlugin<T> (PhantomData<T>);

impl<T> Plugin for InputAssetPlugin<T>
where T: Sized + Hash + Eq + Reflect + TypePath + Actionlike + DeserializeOwned
{
    fn build(&self, app: &mut App) {
        app.add_plugins(TomlAssetPlugin::<InputAsset<T>>::new(&INPUT_ASSET_EXTENSIONS));
    }
}