steel_registry/data_components/components/
enchantments.rs

1use rustc_hash::FxHashMap;
2use simdnbt::owned::{NbtCompound, NbtTag};
3use simdnbt::{FromNbtTag, ToNbtTag};
4use steel_utils::Identifier;
5use steel_utils::codec::VarInt;
6use steel_utils::hash::{ComponentHasher, HashComponent, HashEntry, sort_map_entries};
7use steel_utils::nbt::NbtNumeric as _;
8use steel_utils::serial::{ReadFrom, WriteTo};
9
10use crate::{REGISTRY, RegistryExt};
11
12/// Enchantments stored on an item. Maps enchantment key to level.
13///
14/// Used by both the `minecraft:enchantments` component (on enchanted items)
15/// and the `minecraft:stored_enchantments` component (on enchanted books).
16///
17/// Vanilla moved tooltip visibility to the separate `TOOLTIP_DISPLAY` component.
18#[derive(Debug, Clone, PartialEq)]
19pub struct ItemEnchantments {
20    pub levels: FxHashMap<Identifier, u32>,
21}
22
23impl ItemEnchantments {
24    #[must_use]
25    pub fn empty() -> Self {
26        Self {
27            levels: FxHashMap::default(),
28        }
29    }
30
31    #[must_use]
32    pub fn get_level(&self, enchantment: &Identifier) -> u32 {
33        self.levels.get(enchantment).copied().unwrap_or(0)
34    }
35
36    pub fn set(&mut self, enchantment: Identifier, level: u32) {
37        if level == 0 {
38            self.levels.remove(&enchantment);
39        } else {
40            self.levels.insert(enchantment, level.min(255));
41        }
42    }
43
44    /// Vanilla `Mutable.upgrade`: keeps the higher of existing vs new level.
45    pub fn upgrade(&mut self, enchantment: Identifier, level: u32) {
46        if level > 0 {
47            let existing = self.get_level(&enchantment);
48            self.levels
49                .insert(enchantment, existing.max(level).min(255));
50        }
51    }
52
53    #[must_use]
54    pub fn is_empty(&self) -> bool {
55        self.levels.is_empty()
56    }
57
58    #[must_use]
59    pub fn len(&self) -> usize {
60        self.levels.len()
61    }
62
63    pub fn iter(&self) -> impl Iterator<Item = (&Identifier, &u32)> {
64        self.levels.iter()
65    }
66}
67
68impl IntoIterator for ItemEnchantments {
69    type Item = (Identifier, u32);
70
71    type IntoIter = <FxHashMap<Identifier, u32> as IntoIterator>::IntoIter;
72
73    fn into_iter(self) -> Self::IntoIter {
74        self.levels.into_iter()
75    }
76}
77
78impl Default for ItemEnchantments {
79    fn default() -> Self {
80        Self::empty()
81    }
82}
83
84/// Network format: `VarInt` count, then (`VarInt` `enchantment_id`, `VarInt` level) pairs.
85impl WriteTo for ItemEnchantments {
86    fn write(&self, writer: &mut impl std::io::Write) -> std::io::Result<()> {
87        let count = i32::try_from(self.levels.len()).map_err(|_| {
88            std::io::Error::other(format!(
89                "Enchantment map too large: {} entries",
90                self.levels.len()
91            ))
92        })?;
93        VarInt(count).write(writer)?;
94        for (key, &level) in &self.levels {
95            let id = REGISTRY
96                .enchantments
97                .id_from_key(key)
98                .ok_or_else(|| std::io::Error::other(format!("Unknown enchantment: {key}")))?;
99            let id = i32::try_from(id).map_err(|_| {
100                std::io::Error::other(format!("Enchantment id out of protocol range: {id}"))
101            })?;
102            if level > 255 {
103                return Err(std::io::Error::other(format!(
104                    "Enchantment {key} has invalid level {level}"
105                )));
106            }
107            VarInt(id).write(writer)?;
108            VarInt(level as i32).write(writer)?;
109        }
110        Ok(())
111    }
112}
113
114impl ReadFrom for ItemEnchantments {
115    fn read(data: &mut std::io::Cursor<&[u8]>) -> std::io::Result<Self> {
116        let count = VarInt::read(data)?.0;
117        let count = usize::try_from(count)
118            .map_err(|_| std::io::Error::other(format!("Negative enchantment count: {count}")))?;
119        let mut levels = FxHashMap::default();
120        levels.reserve(count.min(65_536));
121        for _ in 0..count {
122            let id = VarInt::read(data)?.0;
123            let id = usize::try_from(id)
124                .map_err(|_| std::io::Error::other(format!("Negative enchantment id: {id}")))?;
125            let level = VarInt::read(data)?.0;
126            if !(0..=255).contains(&level) {
127                return Err(std::io::Error::other(format!(
128                    "Enchantment level out of range: {level}"
129                )));
130            }
131            let enchantment = REGISTRY
132                .enchantments
133                .by_id(id)
134                .ok_or_else(|| std::io::Error::other(format!("Unknown enchantment id: {id}")))?;
135            levels.insert(enchantment.key.clone(), level as u32);
136        }
137        Ok(Self { levels })
138    }
139}
140
141/// NBT format: compound with enchantment identifiers as keys and int levels as values.
142impl ToNbtTag for ItemEnchantments {
143    fn to_nbt_tag(self) -> NbtTag {
144        let mut compound = NbtCompound::new();
145        for (key, level) in &self.levels {
146            compound.insert(key.to_string(), NbtTag::Int(*level as i32));
147        }
148        NbtTag::Compound(compound)
149    }
150}
151
152impl FromNbtTag for ItemEnchantments {
153    fn from_nbt_tag(tag: simdnbt::borrow::NbtTag) -> Option<Self> {
154        let compound = tag.compound()?;
155        let mut levels = FxHashMap::default();
156        for (key, value) in compound.iter() {
157            let ident = key.to_str().parse::<Identifier>().ok()?;
158            REGISTRY.enchantments.by_key(&ident)?;
159            let level = value.codec_i32()?;
160            if !(1..=255).contains(&level) {
161                return None;
162            }
163            if levels.insert(ident, level as u32).is_some() {
164                return None;
165            }
166        }
167        Some(Self { levels })
168    }
169}
170
171impl HashComponent for ItemEnchantments {
172    fn hash_component(&self, hasher: &mut ComponentHasher) {
173        hasher.start_map();
174        let mut entries: Vec<_> = self
175            .levels
176            .iter()
177            .map(|(key, &level)| {
178                let mut key_hasher = ComponentHasher::new();
179                key_hasher.put_string(&key.to_string());
180                let mut value_hasher = ComponentHasher::new();
181                value_hasher.put_int(level as i32);
182                HashEntry::new(key_hasher, value_hasher)
183            })
184            .collect();
185        sort_map_entries(&mut entries);
186        for entry in &entries {
187            hasher.put_raw_bytes(&entry.key_bytes);
188            hasher.put_raw_bytes(&entry.value_bytes);
189        }
190        hasher.end_map();
191    }
192}
193
194#[cfg(test)]
195mod tests {
196    use std::io::Cursor;
197
198    use simdnbt::FromNbtTag;
199    use simdnbt::borrow::{NbtTag as BorrowedNbtTag, read_tag};
200    use simdnbt::owned::{NbtCompound, NbtTag};
201
202    use super::ItemEnchantments;
203    use crate::init_vanilla_registry;
204
205    fn with_borrowed_tag<R>(tag: NbtTag, visitor: impl FnOnce(BorrowedNbtTag<'_, '_>) -> R) -> R {
206        let mut bytes = Vec::new();
207        tag.write(&mut bytes);
208        let borrowed =
209            read_tag(&mut Cursor::new(bytes.as_slice())).expect("owned test tag should parse");
210        visitor(borrowed.as_tag())
211    }
212
213    fn parse_enchantments(compound: NbtCompound) -> Option<ItemEnchantments> {
214        with_borrowed_tag(NbtTag::Compound(compound), ItemEnchantments::from_nbt_tag)
215    }
216
217    #[test]
218    fn enchantment_nbt_requires_known_keys_and_vanilla_levels() {
219        init_vanilla_registry();
220        let mut valid = NbtCompound::new();
221        valid.insert("minecraft:efficiency", 5_i8);
222        assert!(parse_enchantments(valid).is_some());
223
224        let mut unknown = NbtCompound::new();
225        unknown.insert("minecraft:not_an_enchantment", 1);
226        assert!(parse_enchantments(unknown).is_none());
227
228        let mut out_of_range = NbtCompound::new();
229        out_of_range.insert("minecraft:efficiency", 256);
230        assert!(parse_enchantments(out_of_range).is_none());
231    }
232}
steel_registry/data_components/components/enchantments.rs

steel_registry/data_components/components/
enchantments.rs
