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::serial::{ReadFrom, WriteTo};
8
9use crate::{REGISTRY, RegistryExt};
10
11/// Enchantments stored on an item. Maps enchantment key to level.
12///
13/// Used by both the `minecraft:enchantments` component (on enchanted items)
14/// and the `minecraft:stored_enchantments` component (on enchanted books).
15///
16/// Vanilla moved tooltip visibility to the separate `TOOLTIP_DISPLAY` component.
17#[derive(Debug, Clone, PartialEq)]
18pub struct ItemEnchantments {
19    pub levels: FxHashMap<Identifier, u32>,
20}
21
22impl ItemEnchantments {
23    #[must_use]
24    pub fn empty() -> Self {
25        Self {
26            levels: FxHashMap::default(),
27        }
28    }
29
30    #[must_use]
31    pub fn get_level(&self, enchantment: &Identifier) -> u32 {
32        self.levels.get(enchantment).copied().unwrap_or(0)
33    }
34
35    pub fn set(&mut self, enchantment: Identifier, level: u32) {
36        if level == 0 {
37            self.levels.remove(&enchantment);
38        } else {
39            self.levels.insert(enchantment, level);
40        }
41    }
42
43    /// Vanilla `Mutable.upgrade`: keeps the higher of existing vs new level.
44    pub fn upgrade(&mut self, enchantment: Identifier, level: u32) {
45        if level > 0 {
46            let existing = self.get_level(&enchantment);
47            self.levels
48                .insert(enchantment, existing.max(level).min(255));
49        }
50    }
51
52    #[must_use]
53    pub fn is_empty(&self) -> bool {
54        self.levels.is_empty()
55    }
56
57    #[must_use]
58    pub fn len(&self) -> usize {
59        self.levels.len()
60    }
61
62    pub fn iter(&self) -> impl Iterator<Item = (&Identifier, &u32)> {
63        self.levels.iter()
64    }
65}
66
67impl Default for ItemEnchantments {
68    fn default() -> Self {
69        Self::empty()
70    }
71}
72
73/// Network format: VarInt count, then (VarInt enchantment_id, VarInt level) pairs.
74impl WriteTo for ItemEnchantments {
75    fn write(&self, writer: &mut impl std::io::Write) -> std::io::Result<()> {
76        VarInt(self.levels.len() as i32).write(writer)?;
77        for (key, &level) in &self.levels {
78            let id = REGISTRY
79                .enchantments
80                .id_from_key(key)
81                .ok_or_else(|| std::io::Error::other(format!("Unknown enchantment: {key}")))?;
82            VarInt(id as i32).write(writer)?;
83            VarInt(level as i32).write(writer)?;
84        }
85        Ok(())
86    }
87}
88
89impl ReadFrom for ItemEnchantments {
90    fn read(data: &mut std::io::Cursor<&[u8]>) -> std::io::Result<Self> {
91        let count = VarInt::read(data)?.0;
92        if !(0..=256).contains(&count) {
93            return Err(std::io::Error::other(format!(
94                "Enchantment count out of range: {count}"
95            )));
96        }
97        let count = count as usize;
98        let mut levels = FxHashMap::default();
99        for _ in 0..count {
100            let id = VarInt::read(data)?.0 as usize;
101            let level = VarInt::read(data)?.0 as u32;
102            let enchantment = REGISTRY
103                .enchantments
104                .by_id(id)
105                .ok_or_else(|| std::io::Error::other(format!("Unknown enchantment id: {id}")))?;
106            levels.insert(enchantment.key.clone(), level);
107        }
108        Ok(Self { levels })
109    }
110}
111
112/// NBT format: compound with enchantment identifiers as keys and int levels as values.
113impl ToNbtTag for ItemEnchantments {
114    fn to_nbt_tag(self) -> NbtTag {
115        let mut compound = NbtCompound::new();
116        for (key, level) in &self.levels {
117            compound.insert(key.to_string(), NbtTag::Int(*level as i32));
118        }
119        NbtTag::Compound(compound)
120    }
121}
122
123impl FromNbtTag for ItemEnchantments {
124    fn from_nbt_tag(tag: simdnbt::borrow::NbtTag) -> Option<Self> {
125        let compound = tag.compound()?;
126        let mut levels = FxHashMap::default();
127        for (key, value) in compound.iter() {
128            let key_str = key.to_str();
129            if let Ok(ident) = key_str.parse::<Identifier>()
130                && let Some(level) = value.int()
131                && level > 0
132            {
133                levels.insert(ident, level as u32);
134            }
135        }
136        Some(Self { levels })
137    }
138}
139
140impl HashComponent for ItemEnchantments {
141    fn hash_component(&self, hasher: &mut ComponentHasher) {
142        hasher.start_map();
143        let mut entries: Vec<_> = self
144            .levels
145            .iter()
146            .map(|(key, &level)| {
147                let mut key_hasher = ComponentHasher::new();
148                key_hasher.put_string(&key.to_string());
149                let mut value_hasher = ComponentHasher::new();
150                value_hasher.put_int(level as i32);
151                HashEntry::new(key_hasher, value_hasher)
152            })
153            .collect();
154        sort_map_entries(&mut entries);
155        for entry in &entries {
156            hasher.put_raw_bytes(&entry.key_bytes);
157            hasher.put_raw_bytes(&entry.value_bytes);
158        }
159        hasher.end_map();
160    }
161}
steel_registry/data_components/components/enchantments.rs

steel_registry/data_components/components/
enchantments.rs
