steel_registry/enchantment/
mod.rs

1pub mod effect;
2
3use crate::equipment::EquipmentSlot;
4pub use crate::equipment::EquipmentSlotGroup;
5use crate::items::ItemRef;
6use crate::{REGISTRY, RegistryEntry, RegistryExt, RegistryTags, TaggedRegistryExt};
7use effect::EnchantmentEffects;
8use rustc_hash::FxHashMap;
9use simdnbt::ToNbtTag;
10use simdnbt::owned::{NbtCompound, NbtList, NbtTag};
11use steel_utils::Identifier;
12
13/// Enchanting cost formula: `base + per_level_above_first * (level - 1)`.
14#[derive(Debug, Clone, Copy)]
15pub struct EnchantmentCost {
16    pub base: i32,
17    pub per_level_above_first: i32,
18}
19
20impl EnchantmentCost {
21    /// Calculates the cost for `level`.
22    #[must_use]
23    pub const fn calculate(self, level: i32) -> i32 {
24        self.base + self.per_level_above_first * (level - 1)
25    }
26}
27
28#[derive(Debug)]
29pub struct Enchantment {
30    pub key: Identifier,
31    pub max_level: u32,
32    pub min_cost: EnchantmentCost,
33    pub max_cost: EnchantmentCost,
34    pub anvil_cost: i32,
35    pub weight: u32,
36    pub slots: &'static [EquipmentSlotGroup],
37    /// Item tag the enchantment applies to (`#tag` in the definition JSON).
38    pub supported_items: Identifier,
39    /// Narrower item tag for primary offers, when set.
40    pub primary_items: Option<Identifier>,
41    /// Enchantment tag this enchantment is incompatible with, when set.
42    pub exclusive_set: Option<Identifier>,
43    pub effects_nbt: fn() -> NbtCompound,
44    pub effects: EnchantmentEffects,
45}
46
47crate::impl_registry_entry_eq!(Enchantment);
48
49impl RegistryEntry for Enchantment {
50    fn key(&self) -> &Identifier {
51        &self.key
52    }
53
54    fn try_id(&self) -> Option<usize> {
55        REGISTRY.enchantments.id_from_key(&self.key)
56    }
57}
58
59impl ToNbtTag for &Enchantment {
60    fn to_nbt_tag(self) -> NbtTag {
61        let mut compound = NbtCompound::new();
62
63        // description: translatable text component {"translate": "enchantment.minecraft.<key>"}
64        let mut desc = NbtCompound::new();
65        desc.insert(
66            "translate",
67            format!("enchantment.{}.{}", self.key.namespace, self.key.path).as_str(),
68        );
69        compound.insert("description", NbtTag::Compound(desc));
70
71        // Definition fields (inlined, not nested)
72        compound.insert("supported_items", tag_ref(&self.supported_items));
73        if let Some(primary) = &self.primary_items {
74            compound.insert("primary_items", tag_ref(primary));
75        }
76        compound.insert("weight", self.weight as i32);
77        compound.insert("max_level", self.max_level as i32);
78
79        let mut min_cost = NbtCompound::new();
80        min_cost.insert("base", self.min_cost.base);
81        min_cost.insert("per_level_above_first", self.min_cost.per_level_above_first);
82        compound.insert("min_cost", NbtTag::Compound(min_cost));
83
84        let mut max_cost = NbtCompound::new();
85        max_cost.insert("base", self.max_cost.base);
86        max_cost.insert("per_level_above_first", self.max_cost.per_level_above_first);
87        compound.insert("max_cost", NbtTag::Compound(max_cost));
88
89        compound.insert("anvil_cost", self.anvil_cost);
90
91        let slots: Vec<String> = self.slots.iter().map(|s| s.as_str().to_owned()).collect();
92        compound.insert("slots", NbtTag::List(NbtList::from(slots)));
93
94        if let Some(exclusive) = &self.exclusive_set {
95            compound.insert("exclusive_set", tag_ref(exclusive));
96        }
97
98        let effects = (self.effects_nbt)();
99        if !effects.is_empty() {
100            compound.insert("effects", NbtTag::Compound(effects));
101        }
102
103        NbtTag::Compound(compound)
104    }
105}
106
107/// Formats a tag id the way enchantment definitions reference it: `#namespace:path`.
108fn tag_ref(tag: &Identifier) -> String {
109    format!("#{tag}")
110}
111
112impl Enchantment {
113    #[must_use]
114    pub fn matching_slot(&self, slot: EquipmentSlot) -> bool {
115        self.slots.iter().any(|group| group.test(slot))
116    }
117
118    /// Checks if this enchantment can be applied to the given item via `supported_items` tag.
119    pub fn can_enchant(&self, item: ItemRef) -> bool {
120        REGISTRY.items.is_in_tag(item, &self.supported_items)
121    }
122
123    /// Minimum enchanting cost for `level`.
124    #[must_use]
125    pub const fn get_min_cost(&self, level: i32) -> i32 {
126        self.min_cost.calculate(level)
127    }
128
129    /// Maximum enchanting cost for `level`.
130    #[must_use]
131    pub const fn get_max_cost(&self, level: i32) -> i32 {
132        self.max_cost.calculate(level)
133    }
134
135    /// Requires a supported item that also matches `primary_items` when set.
136    pub fn is_primary_item(&self, item: ItemRef) -> bool {
137        if !self.can_enchant(item) {
138            return false;
139        }
140        let Some(primary_items) = &self.primary_items else {
141            return true;
142        };
143        REGISTRY.items.is_in_tag(item, primary_items)
144    }
145
146    /// Checks if two enchantments are compatible (neither's `exclusive_set` contains the other).
147    #[must_use]
148    pub fn are_compatible(a: EnchantmentRef, b: EnchantmentRef) -> bool {
149        if a == b {
150            return false;
151        }
152        if let Some(set) = &a.exclusive_set
153            && REGISTRY.enchantments.is_in_tag(b, set)
154        {
155            return false;
156        }
157        if let Some(set) = &b.exclusive_set
158            && REGISTRY.enchantments.is_in_tag(a, set)
159        {
160            return false;
161        }
162        true
163    }
164
165    /// Checks if this enchantment is compatible with all existing enchantments on an item.
166    #[must_use]
167    pub fn is_compatible_with_existing(
168        enchantment: EnchantmentRef,
169        item: &crate::item_stack::ItemStack,
170    ) -> bool {
171        let Some(enchantments) = item.get_enchantments_for_crafting() else {
172            return true;
173        };
174        for (existing_key, _) in enchantments.iter() {
175            let Some(existing) = REGISTRY.enchantments.by_key(existing_key) else {
176                continue;
177            };
178            if !Self::are_compatible(enchantment, existing) {
179                return false;
180            }
181        }
182        true
183    }
184}
185
186pub type EnchantmentRef = &'static Enchantment;
187
188pub struct EnchantmentRegistry {
189    enchantments_by_id: Vec<EnchantmentRef>,
190    enchantments_by_key: FxHashMap<Identifier, usize>,
191    tags: RegistryTags,
192    allows_registering: bool,
193}
194
195impl EnchantmentRegistry {
196    #[must_use]
197    pub fn new() -> Self {
198        Self {
199            enchantments_by_id: Vec::new(),
200            enchantments_by_key: FxHashMap::default(),
201            tags: RegistryTags::default(),
202            allows_registering: true,
203        }
204    }
205}
206
207crate::impl_registry_ext!(
208    EnchantmentRegistry,
209    Enchantment,
210    enchantments_by_id,
211    enchantments_by_key
212);
213
214crate::impl_standard_methods!(
215    EnchantmentRegistry,
216    EnchantmentRef,
217    enchantments_by_id,
218    enchantments_by_key,
219    allows_registering
220);
221
222crate::impl_tagged_registry!(EnchantmentRegistry, enchantments_by_key, "enchantment");
223
224#[cfg(test)]
225mod tests {
226    use super::Enchantment;
227    use crate::enchantment_effect::{
228        DamageSourcePredicate, EnchantmentEffectComponent, EnchantmentEffectRequirements,
229        EnchantmentEntityEffect, EnchantmentTarget,
230    };
231    use crate::equipment::EquipmentSlot;
232    use crate::item_stack::ItemStack;
233    use crate::{init_vanilla_registry, vanilla_enchantments, vanilla_items};
234    use simdnbt::ToNbtTag;
235    use simdnbt::owned::{NbtList, NbtTag};
236    use steel_utils::Identifier;
237
238    #[test]
239    fn sharpness_is_primary_for_swords_but_only_supported_for_axes() {
240        init_vanilla_registry();
241        let sharpness = &vanilla_enchantments::SHARPNESS;
242        assert!(sharpness.can_enchant(&vanilla_items::DIAMOND_AXE));
243        assert!(!sharpness.is_primary_item(&vanilla_items::DIAMOND_AXE));
244        assert!(sharpness.is_primary_item(&vanilla_items::DIAMOND_SWORD));
245        assert!(!sharpness.is_primary_item(&vanilla_items::STONE));
246    }
247
248    #[test]
249    fn cost_windows_follow_the_vanilla_formula() {
250        init_vanilla_registry();
251        // Sharpness: min 1 + 11/level, max 21 + 11/level.
252        let sharpness = &vanilla_enchantments::SHARPNESS;
253        assert_eq!(sharpness.get_min_cost(1), 1);
254        assert_eq!(sharpness.get_min_cost(5), 45);
255        assert_eq!(sharpness.get_max_cost(3), 43);
256    }
257
258    #[test]
259    fn binding_curse_has_prevent_armor_change_effect() {
260        assert!(
261            vanilla_enchantments::BINDING_CURSE
262                .effects
263                .has(EnchantmentEffectComponent::PreventArmorChange)
264        );
265    }
266
267    #[test]
268    fn enchantment_matching_slot_uses_slot_groups() {
269        assert!(vanilla_enchantments::BINDING_CURSE.matching_slot(EquipmentSlot::Head));
270        assert!(!vanilla_enchantments::BINDING_CURSE.matching_slot(EquipmentSlot::MainHand));
271    }
272
273    #[test]
274    fn existing_identical_or_exclusive_enchantments_are_incompatible() {
275        init_vanilla_registry();
276        let mut sword = ItemStack::new(&vanilla_items::DIAMOND_SWORD);
277        sword.upgrade_enchantment(vanilla_enchantments::SHARPNESS.key.clone(), 1);
278
279        assert!(!Enchantment::is_compatible_with_existing(
280            &vanilla_enchantments::SHARPNESS,
281            &sword
282        ));
283        assert!(!Enchantment::is_compatible_with_existing(
284            &vanilla_enchantments::SMITE,
285            &sword
286        ));
287        assert!(Enchantment::is_compatible_with_existing(
288            &vanilla_enchantments::UNBREAKING,
289            &sword
290        ));
291    }
292
293    #[test]
294    fn compatibility_checks_stored_enchantments_on_enchanted_books() {
295        init_vanilla_registry();
296        let mut book = ItemStack::new(&vanilla_items::ENCHANTED_BOOK);
297        book.upgrade_enchantment(vanilla_enchantments::SHARPNESS.key.clone(), 1);
298
299        assert!(!Enchantment::is_compatible_with_existing(
300            &vanilla_enchantments::SMITE,
301            &book
302        ));
303        assert!(Enchantment::is_compatible_with_existing(
304            &vanilla_enchantments::UNBREAKING,
305            &book
306        ));
307    }
308
309    #[test]
310    fn enchantment_nbt_includes_raw_effect_payloads() {
311        let NbtTag::Compound(compound) = (&vanilla_enchantments::LUNGE).to_nbt_tag() else {
312            panic!("enchantment NBT should be a compound");
313        };
314        let Some(NbtTag::Compound(effects)) = compound.get("effects") else {
315            panic!("enchantment NBT should include effects");
316        };
317        let Some(NbtTag::List(NbtList::Compound(post_piercing))) =
318            effects.get("minecraft:post_piercing_attack")
319        else {
320            panic!("Lunge should include post-piercing attack effects");
321        };
322
323        assert_eq!(post_piercing.len(), 1);
324        let Some(NbtTag::Compound(effect)) = post_piercing[0].get("effect") else {
325            panic!("post-piercing entry should include an effect compound");
326        };
327        assert_eq!(
328            effect.string("type").map(ToString::to_string).as_deref(),
329            Some("minecraft:all_of")
330        );
331
332        let Some(NbtTag::List(NbtList::Compound(children))) = effect.get("effects") else {
333            panic!("Lunge all_of effect should include child effects");
334        };
335        let impulse = children
336            .iter()
337            .find(|child| {
338                child.string("type").map(ToString::to_string).as_deref()
339                    == Some("minecraft:apply_impulse")
340            })
341            .expect("Lunge should include apply_impulse");
342
343        assert!(matches!(
344            impulse.get("direction"),
345            Some(NbtTag::List(NbtList::Double(values))) if values == &[0.0, 0.0, 1.0]
346        ));
347        assert!(matches!(
348            impulse.get("coordinate_scale"),
349            Some(NbtTag::List(NbtList::Double(values))) if values == &[1.0, 0.0, 1.0]
350        ));
351        assert!(matches!(
352            impulse.get("magnitude").and_then(|tag| tag.compound()).and_then(|compound| compound.get("base")),
353            Some(NbtTag::Float(value)) if value.to_bits() == 0.458_f32.to_bits()
354        ));
355    }
356
357    #[test]
358    fn movement_requirements_use_vanilla_double_bounds_in_nbt() {
359        let NbtTag::Compound(compound) = (&vanilla_enchantments::WIND_BURST).to_nbt_tag() else {
360            panic!("enchantment NBT should be a compound");
361        };
362        let Some(NbtTag::Compound(effects)) = compound.get("effects") else {
363            panic!("enchantment NBT should include effects");
364        };
365        let Some(NbtTag::List(NbtList::Compound(post_attack))) =
366            effects.get("minecraft:post_attack")
367        else {
368            panic!("Wind Burst should include post-attack effects");
369        };
370        let Some(NbtTag::Compound(requirements)) = post_attack[0].get("requirements") else {
371            panic!("Wind Burst effect should include requirements");
372        };
373        let Some(NbtTag::Compound(predicate)) = requirements.get("predicate") else {
374            panic!("entity_properties requirements should include predicate");
375        };
376        let Some(NbtTag::Compound(movement)) = predicate.get("minecraft:movement") else {
377            panic!("Wind Burst predicate should include movement bounds");
378        };
379        let Some(NbtTag::Compound(fall_distance)) = movement.get("fall_distance") else {
380            panic!("movement predicate should include fall distance bounds");
381        };
382
383        assert!(matches!(
384            fall_distance.get("min"),
385            Some(NbtTag::Double(value)) if value.to_bits() == 1.5_f64.to_bits()
386        ));
387    }
388
389    #[test]
390    fn unconditional_value_effects_modify_values() {
391        assert_eq!(vanilla_enchantments::KNOCKBACK.effects.knockback.len(), 1);
392        let knockback = &vanilla_enchantments::KNOCKBACK.effects.knockback[0];
393
394        assert!(knockback.is_unconditional());
395        assert_eq!(
396            knockback
397                .effect
398                .process_without_random(2, 0.0)
399                .map(f32::to_bits),
400            Some(2.0_f32.to_bits())
401        );
402
403        assert_eq!(vanilla_enchantments::SHARPNESS.effects.damage.len(), 1);
404        let damage = &vanilla_enchantments::SHARPNESS.effects.damage[0];
405
406        assert!(damage.is_unconditional());
407        assert_eq!(
408            damage
409                .effect
410                .process_without_random(5, 7.0)
411                .map(f32::to_bits),
412            Some(10.0_f32.to_bits())
413        );
414    }
415
416    #[test]
417    fn conditional_value_effects_are_not_applied_without_context() {
418        assert_eq!(vanilla_enchantments::PUNCH.effects.knockback.len(), 1);
419        assert!(!vanilla_enchantments::PUNCH.effects.knockback[0].is_unconditional());
420    }
421
422    #[test]
423    fn looting_equipment_drops_preserves_attacker_target() {
424        let effects = vanilla_enchantments::LOOTING.effects.equipment_drops;
425
426        assert_eq!(effects.len(), 1);
427        assert_eq!(effects[0].enchanted, EnchantmentTarget::Attacker);
428        assert_eq!(effects[0].affected, EnchantmentTarget::Victim);
429    }
430
431    #[test]
432    fn frost_walker_damage_immunity_preserves_damage_source_requirements() {
433        assert!(
434            vanilla_enchantments::FROST_WALKER
435                .effects
436                .has(EnchantmentEffectComponent::DamageImmunity)
437        );
438
439        let effects = vanilla_enchantments::FROST_WALKER.effects.damage_immunity;
440        assert_eq!(effects.len(), 1);
441        let Some(requirements) = effects[0].requirements else {
442            panic!("Frost Walker damage immunity should have requirements");
443        };
444        let EnchantmentEffectRequirements::DamageSourceProperties(DamageSourcePredicate {
445            tags,
446            is_direct,
447        }) = requirements
448        else {
449            panic!("Frost Walker damage immunity should use damage-source requirements");
450        };
451
452        assert_eq!(*is_direct, None);
453        assert!(tags.iter().any(|tag| {
454            tag.tag == Identifier::vanilla_static("burn_from_stepping") && tag.expected
455        }));
456        assert!(tags.iter().any(|tag| {
457            tag.tag == Identifier::vanilla_static("bypasses_invulnerability") && !tag.expected
458        }));
459    }
460
461    #[test]
462    fn lunge_post_piercing_attack_preserves_entity_effects() {
463        assert!(
464            vanilla_enchantments::LUNGE
465                .effects
466                .has(EnchantmentEffectComponent::PostPiercingAttack)
467        );
468
469        let effects = vanilla_enchantments::LUNGE.effects.post_piercing_attack;
470        assert_eq!(effects.len(), 1);
471        let EnchantmentEntityEffect::AllOf(children) = &effects[0].effect else {
472            panic!("Lunge post-piercing effect should be an all_of entity effect");
473        };
474
475        assert_eq!(children.len(), 4);
476        assert!(
477            children
478                .iter()
479                .any(|effect| matches!(effect, EnchantmentEntityEffect::ChangeItemDamage { .. }))
480        );
481        assert!(
482            children
483                .iter()
484                .any(|effect| matches!(effect, EnchantmentEntityEffect::ApplyExhaustion { .. }))
485        );
486        assert!(
487            children
488                .iter()
489                .any(|effect| matches!(effect, EnchantmentEntityEffect::ApplyImpulse { .. }))
490        );
491        assert!(
492            children
493                .iter()
494                .any(|effect| matches!(effect, EnchantmentEntityEffect::PlaySound { .. }))
495        );
496        assert!(effects[0].requirements.is_some());
497    }
498
499    #[test]
500    fn thorns_post_attack_preserves_random_damage_and_item_damage_effects() {
501        assert!(
502            vanilla_enchantments::THORNS
503                .effects
504                .has(EnchantmentEffectComponent::PostAttack)
505        );
506
507        let effects = vanilla_enchantments::THORNS.effects.post_attack;
508        assert!(effects.iter().any(|effect| {
509            effect
510                .requirements
511                .is_some_and(requirements_contain_random_chance)
512        }));
513        assert!(
514            effects
515                .iter()
516                .any(|effect| entity_effect_contains_damage_entity(&effect.effect))
517        );
518        assert!(
519            effects
520                .iter()
521                .any(|effect| entity_effect_contains_change_item_damage(&effect.effect))
522        );
523    }
524
525    fn requirements_contain_random_chance(requirements: &EnchantmentEffectRequirements) -> bool {
526        match requirements {
527            EnchantmentEffectRequirements::AllOf(children)
528            | EnchantmentEffectRequirements::AnyOf(children) => children
529                .iter()
530                .any(|child| requirements_contain_random_chance(child)),
531            EnchantmentEffectRequirements::Inverted(child) => {
532                requirements_contain_random_chance(child)
533            }
534            EnchantmentEffectRequirements::RandomChance { .. } => true,
535            EnchantmentEffectRequirements::EntityProperties { .. }
536            | EnchantmentEffectRequirements::DamageSourceProperties(_)
537            | EnchantmentEffectRequirements::MatchTool { .. }
538            | EnchantmentEffectRequirements::Unsupported { .. } => false,
539        }
540    }
541
542    fn entity_effect_contains_damage_entity(effect: &EnchantmentEntityEffect) -> bool {
543        match effect {
544            EnchantmentEntityEffect::AllOf(children) => children
545                .iter()
546                .any(|child| entity_effect_contains_damage_entity(child)),
547            EnchantmentEntityEffect::DamageEntity { .. } => true,
548            EnchantmentEntityEffect::ChangeItemDamage { .. }
549            | EnchantmentEntityEffect::ApplyExhaustion { .. }
550            | EnchantmentEntityEffect::ApplyImpulse { .. }
551            | EnchantmentEntityEffect::PlaySound { .. }
552            | EnchantmentEntityEffect::Ignite { .. }
553            | EnchantmentEntityEffect::ApplyMobEffect { .. }
554            | EnchantmentEntityEffect::Unsupported { .. } => false,
555        }
556    }
557
558    fn entity_effect_contains_change_item_damage(effect: &EnchantmentEntityEffect) -> bool {
559        match effect {
560            EnchantmentEntityEffect::AllOf(children) => children
561                .iter()
562                .any(|child| entity_effect_contains_change_item_damage(child)),
563            EnchantmentEntityEffect::ChangeItemDamage { .. } => true,
564            EnchantmentEntityEffect::ApplyExhaustion { .. }
565            | EnchantmentEntityEffect::ApplyImpulse { .. }
566            | EnchantmentEntityEffect::PlaySound { .. }
567            | EnchantmentEntityEffect::DamageEntity { .. }
568            | EnchantmentEntityEffect::Ignite { .. }
569            | EnchantmentEntityEffect::ApplyMobEffect { .. }
570            | EnchantmentEntityEffect::Unsupported { .. } => false,
571        }
572    }
573}
steel_registry/enchantment/mod.rs

steel_registry/enchantment/
mod.rs
