Skip to main content

steel_core/player/player_inventory/
equipment.rs

1use std::mem;
2
3use steel_registry::{
4    REGISTRY, RegistryExt, enchantment_effect::EnchantmentEffectComponent, item_stack::ItemStack,
5    items::ItemRef,
6};
7use steel_utils::types::InteractionHand;
8
9use crate::inventory::{
10    container::Container,
11    equipment::{EntityEquipment, EquipmentSlot},
12};
13
14use super::core::PlayerInventory;
15
16/// Result of swapping a held item with an equipment slot.
17#[derive(Debug, PartialEq)]
18pub enum EquipmentSwapResult {
19    /// The swap succeeded. Contains an overflow stack that should be dropped if non-empty.
20    Success(ItemStack),
21    /// The swap is blocked by vanilla equipment rules.
22    Fail,
23}
24
25const fn hand_to_equipment_slot(hand: InteractionHand) -> EquipmentSlot {
26    match hand {
27        InteractionHand::MainHand => EquipmentSlot::MainHand,
28        InteractionHand::OffHand => EquipmentSlot::OffHand,
29    }
30}
31
32const fn equipment_to_slot(slot: EquipmentSlot, selected: u8) -> usize {
33    match slot {
34        EquipmentSlot::MainHand => selected as usize,
35        EquipmentSlot::OffHand => 40,
36        EquipmentSlot::Feet => 36,
37        EquipmentSlot::Legs => 37,
38        EquipmentSlot::Chest => 38,
39        EquipmentSlot::Head => 39,
40        EquipmentSlot::Body => 41,
41        EquipmentSlot::Saddle => 42,
42    }
43}
44
45impl PlayerInventory {
46    /// Gets the item in the specified hand.
47    #[must_use]
48    pub fn get_item_in_hand(&self, hand: InteractionHand) -> &ItemStack {
49        match hand {
50            InteractionHand::MainHand => self.get_selected_item(),
51            InteractionHand::OffHand => self.get_offhand_item(),
52        }
53    }
54
55    /// Gets the item in the specified hand.
56    #[must_use]
57    pub fn get_item_in_hand_mut(&mut self, hand: InteractionHand) -> &mut ItemStack {
58        match hand {
59            InteractionHand::MainHand => self.get_selected_item_mut(),
60            InteractionHand::OffHand => self.get_offhand_item_mut(),
61        }
62    }
63
64    /// Sets the item in the specified hand.
65    pub fn set_item_in_hand(&mut self, hand: InteractionHand, item: ItemStack) {
66        match hand {
67            InteractionHand::MainHand => self.set_selected_item(item),
68            InteractionHand::OffHand => self.set_offhand_item(item),
69        }
70    }
71
72    /// Shrinks the item in the specified hand and records inventory/equipment changes.
73    pub fn shrink_item_in_hand(&mut self, hand: InteractionHand, amount: i32) {
74        if amount <= 0 || self.get_item_in_hand(hand).is_empty() {
75            return;
76        }
77
78        self.get_item_in_hand_mut(hand).shrink(amount);
79        self.set_changed();
80    }
81
82    /// Splits items from the specified hand and records inventory/equipment changes.
83    pub fn split_item_in_hand(&mut self, hand: InteractionHand, amount: i32) -> ItemStack {
84        if amount <= 0 || self.get_item_in_hand(hand).is_empty() {
85            return ItemStack::empty();
86        }
87
88        let result = self.get_item_in_hand_mut(hand).split(amount);
89        self.set_changed();
90        result
91    }
92
93    /// Damages the held item and records inventory/equipment changes.
94    pub fn hurt_item_in_hand(
95        &mut self,
96        hand: InteractionHand,
97        amount: i32,
98        has_infinite_materials: bool,
99    ) {
100        if amount <= 0 || self.get_item_in_hand(hand).is_empty() {
101            return;
102        }
103
104        let changed = {
105            let item = self.get_item_in_hand_mut(hand);
106            let previous_item = item.item();
107            let previous_count = item.count();
108            let previous_damage = item.get_damage_value();
109
110            let _ = item.hurt_and_break(amount, has_infinite_materials);
111
112            item.item() != previous_item
113                || item.count() != previous_count
114                || item.get_damage_value() != previous_damage
115        };
116
117        if changed {
118            self.set_changed();
119        }
120    }
121
122    /// Mutates the held item and records inventory/equipment changes if its stack state changed.
123    pub fn mutate_item_in_hand<R>(
124        &mut self,
125        hand: InteractionHand,
126        f: impl FnOnce(&mut ItemStack) -> R,
127    ) -> R {
128        self.with_equipment_item_mut(hand_to_equipment_slot(hand), f)
129    }
130
131    /// Damages the held item and converts it to `replacement_item` if it breaks.
132    ///
133    /// Mirrors vanilla `ItemStack.hurtAndConvertOnBreak` for hand-held player items.
134    pub fn hurt_and_convert_item_in_hand_on_break(
135        &mut self,
136        hand: InteractionHand,
137        amount: i32,
138        replacement_item: ItemRef,
139        has_infinite_materials: bool,
140    ) {
141        if amount <= 0 || self.get_item_in_hand(hand).is_empty() {
142            return;
143        }
144
145        let changed = {
146            let item = self.get_item_in_hand_mut(hand);
147            let previous_item = item.item();
148            let previous_count = item.count();
149            let previous_damage = item.get_damage_value();
150
151            if item.hurt_and_break(amount, has_infinite_materials) && item.is_empty() {
152                item.set_item(&replacement_item.key);
153                item.set_count(1);
154                if item.is_damageable_item() {
155                    item.set_damage_value(0);
156                }
157            }
158
159            item.item() != previous_item
160                || item.count() != previous_count
161                || item.get_damage_value() != previous_damage
162        };
163
164        if changed {
165            self.set_changed();
166        }
167    }
168
169    /// Swaps the selected main-hand item with the offhand item.
170    ///
171    /// Returns true when the visible hand contents changed.
172    pub fn swap_hands(&mut self) -> bool {
173        if ItemStack::matches(self.get_selected_item(), self.get_offhand_item()) {
174            return false;
175        }
176
177        let main_hand = EntityEquipment::take(self, EquipmentSlot::MainHand);
178        let offhand = EntityEquipment::take(self, EquipmentSlot::OffHand);
179        let _ = EntityEquipment::set(self, EquipmentSlot::MainHand, offhand);
180        let _ = EntityEquipment::set(self, EquipmentSlot::OffHand, main_hand);
181        true
182    }
183
184    /// Attempts to equip the held item into the target equipment slot.
185    pub fn try_swap_with_equipment_slot(
186        &mut self,
187        hand: InteractionHand,
188        slot: EquipmentSlot,
189        has_infinite_materials: bool,
190    ) -> EquipmentSwapResult {
191        let in_hand = self.get_item_in_hand(hand);
192        if in_hand.is_empty() {
193            return EquipmentSwapResult::Fail;
194        }
195
196        let in_equipment_slot = EntityEquipment::get_ref(self, slot);
197        if ItemStack::is_same_item_same_components(in_hand, in_equipment_slot) {
198            return EquipmentSwapResult::Fail;
199        }
200
201        if !has_infinite_materials
202            && in_equipment_slot
203                .has_enchantment_effect(EnchantmentEffectComponent::PreventArmorChange)
204        {
205            return EquipmentSwapResult::Fail;
206        }
207
208        if in_hand.count() <= 1 {
209            self.swap_single_item_with_equipment_slot(hand, slot, has_infinite_materials);
210            return EquipmentSwapResult::Success(ItemStack::empty());
211        }
212
213        let to_equip = in_hand.copy_with_count(1);
214        if !has_infinite_materials {
215            self.get_item_in_hand_mut(hand).shrink(1);
216        }
217        let mut overflow = EntityEquipment::set(self, slot, to_equip);
218        if !overflow.is_empty() && self.add(&mut overflow) {
219            overflow = ItemStack::empty();
220        }
221
222        EquipmentSwapResult::Success(overflow)
223    }
224
225    /// Repairs a random damaged equipped item with `REPAIR_WITH_XP`, returning leftover XP.
226    pub fn repair_random_equipped_item_with_xp(&mut self, amount: i32) -> i32 {
227        let mut remaining = amount;
228
229        loop {
230            let candidates = self.repair_with_xp_candidate_slots();
231            if candidates.is_empty() {
232                return remaining;
233            }
234
235            let selected = rand::random_range(0..candidates.len());
236            let slot = candidates[selected];
237            let item = EntityEquipment::get_mut(self, slot);
238            let to_repair = item
239                .apply_unconditional_enchantment_value_effects(
240                    EnchantmentEffectComponent::RepairWithXp,
241                    remaining as f32,
242                )
243                .max(0.0) as i32;
244            if to_repair <= 0 {
245                return 0;
246            }
247
248            let damage = item.get_damage_value();
249            let repair = to_repair.min(damage);
250            if repair <= 0 {
251                return 0;
252            }
253
254            item.set_damage_value(damage - repair);
255            self.set_changed();
256
257            remaining -= repair * remaining / to_repair;
258            if remaining <= 0 {
259                return 0;
260            }
261        }
262    }
263
264    fn swap_single_item_with_equipment_slot(
265        &mut self,
266        hand: InteractionHand,
267        slot: EquipmentSlot,
268        has_infinite_materials: bool,
269    ) {
270        if has_infinite_materials {
271            let held = self
272                .get_item_in_hand(hand)
273                .copy_with_count(self.get_item_in_hand(hand).count());
274            let previous = EntityEquipment::set(self, slot, held);
275            if !previous.is_empty() {
276                self.set_item_in_hand(hand, previous);
277            }
278            return;
279        }
280
281        let held = self.take_item_in_hand(hand);
282        let previous = EntityEquipment::set(self, slot, held);
283        self.set_item_in_hand(hand, previous);
284    }
285
286    fn repair_with_xp_candidate_slots(&self) -> Vec<EquipmentSlot> {
287        let mut slots = Vec::new();
288        for slot in EquipmentSlot::ALL {
289            let item = EntityEquipment::get_ref(self, slot);
290            if !item.is_damaged() {
291                continue;
292            }
293
294            let Some(enchantments) = item.get_enchantments() else {
295                continue;
296            };
297            for (key, level) in enchantments.iter() {
298                if *level == 0 {
299                    continue;
300                }
301                let Some(enchantment) = REGISTRY.enchantments.by_key(key) else {
302                    continue;
303                };
304                if enchantment
305                    .effects
306                    .has(EnchantmentEffectComponent::RepairWithXp)
307                    && enchantment.matching_slot(slot)
308                {
309                    slots.push(slot);
310                }
311            }
312        }
313        slots
314    }
315
316    fn take_item_in_hand(&mut self, hand: InteractionHand) -> ItemStack {
317        match hand {
318            InteractionHand::MainHand => EntityEquipment::take(self, EquipmentSlot::MainHand),
319            InteractionHand::OffHand => EntityEquipment::take(self, EquipmentSlot::OffHand),
320        }
321    }
322}
323
324impl PlayerInventory {
325    pub(super) const fn equipment_slot_index(&self, slot: EquipmentSlot) -> usize {
326        equipment_to_slot(slot, self.selected)
327    }
328}
329
330impl EntityEquipment for PlayerInventory {
331    fn get_ref(&self, slot: EquipmentSlot) -> &ItemStack {
332        &self.items[self.equipment_slot_index(slot)]
333    }
334
335    fn get_mut(&mut self, slot: EquipmentSlot) -> &mut ItemStack {
336        let inventory_index = self.equipment_slot_index(slot);
337        &mut self.items[inventory_index]
338    }
339
340    fn set(&mut self, slot: EquipmentSlot, stack: ItemStack) -> ItemStack {
341        let inventory_index = self.equipment_slot_index(slot);
342        let old = mem::replace(&mut self.items[inventory_index], stack);
343        Container::set_changed(self);
344        old
345    }
346
347    fn take(&mut self, slot: EquipmentSlot) -> ItemStack {
348        let inventory_index = self.equipment_slot_index(slot);
349        let old = mem::take(&mut self.items[inventory_index]);
350        if !old.is_empty() {
351            Container::set_changed(self);
352        }
353        old
354    }
355
356    fn clear(&mut self) {
357        let mut changed = false;
358        for slot in EquipmentSlot::ALL {
359            let inventory_index = self.equipment_slot_index(slot);
360            if self.items[inventory_index].is_empty() {
361                continue;
362            }
363
364            self.items[inventory_index] = ItemStack::empty();
365            changed = true;
366        }
367        if changed {
368            Container::set_changed(self);
369        }
370    }
371}