1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142

// Docs are missing from generated types :(
#![allow(missing_docs)]

use crate::{charges::ChargeState, cooldown::CooldownState, Abilitylike};
use bevy::ecs::query::WorldQuery;
use leafwing_input_manager::action_state::ActionState;

/// A custom [`WorldQuery`](bevy::ecs::query::WorldQuery) type that fetches all ability relevant data for you.
///
/// This type is intended to make collecting the data for [`Abilitylike`] methods easier when working with a full [`AbilitiesBundle`](crate::AbilitiesBundle`).
/// This struct can be used as the first type parameter in a [`Query`](bevy::ecs::system::Query) to fetch the appropriate data.
///
/// Once you have a [`AbilityStateItem`] by calling `.iter_mut()` or `.single_mut` on your query
/// (or a [`AbilityStateReadOnlyItem`] by calling `.iter()` or `.single`),
/// you can use the methods defined there to perform common tasks quickly and reliably.
#[derive(WorldQuery)]
#[world_query(mutable)]
pub struct AbilityState<A: Abilitylike> {
    /// The [`ActionState`] of the abilities of this entity of type `A`
    pub action_state: &'static ActionState<A>,
    /// The [`ChargeState`] associated with each action of type `A` for this entity
    pub charges: &'static mut ChargeState<A>,
    /// The [`CooldownState`] associated with each action of type `A` for this entity
    pub cooldowns: &'static mut CooldownState<A>,
}

impl<A: Abilitylike> AbilityStateItem<'_, A> {
    /// Is this ability ready?
    ///
    /// Calls [`Abilitylike::ready`] on the specified action.
    #[inline]
    #[must_use]
    pub fn ready(&self, action: A) -> bool {
        action.ready(&*self.charges, &*self.cooldowns)
    }

    /// Is this ability both ready and pressed?
    #[inline]
    pub fn ready_and_pressed(&self, action: A) -> bool {
        self.action_state.pressed(action.clone()) && self.ready(action)
    }

    /// Is this ability both ready and just pressed?
    #[inline]
    pub fn ready_and_just_pressed(&self, action: A) -> bool {
        self.action_state.just_pressed(action.clone()) && self.ready(action)
    }

    /// Triggers this ability, depleting a charge if available.
    ///
    /// Calls [`Abilitylike::trigger`] on the specified action.
    #[inline]
    pub fn trigger(&mut self, action: A) -> bool {
        action.trigger(&mut *self.charges, &mut *self.cooldowns)
    }

    /// Triggers this ability (and depletes available charges), if action is pressed.
    ///
    /// Calls [`Abilitylike::trigger`] on the specified action.
    #[inline]
    pub fn trigger_if_pressed(&mut self, action: A) -> bool {
        if self.action_state.just_pressed(action.clone()) {
            action.trigger(&mut *self.charges, &mut *self.cooldowns)
        } else {
            false
        }
    }

    /// Triggers this ability (and depletes available charges), if action was just pressed.
    ///
    /// Calls [`Abilitylike::trigger`] on the specified action.
    #[inline]
    pub fn trigger_if_just_pressed(&mut self, action: A) -> bool {
        if self.action_state.just_pressed(action.clone()) {
            action.trigger(&mut *self.charges, &mut *self.cooldowns)
        } else {
            false
        }
    }
}

impl<A: Abilitylike> AbilityStateReadOnlyItem<'_, A> {
    /// Is this ability ready?
    ///
    /// Calls [`Abilitylike::ready`] on the specified action.
    #[inline]
    #[must_use]
    pub fn ready(&self, action: A) -> bool {
        action.ready(self.charges, self.cooldowns)
    }

    /// Is this ability both ready and pressed?
    #[inline]
    pub fn ready_and_pressed(&self, action: A) -> bool {
        self.action_state.pressed(action.clone()) && self.ready(action)
    }

    /// Is this ability both ready and just pressed?
    #[inline]
    pub fn ready_and_just_pressed(&self, action: A) -> bool {
        self.action_state.just_pressed(action.clone()) && self.ready(action)
    }
}

#[cfg(test)]
mod tests {
    use crate as leafwing_abilities;
    use crate::{AbilitiesBundle, AbilityState, Abilitylike};
    use bevy::prelude::*;
    use leafwing_input_manager::{action_state::ActionState, Actionlike};

    #[derive(Actionlike, Abilitylike, Clone, Debug)]
    enum TestAction {
        Duck,
        Cover,
    }

    #[test]
    fn ability_state_methods_are_visible_from_query() {
        fn simple_system(mut query: Query<AbilityState<TestAction>>) {
            let mut ability_state = query.single_mut();
            if ability_state.ready(TestAction::Duck) {
                ability_state.trigger(TestAction::Duck);
            }
        }

        let mut app = App::new();
        app.add_system(simple_system);
    }

    #[test]
    fn ability_state_fetches_abilities_bundle() {
        let mut world = World::new();
        world
            .spawn()
            .insert_bundle(AbilitiesBundle::<TestAction>::default())
            .insert(ActionState::<TestAction>::default());

        let mut query_state = world.query::<AbilityState<TestAction>>();
        assert_eq!(query_state.iter(&world).len(), 1);
    }
}