steel_core/player/
item_cooldowns.rs

1use rustc_hash::FxHashMap;
2use steel_protocol::packets::game::CCooldown;
3use steel_registry::data_components::vanilla_components::USE_COOLDOWN;
4use steel_registry::item_stack::ItemStack;
5use steel_utils::Identifier;
6
7use super::Player;
8
9#[derive(Clone, Copy)]
10struct CooldownInstance {
11    end_time: i32,
12}
13
14#[derive(Default)]
15pub(super) struct ItemCooldowns {
16    cooldowns: FxHashMap<Identifier, CooldownInstance>,
17    tick_count: i32,
18}
19
20impl ItemCooldowns {
21    #[must_use]
22    pub(super) fn is_on_cooldown(&self, stack: &ItemStack) -> bool {
23        let group = cooldown_group(stack);
24        self.cooldowns
25            .get(&group)
26            .is_some_and(|cooldown| cooldown.end_time > self.tick_count)
27    }
28
29    pub(super) fn tick(&mut self) -> Vec<Identifier> {
30        self.tick_count = self.tick_count.wrapping_add(1);
31        let mut ended = Vec::new();
32        self.cooldowns.retain(|group, cooldown| {
33            if cooldown.end_time <= self.tick_count {
34                ended.push(group.clone());
35                false
36            } else {
37                true
38            }
39        });
40        ended
41    }
42
43    pub(super) fn add_from_stack(&mut self, stack: &ItemStack) -> Option<(Identifier, i32)> {
44        let cooldown = stack.get(USE_COOLDOWN)?;
45        let duration = cooldown.ticks();
46        if duration <= 0 {
47            return None;
48        }
49        let group = cooldown_group(stack);
50        self.cooldowns.insert(
51            group.clone(),
52            CooldownInstance {
53                end_time: self.tick_count + duration,
54            },
55        );
56        Some((group, duration))
57    }
58}
59
60fn cooldown_group(stack: &ItemStack) -> Identifier {
61    stack
62        .get(USE_COOLDOWN)
63        .and_then(|cooldown| cooldown.cooldown_group.clone())
64        .unwrap_or_else(|| stack.item().key.clone())
65}
66
67impl Player {
68    /// Returns whether the stack's vanilla cooldown group is currently active.
69    pub fn is_item_on_cooldown(&self, stack: &ItemStack) -> bool {
70        self.item_cooldowns.lock().is_on_cooldown(stack)
71    }
72
73    /// Starts the stack's vanilla `use_cooldown`, if it has one.
74    pub fn apply_item_use_cooldown(&self, stack: &ItemStack) {
75        let cooldown = self.item_cooldowns.lock().add_from_stack(stack);
76        if let Some((cooldown_group, duration)) = cooldown {
77            self.send_packet(CCooldown {
78                cooldown_group,
79                duration,
80            });
81        }
82    }
83
84    pub(super) fn tick_item_cooldowns(&self) {
85        let ended = self.item_cooldowns.lock().tick();
86        for cooldown_group in ended {
87            self.send_packet(CCooldown {
88                cooldown_group,
89                duration: 0,
90            });
91        }
92    }
93}
94
95#[cfg(test)]
96mod tests {
97    use steel_registry::data_components::vanilla_components::{USE_COOLDOWN, UseCooldown};
98    use steel_registry::item_stack::ItemStack;
99    use steel_registry::{init_vanilla_registry, vanilla_items};
100
101    use super::ItemCooldowns;
102
103    #[test]
104    fn cooldown_blocks_until_duration_ticks_pass() {
105        init_vanilla_registry();
106
107        let stack = ItemStack::with_count(&vanilla_items::ENDER_PEARL, 1);
108        let mut cooldowns = ItemCooldowns::default();
109
110        let Some((group, duration)) = cooldowns.add_from_stack(&stack) else {
111            panic!("ender pearl should have a use cooldown");
112        };
113
114        assert_eq!(group, vanilla_items::ENDER_PEARL.key);
115        assert_eq!(duration, 20);
116        assert!(cooldowns.is_on_cooldown(&stack));
117
118        for _ in 0..19 {
119            assert!(cooldowns.tick().is_empty());
120            assert!(cooldowns.is_on_cooldown(&stack));
121        }
122
123        assert_eq!(
124            cooldowns.tick(),
125            vec![vanilla_items::ENDER_PEARL.key.clone()]
126        );
127        assert!(!cooldowns.is_on_cooldown(&stack));
128    }
129
130    #[test]
131    fn explicit_group_is_shared_between_items() {
132        init_vanilla_registry();
133
134        let group = steel_utils::Identifier::vanilla_static("test_group");
135        let mut stack = ItemStack::with_count(&vanilla_items::ENDER_PEARL, 1);
136        stack.set(USE_COOLDOWN, UseCooldown::new(0.5, Some(group.clone())));
137        let mut other = ItemStack::with_count(&vanilla_items::CHORUS_FRUIT, 1);
138        other.set(USE_COOLDOWN, UseCooldown::new(1.0, Some(group.clone())));
139        let unrelated = ItemStack::with_count(&vanilla_items::WIND_CHARGE, 1);
140        let mut cooldowns = ItemCooldowns::default();
141
142        let Some((started_group, duration)) = cooldowns.add_from_stack(&stack) else {
143            panic!("stack should have a use cooldown");
144        };
145
146        assert_eq!(started_group, group);
147        assert_eq!(duration, 10);
148        assert!(cooldowns.is_on_cooldown(&stack));
149        assert!(cooldowns.is_on_cooldown(&other));
150        assert!(!cooldowns.is_on_cooldown(&unrelated));
151    }
152}
steel_core/player/item_cooldowns.rs

steel_core/player/
item_cooldowns.rs
