azalea_client/plugins/inventory/
equipment_effects.rs

1//! Support for enchantments and items with attribute modifiers.
2
3use std::collections::HashMap;
4
5use azalea_core::{data_registry::ResolvableDataRegistry, registry_holder::value::AttributeEffect};
6use azalea_entity::{Attributes, inventory::Inventory};
7use azalea_inventory::{
8    ItemStack,
9    components::{self, AttributeModifier, EquipmentSlot},
10};
11use azalea_registry::identifier::Identifier;
12use bevy_ecs::{
13    component::Component,
14    entity::Entity,
15    event::EntityEvent,
16    observer::On,
17    query::With,
18    system::{Commands, Query},
19};
20use tracing::{debug, error, warn};
21
22use crate::local_player::InstanceHolder;
23
24/// A component that contains the equipment slots that we had last tick.
25///
26/// This is used by [`collect_equipment_changes`] for applying enchantments.
27#[derive(Component, Debug, Default)]
28pub struct LastEquipmentItems {
29    pub map: HashMap<EquipmentSlot, ItemStack>,
30}
31
32pub fn collect_equipment_changes(
33    mut commands: Commands,
34    mut query: Query<(Entity, &Inventory, Option<&LastEquipmentItems>), With<Attributes>>,
35) {
36    for (entity, inventory, last_equipment_items) in &mut query {
37        let last_equipment_items = if let Some(e) = last_equipment_items {
38            e
39        } else {
40            commands
41                .entity(entity)
42                .insert(LastEquipmentItems::default());
43            continue;
44        };
45
46        let mut changes = HashMap::new();
47
48        for equipment_slot in EquipmentSlot::values() {
49            let current_item = inventory
50                .get_equipment(equipment_slot)
51                .unwrap_or(&ItemStack::Empty);
52            let last_item = last_equipment_items
53                .map
54                .get(&equipment_slot)
55                .unwrap_or(&ItemStack::Empty);
56
57            if current_item == last_item {
58                // item hasn't changed, nothing to do
59                continue;
60            }
61
62            changes.insert(
63                equipment_slot,
64                EquipmentChange {
65                    old: last_item.clone(),
66                    new: current_item.clone(),
67                },
68            );
69        }
70
71        if changes.is_empty() {
72            continue;
73        }
74        commands.trigger(EquipmentChangesEvent {
75            entity,
76            map: changes,
77        });
78    }
79}
80
81#[derive(Debug, EntityEvent)]
82pub struct EquipmentChangesEvent {
83    pub entity: Entity,
84    pub map: HashMap<EquipmentSlot, EquipmentChange>,
85}
86#[derive(Debug)]
87pub struct EquipmentChange {
88    pub old: ItemStack,
89    pub new: ItemStack,
90}
91
92pub fn handle_equipment_changes(
93    equipment_changes: On<EquipmentChangesEvent>,
94    mut query: Query<(&InstanceHolder, &mut LastEquipmentItems, &mut Attributes)>,
95) {
96    let Ok((instance_holder, mut last_equipment_items, mut attributes)) =
97        query.get_mut(equipment_changes.entity)
98    else {
99        error!(
100            "got EquipmentChangesEvent with unknown entity {}",
101            equipment_changes.entity
102        );
103        return;
104    };
105
106    if !equipment_changes.map.is_empty() {
107        debug!("equipment changes: {:?}", equipment_changes.map);
108    }
109
110    for (&slot, change) in &equipment_changes.map {
111        if change.old.is_present() {
112            // stopLocationBasedEffects
113
114            for (attribute, modifier) in
115                collect_attribute_modifiers_from_item(slot, &change.old, instance_holder)
116            {
117                if let Some(attribute) = attributes.get_mut(attribute) {
118                    attribute.remove(&modifier.id);
119                }
120            }
121
122            last_equipment_items.map.remove(&slot);
123        }
124
125        if change.new.is_present() {
126            // see ItemStack.forEachModifier in vanilla
127
128            for (attribute, modifier) in
129                collect_attribute_modifiers_from_item(slot, &change.new, instance_holder)
130            {
131                if let Some(attribute) = attributes.get_mut(attribute) {
132                    attribute.remove(&modifier.id);
133                    attribute.insert(modifier);
134                }
135            }
136
137            // runLocationChangedEffects
138
139            last_equipment_items.map.insert(slot, change.new.clone());
140        }
141    }
142}
143
144fn collect_attribute_modifiers_from_item(
145    slot: EquipmentSlot,
146    item: &ItemStack,
147    instance_holder: &InstanceHolder,
148) -> Vec<(azalea_registry::builtin::Attribute, AttributeModifier)> {
149    let mut modifiers = Vec::new();
150
151    // handle the attribute_modifiers component first
152    let attribute_modifiers = item
153        .get_component::<components::AttributeModifiers>()
154        .unwrap_or_default();
155    for modifier in &attribute_modifiers.modifiers {
156        modifiers.push((modifier.kind, modifier.modifier.clone()));
157    }
158
159    // now handle enchants
160    let enchants = item
161        .get_component::<components::Enchantments>()
162        .unwrap_or_default();
163    if !enchants.levels.is_empty() {
164        let registry_holder = &instance_holder.instance.read().registries;
165        for (enchant, &level) in &enchants.levels {
166            let Some((_enchant_id, enchant_definition)) = enchant.resolve(registry_holder) else {
167                warn!(
168                    "Got equipment with an enchantment that wasn't in the registry, so it couldn't be resolved to an ID"
169                );
170                continue;
171            };
172
173            let effects = enchant_definition.get::<AttributeEffect>();
174            for effect in effects.unwrap_or_default() {
175                // TODO: check if the effect definition allows the slot
176
177                let modifier = AttributeModifier {
178                    id: Identifier::new(format!("{}/{slot}", effect.id)),
179                    amount: effect.amount.calculate(level) as f64,
180                    operation: effect.operation,
181                };
182
183                modifiers.push((effect.attribute, modifier));
184            }
185        }
186    }
187
188    modifiers
189}
azalea_client/plugins/inventory/equipment_effects.rs

azalea_client/plugins/inventory/
equipment_effects.rs
