1use 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
31pub trait RegistryHolderSetEntry: RegistryEntry + Debug + Send + Sync {
37 const REGISTRY_NAME: &'static str;
39
40 fn holder_set_by_id(id: usize) -> Option<&'static Self>;
42
43 fn holder_set_by_key(key: &Identifier) -> Option<&'static Self>;
45
46 fn holder_set_tag_exists(tag: &Identifier) -> bool;
48
49 fn holder_set_tag_contains(entry: &'static Self, tag: &Identifier) -> bool;
51}
52
53#[derive(Debug, PartialEq)]
55pub enum RegistryHolderSet<T: RegistryHolderSetEntry> {
56 Tag(Identifier),
58 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 #[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}