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

steel_registry/enchantment/
mod.rs
