Skip to main content

steel_registry/registry/
holder_set.rs

1//! Registry-backed holder sets used by vanilla codecs.
2
3use std::fmt::Debug;
4use std::io::{Cursor, Error, Result, Write};
5use std::str::FromStr;
6
7use simdnbt::owned::{NbtList, NbtTag};
8use simdnbt::{FromNbtTag, ToNbtTag};
9use steel_utils::{
10    Identifier,
11    codec::VarInt,
12    hash::{ComponentHasher, HashComponent},
13    serial::{ReadFrom, WriteTo},
14};
15
16use crate::attribute::Attribute;
17use crate::banner_pattern::BannerPattern;
18use crate::blocks::Block;
19use crate::damage_type::DamageType;
20use crate::enchantment::Enchantment;
21use crate::entity_type::EntityType;
22use crate::items::Item;
23use crate::jukebox_song::JukeboxSong;
24use crate::mob_effect::MobEffect;
25use crate::potion::Potion;
26use crate::trim_material::TrimMaterial;
27use crate::trim_pattern::TrimPattern;
28use crate::villager_type::VillagerType;
29use crate::{REGISTRY, RegistryEntry, RegistryExt, TaggedRegistryExt};
30
31/// Registry operations required by a [`RegistryHolderSet`].
32///
33/// The trait keeps the holder-set codec independent from Steel's concrete
34/// registries. Plugin-owned entry types can provide the same operations through
35/// their own registry implementation.
36pub trait RegistryHolderSetEntry: RegistryEntry + Debug + Send + Sync {
37    /// Human-readable registry name used in codec errors.
38    const REGISTRY_NAME: &'static str;
39
40    /// Looks up an entry by its protocol registry ID.
41    fn holder_set_by_id(id: usize) -> Option<&'static Self>;
42
43    /// Looks up an entry by its registry key.
44    fn holder_set_by_key(key: &Identifier) -> Option<&'static Self>;
45
46    /// Returns whether the registry contains this tag.
47    fn holder_set_tag_exists(tag: &Identifier) -> bool;
48
49    /// Returns whether an entry belongs to this tag.
50    fn holder_set_tag_contains(entry: &'static Self, tag: &Identifier) -> bool;
51}
52
53/// Vanilla's homogeneous holder-set representation for a registry.
54#[derive(Debug, PartialEq)]
55pub enum RegistryHolderSet<T: RegistryHolderSetEntry> {
56    /// A named registry tag.
57    Tag(Identifier),
58    /// An ordered list of direct registry references.
59    Direct(Vec<&'static T>),
60}
61
62impl<T: RegistryHolderSetEntry> Clone for RegistryHolderSet<T> {
63    fn clone(&self) -> Self {
64        match self {
65            Self::Tag(tag) => Self::Tag(tag.clone()),
66            Self::Direct(entries) => Self::Direct(entries.clone()),
67        }
68    }
69}
70
71impl<T: RegistryHolderSetEntry> RegistryHolderSet<T> {
72    /// Returns whether this holder set contains `entry`.
73    #[must_use]
74    pub fn contains(&self, entry: &'static T) -> bool {
75        match self {
76            Self::Tag(tag) => T::holder_set_tag_contains(entry, tag),
77            Self::Direct(entries) => entries.contains(&entry),
78        }
79    }
80
81    pub(crate) fn from_owned_nbt(tag: &NbtTag) -> Option<Self> {
82        if let Some(value) = tag.string() {
83            let value = value.to_string();
84            if let Some(tag) = value.strip_prefix('#') {
85                let tag = Identifier::from_str(tag).ok()?;
86                if !T::holder_set_tag_exists(&tag) {
87                    return None;
88                }
89                return Some(Self::Tag(tag));
90            }
91
92            let key = Identifier::from_str(&value).ok()?;
93            return Some(Self::Direct(vec![T::holder_set_by_key(&key)?]));
94        }
95
96        let list = tag.list()?;
97        if list.as_nbt_tags().is_empty() {
98            return Some(Self::Direct(Vec::new()));
99        }
100        let values = list.strings()?;
101        let mut entries = Vec::with_capacity(values.len());
102        for value in values {
103            let key = Identifier::from_str(&value.to_string()).ok()?;
104            entries.push(T::holder_set_by_key(&key)?);
105        }
106        Some(Self::Direct(entries))
107    }
108}
109
110impl<T: RegistryHolderSetEntry> WriteTo for RegistryHolderSet<T> {
111    fn write(&self, writer: &mut impl Write) -> Result<()> {
112        match self {
113            Self::Tag(tag) => {
114                if !T::holder_set_tag_exists(tag) {
115                    return Err(Error::other(format!(
116                        "Unknown {} tag: {tag}",
117                        T::REGISTRY_NAME
118                    )));
119                }
120                VarInt(0).write(writer)?;
121                tag.write(writer)
122            }
123            Self::Direct(entries) => {
124                let count = i32::try_from(entries.len()).map_err(|_| {
125                    Error::other(format!(
126                        "{} holder set too large: {}",
127                        T::REGISTRY_NAME,
128                        entries.len()
129                    ))
130                })?;
131                let encoded_count = count.checked_add(1).ok_or_else(|| {
132                    Error::other(format!(
133                        "{} holder set count exceeds protocol range",
134                        T::REGISTRY_NAME
135                    ))
136                })?;
137                VarInt(encoded_count).write(writer)?;
138                for entry in entries {
139                    let id = entry.try_id().ok_or_else(|| {
140                        Error::other(format!("Unknown {}: {}", T::REGISTRY_NAME, entry.key()))
141                    })?;
142                    let id = i32::try_from(id).map_err(|_| {
143                        Error::other(format!(
144                            "{} id out of protocol range: {id}",
145                            T::REGISTRY_NAME
146                        ))
147                    })?;
148                    VarInt(id).write(writer)?;
149                }
150                Ok(())
151            }
152        }
153    }
154}
155
156impl<T: RegistryHolderSetEntry> ReadFrom for RegistryHolderSet<T> {
157    fn read(data: &mut Cursor<&[u8]>) -> Result<Self> {
158        let encoded_count = VarInt::read(data)?.0;
159        if encoded_count == 0 {
160            let tag = Identifier::read(data)?;
161            if !T::holder_set_tag_exists(&tag) {
162                return Err(Error::other(format!(
163                    "Unknown {} tag: {tag}",
164                    T::REGISTRY_NAME
165                )));
166            }
167            return Ok(Self::Tag(tag));
168        }
169
170        let count = encoded_count
171            .checked_sub(1)
172            .and_then(|count| usize::try_from(count).ok())
173            .ok_or_else(|| {
174                Error::other(format!(
175                    "Invalid {} holder set count: {encoded_count}",
176                    T::REGISTRY_NAME
177                ))
178            })?;
179        let mut entries = Vec::with_capacity(count.min(65_536));
180        for _ in 0..count {
181            let id = VarInt::read(data)?.0;
182            let id = usize::try_from(id)
183                .map_err(|_| Error::other(format!("Negative {} id: {id}", T::REGISTRY_NAME)))?;
184            let entry = T::holder_set_by_id(id)
185                .ok_or_else(|| Error::other(format!("Unknown {} id: {id}", T::REGISTRY_NAME)))?;
186            entries.push(entry);
187        }
188        Ok(Self::Direct(entries))
189    }
190}
191
192impl<T: RegistryHolderSetEntry> ToNbtTag for RegistryHolderSet<T> {
193    fn to_nbt_tag(self) -> NbtTag {
194        match self {
195            Self::Tag(tag) => NbtTag::String(format!("#{tag}").into()),
196            Self::Direct(entries) if entries.is_empty() => NbtTag::List(NbtList::Empty),
197            Self::Direct(entries) if entries.len() == 1 => {
198                NbtTag::String(entries[0].key().to_string().into())
199            }
200            Self::Direct(entries) => NbtTag::List(NbtList::String(
201                entries
202                    .into_iter()
203                    .map(|entry| entry.key().to_string().into())
204                    .collect(),
205            )),
206        }
207    }
208}
209
210impl<T: RegistryHolderSetEntry> FromNbtTag for RegistryHolderSet<T> {
211    fn from_nbt_tag(tag: simdnbt::borrow::NbtTag) -> Option<Self> {
212        Self::from_owned_nbt(&tag.to_owned())
213    }
214}
215
216impl<T: RegistryHolderSetEntry> HashComponent for RegistryHolderSet<T> {
217    fn hash_component(&self, hasher: &mut ComponentHasher) {
218        match self {
219            Self::Tag(tag) => hasher.put_string(&format!("#{tag}")),
220            Self::Direct(entries) if entries.len() == 1 => {
221                hasher.put_string(&entries[0].key().to_string());
222            }
223            Self::Direct(entries) => {
224                hasher.start_list();
225                for entry in entries {
226                    hasher.put_component_hash(&entry.key().to_string());
227                }
228                hasher.end_list();
229            }
230        }
231    }
232}
233
234macro_rules! impl_registry_holder_set_entry {
235    ($entry:ty, $registry:ident, $name:literal) => {
236        impl RegistryHolderSetEntry for $entry {
237            const REGISTRY_NAME: &'static str = $name;
238
239            fn holder_set_by_id(id: usize) -> Option<&'static Self> {
240                REGISTRY.$registry.by_id(id)
241            }
242
243            fn holder_set_by_key(key: &Identifier) -> Option<&'static Self> {
244                REGISTRY.$registry.by_key(key)
245            }
246
247            fn holder_set_tag_exists(tag: &Identifier) -> bool {
248                REGISTRY.$registry.get_tag(tag).is_some()
249            }
250
251            fn holder_set_tag_contains(entry: &'static Self, tag: &Identifier) -> bool {
252                REGISTRY.$registry.is_in_tag(entry, tag)
253            }
254        }
255    };
256}
257
258impl_registry_holder_set_entry!(Block, blocks, "block");
259impl_registry_holder_set_entry!(BannerPattern, banner_patterns, "banner pattern");
260impl_registry_holder_set_entry!(EntityType, entity_types, "entity type");
261impl_registry_holder_set_entry!(Item, items, "item");
262impl_registry_holder_set_entry!(DamageType, damage_types, "damage type");
263impl_registry_holder_set_entry!(MobEffect, mob_effects, "mob effect");
264impl_registry_holder_set_entry!(Enchantment, enchantments, "enchantment");
265impl_registry_holder_set_entry!(Potion, potions, "potion");
266impl_registry_holder_set_entry!(Attribute, attributes, "attribute");
267impl_registry_holder_set_entry!(TrimMaterial, trim_materials, "trim material");
268impl_registry_holder_set_entry!(TrimPattern, trim_patterns, "trim pattern");
269impl_registry_holder_set_entry!(JukeboxSong, jukebox_songs, "jukebox song");
270impl_registry_holder_set_entry!(VillagerType, villager_types, "villager type");
271
272#[cfg(test)]
273mod tests {
274    use std::io::Cursor;
275
276    use simdnbt::borrow::{NbtTag as BorrowedNbtTag, read_tag};
277    use simdnbt::owned::{NbtList, NbtTag};
278    use simdnbt::{FromNbtTag, ToNbtTag};
279    use steel_utils::Identifier;
280    use steel_utils::codec::VarInt;
281    use steel_utils::hash::HashComponent as _;
282    use steel_utils::serial::{ReadFrom, WriteTo};
283
284    use super::RegistryHolderSet;
285    use crate::init_vanilla_registry;
286    use crate::items::Item;
287    use crate::vanilla_item_tags::ItemTag;
288    use crate::vanilla_items;
289
290    fn with_borrowed_tag<R>(tag: NbtTag, visitor: impl FnOnce(BorrowedNbtTag<'_, '_>) -> R) -> R {
291        let mut bytes = Vec::new();
292        tag.write(&mut bytes);
293        let borrowed =
294            read_tag(&mut Cursor::new(bytes.as_slice())).expect("owned test tag should parse");
295        visitor(borrowed.as_tag())
296    }
297
298    #[test]
299    fn persistent_codec_matches_vanilla_holder_set_shapes() {
300        init_vanilla_registry();
301
302        let tag = RegistryHolderSet::<Item>::Tag(ItemTag::WOOL);
303        assert_eq!(
304            tag.clone().to_nbt_tag(),
305            NbtTag::String("#minecraft:wool".into())
306        );
307        assert_eq!(
308            with_borrowed_tag(
309                tag.clone().to_nbt_tag(),
310                RegistryHolderSet::<Item>::from_nbt_tag
311            ),
312            Some(tag)
313        );
314        assert_eq!(
315            RegistryHolderSet::<Item>::Tag(ItemTag::WOOL).compute_hash(),
316            NbtTag::String("#minecraft:wool".into()).compute_hash()
317        );
318
319        let singleton: RegistryHolderSet<Item> =
320            RegistryHolderSet::Direct(vec![&vanilla_items::STICK]);
321        assert_eq!(
322            singleton.clone().to_nbt_tag(),
323            NbtTag::String("minecraft:stick".into())
324        );
325        assert_eq!(
326            with_borrowed_tag(
327                singleton.clone().to_nbt_tag(),
328                RegistryHolderSet::<Item>::from_nbt_tag
329            ),
330            Some(singleton)
331        );
332
333        let direct: RegistryHolderSet<Item> =
334            RegistryHolderSet::Direct(vec![&vanilla_items::STICK, &vanilla_items::DIAMOND]);
335        assert_eq!(
336            with_borrowed_tag(
337                direct.clone().to_nbt_tag(),
338                RegistryHolderSet::<Item>::from_nbt_tag
339            ),
340            Some(direct.clone())
341        );
342        assert_eq!(
343            direct.compute_hash(),
344            direct.clone().to_nbt_tag().compute_hash()
345        );
346
347        let empty = RegistryHolderSet::<Item>::Direct(Vec::new());
348        assert_eq!(empty.clone().to_nbt_tag(), NbtTag::List(NbtList::Empty));
349        assert_eq!(
350            with_borrowed_tag(
351                NbtTag::List(NbtList::Empty),
352                RegistryHolderSet::<Item>::from_nbt_tag
353            ),
354            Some(empty)
355        );
356    }
357
358    #[test]
359    fn network_codec_round_trips_tag_direct_and_empty_sets() {
360        init_vanilla_registry();
361
362        for holder_set in [
363            RegistryHolderSet::<Item>::Tag(ItemTag::WOOL),
364            RegistryHolderSet::Direct(vec![&vanilla_items::STICK, &vanilla_items::DIAMOND]),
365            RegistryHolderSet::Direct(Vec::new()),
366        ] {
367            let mut bytes = Vec::new();
368            holder_set
369                .write(&mut bytes)
370                .expect("holder set should write");
371            assert_eq!(
372                RegistryHolderSet::<Item>::read(&mut Cursor::new(bytes.as_slice()))
373                    .expect("holder set should read"),
374                holder_set
375            );
376        }
377    }
378
379    #[test]
380    fn codecs_reject_unknown_registry_values() {
381        init_vanilla_registry();
382
383        let unknown_tag = NbtTag::String("#steel:missing".into());
384        assert_eq!(
385            with_borrowed_tag(unknown_tag, RegistryHolderSet::<Item>::from_nbt_tag),
386            None
387        );
388        let unknown_entry = NbtTag::String("steel:missing".into());
389        assert_eq!(
390            with_borrowed_tag(unknown_entry, RegistryHolderSet::<Item>::from_nbt_tag),
391            None
392        );
393
394        let mut invalid_count = Vec::new();
395        VarInt(-1)
396            .write(&mut invalid_count)
397            .expect("test count should write");
398        assert!(
399            RegistryHolderSet::<Item>::read(&mut Cursor::new(invalid_count.as_slice())).is_err()
400        );
401
402        let mut unknown_id = Vec::new();
403        VarInt(2)
404            .write(&mut unknown_id)
405            .expect("test count should write");
406        VarInt(i32::MAX)
407            .write(&mut unknown_id)
408            .expect("test id should write");
409        assert!(RegistryHolderSet::<Item>::read(&mut Cursor::new(unknown_id.as_slice())).is_err());
410    }
411
412    #[test]
413    fn contains_resolves_tags_and_direct_entries() {
414        init_vanilla_registry();
415
416        let tag = RegistryHolderSet::<Item>::Tag(ItemTag::WOOL);
417        assert!(tag.contains(&vanilla_items::WHITE_WOOL));
418        assert!(!tag.contains(&vanilla_items::STICK));
419
420        let direct: RegistryHolderSet<Item> =
421            RegistryHolderSet::Direct(vec![&vanilla_items::STICK]);
422        assert!(direct.contains(&vanilla_items::STICK));
423        assert!(!direct.contains(&vanilla_items::DIAMOND));
424    }
425
426    #[test]
427    fn string_codec_rejects_malformed_identifiers() {
428        init_vanilla_registry();
429
430        let malformed = NbtTag::String("not an identifier".into());
431        assert_eq!(
432            with_borrowed_tag(malformed, RegistryHolderSet::<Item>::from_nbt_tag),
433            None
434        );
435        let missing_tag = Identifier::new_static("steel", "missing");
436        assert!(
437            RegistryHolderSet::<Item>::Tag(missing_tag)
438                .write(&mut Vec::new())
439                .is_err()
440        );
441    }
442}