steel_registry/data_components/components/
fireworks.rs

1//! Vanilla firework explosion and rocket components.
2
3use std::io::{Cursor, Error, Result, Write};
4
5use simdnbt::owned::{NbtCompound, NbtList, NbtTag};
6use simdnbt::{FromNbtTag, ToNbtTag};
7use steel_utils::codec::VarInt;
8use steel_utils::hash::{ComponentHasher, HashComponent, HashEntry, sort_map_entries};
9use steel_utils::nbt::NbtNumeric as _;
10use steel_utils::serial::{ReadFrom, WriteTo};
11
12/// Shape rendered by a firework explosion.
13#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
14pub enum FireworkExplosionShape {
15    #[default]
16    SmallBall,
17    LargeBall,
18    Star,
19    Creeper,
20    Burst,
21}
22
23impl FireworkExplosionShape {
24    #[must_use]
25    pub const fn id(self) -> i32 {
26        match self {
27            Self::SmallBall => 0,
28            Self::LargeBall => 1,
29            Self::Star => 2,
30            Self::Creeper => 3,
31            Self::Burst => 4,
32        }
33    }
34
35    #[must_use]
36    pub const fn serialized_name(self) -> &'static str {
37        match self {
38            Self::SmallBall => "small_ball",
39            Self::LargeBall => "large_ball",
40            Self::Star => "star",
41            Self::Creeper => "creeper",
42            Self::Burst => "burst",
43        }
44    }
45
46    #[must_use]
47    pub const fn by_id(id: i32) -> Self {
48        match id {
49            1 => Self::LargeBall,
50            2 => Self::Star,
51            3 => Self::Creeper,
52            4 => Self::Burst,
53            _ => Self::SmallBall,
54        }
55    }
56
57    const fn from_serialized_name(name: &str) -> Option<Self> {
58        match name {
59            "small_ball" => Some(Self::SmallBall),
60            "large_ball" => Some(Self::LargeBall),
61            "star" => Some(Self::Star),
62            "creeper" => Some(Self::Creeper),
63            "burst" => Some(Self::Burst),
64            _ => None,
65        }
66    }
67}
68
69/// Visual data for one firework explosion.
70#[derive(Debug, Default, Clone, PartialEq, Eq)]
71pub struct FireworkExplosion {
72    shape: FireworkExplosionShape,
73    colors: Vec<i32>,
74    fade_colors: Vec<i32>,
75    has_trail: bool,
76    has_twinkle: bool,
77}
78
79impl FireworkExplosion {
80    #[must_use]
81    pub const fn new(
82        shape: FireworkExplosionShape,
83        colors: Vec<i32>,
84        fade_colors: Vec<i32>,
85        has_trail: bool,
86        has_twinkle: bool,
87    ) -> Self {
88        Self {
89            shape,
90            colors,
91            fade_colors,
92            has_trail,
93            has_twinkle,
94        }
95    }
96
97    #[must_use]
98    pub const fn shape(&self) -> FireworkExplosionShape {
99        self.shape
100    }
101
102    #[must_use]
103    pub fn colors(&self) -> &[i32] {
104        &self.colors
105    }
106
107    #[must_use]
108    pub fn fade_colors(&self) -> &[i32] {
109        &self.fade_colors
110    }
111
112    #[must_use]
113    pub const fn has_trail(&self) -> bool {
114        self.has_trail
115    }
116
117    #[must_use]
118    pub const fn has_twinkle(&self) -> bool {
119        self.has_twinkle
120    }
121
122    #[must_use]
123    pub fn with_fade_colors(&self, fade_colors: Vec<i32>) -> Self {
124        Self {
125            shape: self.shape,
126            colors: self.colors.clone(),
127            fade_colors,
128            has_trail: self.has_trail,
129            has_twinkle: self.has_twinkle,
130        }
131    }
132
133    fn to_nbt_tag_ref(&self) -> NbtTag {
134        let mut compound = NbtCompound::new();
135        compound.insert("shape", self.shape.serialized_name());
136        if !self.colors.is_empty() {
137            compound.insert("colors", int_list_nbt(&self.colors));
138        }
139        if !self.fade_colors.is_empty() {
140            compound.insert("fade_colors", int_list_nbt(&self.fade_colors));
141        }
142        if self.has_trail {
143            compound.insert("has_trail", true);
144        }
145        if self.has_twinkle {
146            compound.insert("has_twinkle", true);
147        }
148        NbtTag::Compound(compound)
149    }
150
151    fn from_owned_nbt(tag: &NbtTag) -> Option<Self> {
152        let compound = tag.compound()?;
153        let shape = FireworkExplosionShape::from_serialized_name(
154            &compound.get("shape")?.string()?.to_string(),
155        )?;
156        let colors = match compound.get("colors") {
157            Some(tag) => int_list_from_nbt(tag)?,
158            None => Vec::new(),
159        };
160        let fade_colors = match compound.get("fade_colors") {
161            Some(tag) => int_list_from_nbt(tag)?,
162            None => Vec::new(),
163        };
164        let has_trail = optional_bool(compound.get("has_trail"), false)?;
165        let has_twinkle = optional_bool(compound.get("has_twinkle"), false)?;
166        Some(Self::new(
167            shape,
168            colors,
169            fade_colors,
170            has_trail,
171            has_twinkle,
172        ))
173    }
174}
175
176impl WriteTo for FireworkExplosion {
177    fn write(&self, writer: &mut impl Write) -> Result<()> {
178        VarInt(self.shape.id()).write(writer)?;
179        write_int_list(&self.colors, writer)?;
180        write_int_list(&self.fade_colors, writer)?;
181        self.has_trail.write(writer)?;
182        self.has_twinkle.write(writer)
183    }
184}
185
186impl ReadFrom for FireworkExplosion {
187    fn read(data: &mut Cursor<&[u8]>) -> Result<Self> {
188        Ok(Self::new(
189            FireworkExplosionShape::by_id(VarInt::read(data)?.0),
190            read_int_list(data)?,
191            read_int_list(data)?,
192            bool::read(data)?,
193            bool::read(data)?,
194        ))
195    }
196}
197
198impl ToNbtTag for FireworkExplosion {
199    fn to_nbt_tag(self) -> NbtTag {
200        self.to_nbt_tag_ref()
201    }
202}
203
204impl FromNbtTag for FireworkExplosion {
205    fn from_nbt_tag(tag: simdnbt::borrow::NbtTag) -> Option<Self> {
206        Self::from_owned_nbt(&tag.to_owned())
207    }
208}
209
210impl HashComponent for FireworkExplosion {
211    fn hash_component(&self, hasher: &mut ComponentHasher) {
212        let mut entries = Vec::with_capacity(5);
213        push_hash_entry(&mut entries, "shape", self.shape.serialized_name());
214        if !self.colors.is_empty() {
215            push_hash_entry(&mut entries, "colors", &CodecIntList(&self.colors));
216        }
217        if !self.fade_colors.is_empty() {
218            push_hash_entry(
219                &mut entries,
220                "fade_colors",
221                &CodecIntList(&self.fade_colors),
222            );
223        }
224        if self.has_trail {
225            push_hash_entry(&mut entries, "has_trail", &true);
226        }
227        if self.has_twinkle {
228            push_hash_entry(&mut entries, "has_twinkle", &true);
229        }
230        hash_entries(hasher, &mut entries);
231    }
232}
233
234/// Flight duration and ordered explosions carried by a firework rocket.
235#[derive(Debug, Clone, PartialEq, Eq)]
236pub struct Fireworks {
237    flight_duration: i32,
238    explosions: Vec<FireworkExplosion>,
239}
240
241impl Fireworks {
242    pub const MAX_EXPLOSIONS: usize = 256;
243
244    /// Creates a firework component accepted by Vanilla's stream codec.
245    pub fn new(flight_duration: i32, explosions: Vec<FireworkExplosion>) -> Result<Self> {
246        if explosions.len() > Self::MAX_EXPLOSIONS {
247            return Err(Error::other(format!(
248                "Got {} explosions, but maximum is {}",
249                explosions.len(),
250                Self::MAX_EXPLOSIONS
251            )));
252        }
253        Ok(Self {
254            flight_duration,
255            explosions,
256        })
257    }
258
259    pub(crate) const fn from_extracted(flight_duration: i32) -> Self {
260        assert!(
261            flight_duration >= 0 && flight_duration <= u8::MAX as i32,
262            "extracted firework flight duration must be in 0..=255"
263        );
264        Self {
265            flight_duration,
266            explosions: Vec::new(),
267        }
268    }
269
270    #[must_use]
271    pub const fn flight_duration(&self) -> i32 {
272        self.flight_duration
273    }
274
275    #[must_use]
276    pub fn explosions(&self) -> &[FireworkExplosion] {
277        &self.explosions
278    }
279
280    fn to_nbt_tag_ref(&self) -> NbtTag {
281        let mut compound = NbtCompound::new();
282        if self.flight_duration != 0 {
283            compound.insert("flight_duration", self.flight_duration as u8 as i8);
284        }
285        if !self.explosions.is_empty() {
286            compound.insert(
287                "explosions",
288                NbtList::Compound(
289                    self.explosions
290                        .iter()
291                        .map(|explosion| match explosion.to_nbt_tag_ref() {
292                            NbtTag::Compound(compound) => compound,
293                            _ => {
294                                unreachable!("firework explosion codec always produces a compound")
295                            }
296                        })
297                        .collect(),
298                ),
299            );
300        }
301        NbtTag::Compound(compound)
302    }
303
304    pub(crate) fn try_to_persistent_nbt(&self) -> Result<NbtTag> {
305        if self.flight_duration > i32::from(u8::MAX) {
306            return Err(Error::other("Firework flight duration exceeds 255"));
307        }
308        Ok(self.to_nbt_tag_ref())
309    }
310
311    fn from_owned_nbt(tag: &NbtTag) -> Option<Self> {
312        let compound = tag.compound()?;
313        let flight_duration = match compound.get("flight_duration") {
314            Some(tag) => i32::from(tag.codec_i32()? as i8 as u8),
315            None => 0,
316        };
317        let explosions = match compound.get("explosions") {
318            Some(tag) => {
319                let list = tag.list()?.as_nbt_tags();
320                if list.len() > Self::MAX_EXPLOSIONS {
321                    return None;
322                }
323                list.iter()
324                    .map(FireworkExplosion::from_owned_nbt)
325                    .collect::<Option<Vec<_>>>()?
326            }
327            None => Vec::new(),
328        };
329        Self::new(flight_duration, explosions).ok()
330    }
331}
332
333impl WriteTo for Fireworks {
334    fn write(&self, writer: &mut impl Write) -> Result<()> {
335        VarInt(self.flight_duration).write(writer)?;
336        write_count(self.explosions.len(), Self::MAX_EXPLOSIONS, writer)?;
337        for explosion in &self.explosions {
338            explosion.write(writer)?;
339        }
340        Ok(())
341    }
342}
343
344impl ReadFrom for Fireworks {
345    fn read(data: &mut Cursor<&[u8]>) -> Result<Self> {
346        let flight_duration = VarInt::read(data)?.0;
347        let count = read_count(data, Self::MAX_EXPLOSIONS)?;
348        let mut explosions = Vec::with_capacity(count);
349        for _ in 0..count {
350            explosions.push(FireworkExplosion::read(data)?);
351        }
352        Self::new(flight_duration, explosions)
353    }
354}
355
356impl ToNbtTag for Fireworks {
357    fn to_nbt_tag(self) -> NbtTag {
358        self.to_nbt_tag_ref()
359    }
360}
361
362impl FromNbtTag for Fireworks {
363    fn from_nbt_tag(tag: simdnbt::borrow::NbtTag) -> Option<Self> {
364        Self::from_owned_nbt(&tag.to_owned())
365    }
366}
367
368impl HashComponent for Fireworks {
369    fn hash_component(&self, hasher: &mut ComponentHasher) {
370        let mut entries = Vec::with_capacity(2);
371        if self.flight_duration != 0 {
372            push_hash_entry(
373                &mut entries,
374                "flight_duration",
375                &(self.flight_duration as u8 as i8),
376            );
377        }
378        if !self.explosions.is_empty() {
379            push_hash_entry(
380                &mut entries,
381                "explosions",
382                &CodecExplosionList(&self.explosions),
383            );
384        }
385        hash_entries(hasher, &mut entries);
386    }
387}
388
389struct CodecIntList<'a>(&'a [i32]);
390
391impl HashComponent for CodecIntList<'_> {
392    fn hash_component(&self, hasher: &mut ComponentHasher) {
393        hasher.start_list();
394        for color in self.0 {
395            hasher.put_component_hash(color);
396        }
397        hasher.end_list();
398    }
399}
400
401struct CodecExplosionList<'a>(&'a [FireworkExplosion]);
402
403impl HashComponent for CodecExplosionList<'_> {
404    fn hash_component(&self, hasher: &mut ComponentHasher) {
405        hasher.start_list();
406        for explosion in self.0 {
407            hasher.put_component_hash(explosion);
408        }
409        hasher.end_list();
410    }
411}
412
413fn int_list_nbt(values: &[i32]) -> NbtList {
414    if values.is_empty() {
415        NbtList::Empty
416    } else {
417        NbtList::Int(values.to_vec())
418    }
419}
420
421fn int_list_from_nbt(tag: &NbtTag) -> Option<Vec<i32>> {
422    tag.list()?
423        .as_nbt_tags()
424        .iter()
425        .map(steel_utils::nbt::NbtNumeric::codec_i32)
426        .collect()
427}
428
429fn optional_bool(tag: Option<&NbtTag>, default: bool) -> Option<bool> {
430    match tag {
431        Some(tag) => tag.codec_bool(),
432        None => Some(default),
433    }
434}
435
436fn write_int_list(values: &[i32], writer: &mut impl Write) -> Result<()> {
437    write_count(values.len(), i32::MAX as usize, writer)?;
438    for value in values {
439        value.write(writer)?;
440    }
441    Ok(())
442}
443
444fn read_int_list(data: &mut Cursor<&[u8]>) -> Result<Vec<i32>> {
445    let count = read_count(data, i32::MAX as usize)?;
446    let mut values = Vec::with_capacity(count.min(65_536));
447    for _ in 0..count {
448        values.push(i32::read(data)?);
449    }
450    Ok(values)
451}
452
453fn write_count(count: usize, max: usize, writer: &mut impl Write) -> Result<()> {
454    if count > max || count > i32::MAX as usize {
455        return Err(Error::other(format!(
456            "Collection size {count} exceeds {max}"
457        )));
458    }
459    VarInt(count as i32).write(writer)
460}
461
462fn read_count(data: &mut Cursor<&[u8]>, max: usize) -> Result<usize> {
463    let count = VarInt::read(data)?.0;
464    let count = usize::try_from(count)
465        .map_err(|_| Error::other(format!("Negative collection size: {count}")))?;
466    if count > max {
467        return Err(Error::other(format!(
468            "Collection size {count} exceeds {max}"
469        )));
470    }
471    Ok(count)
472}
473
474fn push_hash_entry<T: HashComponent + ?Sized>(entries: &mut Vec<HashEntry>, key: &str, value: &T) {
475    let mut key_hasher = ComponentHasher::new();
476    key_hasher.put_string(key);
477    let mut value_hasher = ComponentHasher::new();
478    value.hash_component(&mut value_hasher);
479    entries.push(HashEntry::new(key_hasher, value_hasher));
480}
481
482fn hash_entries(hasher: &mut ComponentHasher, entries: &mut [HashEntry]) {
483    sort_map_entries(entries);
484    hasher.start_map();
485    for entry in entries {
486        hasher.put_raw_bytes(&entry.key_bytes);
487        hasher.put_raw_bytes(&entry.value_bytes);
488    }
489    hasher.end_map();
490}
491
492#[cfg(test)]
493mod tests {
494    use std::io::Cursor;
495
496    use simdnbt::ToNbtTag as _;
497    use simdnbt::owned::{NbtCompound, NbtList, NbtTag};
498    use steel_utils::hash::HashComponent as _;
499    use steel_utils::serial::{ReadFrom as _, WriteTo as _};
500
501    use super::{FireworkExplosion, FireworkExplosionShape, Fireworks};
502    use crate::data_components::vanilla_components::FIREWORKS;
503    use crate::init_vanilla_registry;
504    use crate::{REGISTRY, RegistryExt};
505
506    fn parse<T: simdnbt::FromNbtTag>(tag: NbtTag) -> Option<T> {
507        let mut bytes = Vec::new();
508        tag.write(&mut bytes);
509        let borrowed = simdnbt::borrow::read_tag(&mut Cursor::new(bytes.as_slice())).ok()?;
510        T::from_nbt_tag(borrowed.as_tag())
511    }
512
513    #[test]
514    fn explosion_codecs_match_shape_and_optional_fields() {
515        let explosion = FireworkExplosion::new(
516            FireworkExplosionShape::Star,
517            vec![0x123456],
518            vec![0x654321],
519            true,
520            false,
521        );
522        let nbt = explosion.clone().to_nbt_tag();
523        assert_eq!(parse(nbt.clone()), Some(explosion.clone()));
524        // HashOps preserves Codec.BOOL while NbtOps represents booleans as bytes.
525        assert_ne!(explosion.compute_hash(), nbt.compute_hash());
526
527        let mut network = Vec::new();
528        explosion
529            .write(&mut network)
530            .expect("explosion should encode");
531        assert_eq!(
532            FireworkExplosion::read(&mut Cursor::new(network.as_slice()))
533                .expect("explosion should decode"),
534            explosion
535        );
536    }
537
538    #[test]
539    fn unknown_stream_shape_ids_fall_back_to_small_ball() {
540        let mut network = Vec::new();
541        steel_utils::codec::VarInt(99)
542            .write(&mut network)
543            .expect("shape should encode");
544        steel_utils::codec::VarInt(0)
545            .write(&mut network)
546            .expect("colors should encode");
547        steel_utils::codec::VarInt(0)
548            .write(&mut network)
549            .expect("fade colors should encode");
550        false.write(&mut network).expect("trail should encode");
551        false.write(&mut network).expect("twinkle should encode");
552        let decoded = FireworkExplosion::read(&mut Cursor::new(network.as_slice()))
553            .expect("explosion should decode");
554        assert_eq!(decoded.shape(), FireworkExplosionShape::SmallBall);
555    }
556
557    #[test]
558    fn fireworks_use_unsigned_byte_persistence_and_bounded_lists() {
559        let firework =
560            Fireworks::new(255, vec![FireworkExplosion::default()]).expect("valid firework");
561        let mut explosion = NbtCompound::new();
562        explosion.insert("shape", "small_ball");
563        let mut expected = NbtCompound::new();
564        expected.insert("flight_duration", -1_i8);
565        expected.insert("explosions", NbtList::Compound(vec![explosion]));
566        let expected = NbtTag::Compound(expected);
567        assert_eq!(firework.clone().to_nbt_tag(), expected);
568        assert_eq!(parse(expected.clone()), Some(firework.clone()));
569        assert_eq!(firework.compute_hash(), expected.compute_hash());
570
571        let mut network = Vec::new();
572        firework
573            .write(&mut network)
574            .expect("firework should encode");
575        assert_eq!(
576            Fireworks::read(&mut Cursor::new(network.as_slice())).expect("firework should decode"),
577            firework
578        );
579        let negative = Fireworks::new(-1, Vec::new()).expect("stream permits negative flight");
580        let oversized = Fireworks::new(256, Vec::new()).expect("stream permits any VarInt flight");
581        for value in [&negative, &oversized] {
582            let mut encoded = Vec::new();
583            value
584                .write(&mut encoded)
585                .expect("stream value should encode");
586            assert_eq!(
587                Fireworks::read(&mut Cursor::new(encoded.as_slice()))
588                    .expect("stream value should decode"),
589                *value
590            );
591        }
592        assert!(negative.try_to_persistent_nbt().is_ok());
593        assert!(oversized.try_to_persistent_nbt().is_err());
594        assert!(Fireworks::new(0, vec![FireworkExplosion::default(); 257]).is_err());
595    }
596
597    #[test]
598    fn extracted_firework_rocket_has_one_unit_of_flight() {
599        init_vanilla_registry();
600        let rocket = REGISTRY
601            .items
602            .by_key(&steel_utils::Identifier::vanilla_static("firework_rocket"))
603            .expect("firework rocket should be registered");
604        assert_eq!(
605            rocket.components.get(FIREWORKS),
606            Some(Fireworks::new(1, Vec::new()).expect("valid default rocket"))
607        );
608    }
609}
steel_registry/data_components/components/fireworks.rs

steel_registry/data_components/components/
fireworks.rs
