steel_registry/data_components/components/
potion_contents.rs

1//! Vanilla `minecraft:potion_contents` item component.
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::{PrefixedRead as _, PrefixedWrite as _, ReadFrom, WriteTo};
11
12use crate::RegistryReference;
13use crate::mob_effect_instance::MobEffectInstance;
14use crate::potion::Potion;
15
16/// A registered base potion plus optional custom display and effect data.
17#[derive(Debug, Default, Clone, PartialEq)]
18pub struct PotionContents {
19    potion: Option<RegistryReference<Potion>>,
20    custom_color: Option<i32>,
21    custom_effects: Vec<MobEffectInstance>,
22    custom_name: Option<String>,
23}
24
25impl PotionContents {
26    const MAX_NETWORK_STRING_LENGTH: usize = 32_767;
27
28    #[must_use]
29    pub const fn empty() -> Self {
30        Self {
31            potion: None,
32            custom_color: None,
33            custom_effects: Vec::new(),
34            custom_name: None,
35        }
36    }
37
38    #[must_use]
39    pub const fn new(
40        potion: Option<RegistryReference<Potion>>,
41        custom_color: Option<i32>,
42        custom_effects: Vec<MobEffectInstance>,
43        custom_name: Option<String>,
44    ) -> Self {
45        Self {
46            potion,
47            custom_color,
48            custom_effects,
49            custom_name,
50        }
51    }
52
53    #[must_use]
54    pub const fn potion(&self) -> Option<RegistryReference<Potion>> {
55        self.potion
56    }
57
58    #[must_use]
59    pub const fn custom_color(&self) -> Option<i32> {
60        self.custom_color
61    }
62
63    #[must_use]
64    pub fn custom_effects(&self) -> &[MobEffectInstance] {
65        &self.custom_effects
66    }
67
68    #[must_use]
69    pub fn custom_name(&self) -> Option<&str> {
70        self.custom_name.as_deref()
71    }
72
73    #[must_use]
74    pub fn is(&self, potion: &Potion) -> bool {
75        self.potion.is_some_and(|p| p.value().key == potion.key) && self.custom_effects.is_empty()
76    }
77
78    /// Returns vanilla `PotionContents.getAllEffects()`: the base potion's
79    /// effects followed by the custom effects.
80    #[must_use]
81    pub fn all_effects(&self) -> Vec<MobEffectInstance> {
82        let mut effects = Vec::with_capacity(self.custom_effects.len() + 1);
83        if let Some(potion) = self.potion {
84            effects.extend(potion.value().effects.iter().map(|effect| {
85                MobEffectInstance::simple(effect.effect, effect.duration, effect.amplifier)
86            }));
87        }
88        effects.extend(self.custom_effects.iter().cloned());
89        effects
90    }
91
92    fn to_nbt_tag_ref(&self) -> NbtTag {
93        let mut compound = NbtCompound::new();
94        if let Some(potion) = self.potion {
95            compound.insert("potion", potion.to_nbt_tag());
96        }
97        if let Some(custom_color) = self.custom_color {
98            compound.insert("custom_color", custom_color);
99        }
100        if !self.custom_effects.is_empty() {
101            compound.insert(
102                "custom_effects",
103                NbtList::Compound(
104                    self.custom_effects
105                        .iter()
106                        .map(|effect| match effect.to_nbt_tag_ref() {
107                            NbtTag::Compound(compound) => compound,
108                            _ => unreachable!("mob effect codec always produces a compound"),
109                        })
110                        .collect(),
111                ),
112            );
113        }
114        if let Some(custom_name) = &self.custom_name {
115            compound.insert("custom_name", custom_name.clone());
116        }
117        NbtTag::Compound(compound)
118    }
119
120    fn from_owned_nbt(tag: &NbtTag) -> Option<Self> {
121        if tag.string().is_some() {
122            return registry_reference_from_owned_nbt(tag)
123                .map(|potion| Self::new(Some(potion), None, Vec::new(), None));
124        }
125
126        let compound = tag.compound()?;
127        let potion = match compound.get("potion") {
128            Some(tag) => Some(registry_reference_from_owned_nbt(tag)?),
129            None => None,
130        };
131        let custom_color = match compound.get("custom_color") {
132            Some(tag) => Some(tag.codec_i32()?),
133            None => None,
134        };
135        let custom_effects = match compound.get("custom_effects") {
136            Some(tag) => tag
137                .list()?
138                .as_nbt_tags()
139                .iter()
140                .map(MobEffectInstance::from_owned_nbt)
141                .collect::<Option<Vec<_>>>()?,
142            None => Vec::new(),
143        };
144        let custom_name = match compound.get("custom_name") {
145            Some(tag) => Some(tag.string()?.to_string()),
146            None => None,
147        };
148        Some(Self::new(potion, custom_color, custom_effects, custom_name))
149    }
150}
151
152impl WriteTo for PotionContents {
153    fn write(&self, writer: &mut impl Write) -> Result<()> {
154        self.potion.is_some().write(writer)?;
155        if let Some(potion) = self.potion {
156            potion.write(writer)?;
157        }
158        self.custom_color.is_some().write(writer)?;
159        if let Some(custom_color) = self.custom_color {
160            custom_color.write(writer)?;
161        }
162        write_count(self.custom_effects.len(), writer)?;
163        for effect in &self.custom_effects {
164            effect.write(writer)?;
165        }
166        self.custom_name.is_some().write(writer)?;
167        if let Some(custom_name) = &self.custom_name {
168            write_network_string(custom_name, writer)?;
169        }
170        Ok(())
171    }
172}
173
174impl ReadFrom for PotionContents {
175    fn read(data: &mut Cursor<&[u8]>) -> Result<Self> {
176        let potion = if bool::read(data)? {
177            Some(RegistryReference::read(data)?)
178        } else {
179            None
180        };
181        let custom_color = if bool::read(data)? {
182            Some(i32::read(data)?)
183        } else {
184            None
185        };
186        let count = read_count(data)?;
187        let mut custom_effects = Vec::with_capacity(count.min(65_536));
188        for _ in 0..count {
189            custom_effects.push(MobEffectInstance::read(data)?);
190        }
191        let custom_name = if bool::read(data)? {
192            Some(read_network_string(data)?)
193        } else {
194            None
195        };
196        Ok(Self::new(potion, custom_color, custom_effects, custom_name))
197    }
198}
199
200impl ToNbtTag for PotionContents {
201    fn to_nbt_tag(self) -> NbtTag {
202        self.to_nbt_tag_ref()
203    }
204}
205
206impl FromNbtTag for PotionContents {
207    fn from_nbt_tag(tag: simdnbt::borrow::NbtTag) -> Option<Self> {
208        Self::from_owned_nbt(&tag.to_owned())
209    }
210}
211
212impl HashComponent for PotionContents {
213    fn hash_component(&self, hasher: &mut ComponentHasher) {
214        let mut entries = Vec::with_capacity(4);
215        if let Some(potion) = &self.potion {
216            push_hash_entry(&mut entries, "potion", potion);
217        }
218        if let Some(custom_color) = self.custom_color {
219            push_hash_entry(&mut entries, "custom_color", &custom_color);
220        }
221        if !self.custom_effects.is_empty() {
222            push_hash_entry(
223                &mut entries,
224                "custom_effects",
225                &MobEffectList(&self.custom_effects),
226            );
227        }
228        if let Some(custom_name) = &self.custom_name {
229            push_hash_entry(&mut entries, "custom_name", custom_name);
230        }
231        sort_map_entries(&mut entries);
232        hasher.start_map();
233        for entry in entries {
234            hasher.put_raw_bytes(&entry.key_bytes);
235            hasher.put_raw_bytes(&entry.value_bytes);
236        }
237        hasher.end_map();
238    }
239}
240
241struct MobEffectList<'a>(&'a [MobEffectInstance]);
242
243impl HashComponent for MobEffectList<'_> {
244    fn hash_component(&self, hasher: &mut ComponentHasher) {
245        hasher.start_list();
246        for effect in self.0 {
247            hasher.put_component_hash(effect);
248        }
249        hasher.end_list();
250    }
251}
252
253fn registry_reference_from_owned_nbt(tag: &NbtTag) -> Option<RegistryReference<Potion>> {
254    let key = tag.string()?.to_string().parse().ok()?;
255    <Potion as crate::RegistryReferenceEntry>::reference_by_key(&key).map(RegistryReference::new)
256}
257
258fn write_count(count: usize, writer: &mut impl Write) -> Result<()> {
259    let count = i32::try_from(count).map_err(|_| Error::other("Effect list is too large"))?;
260    VarInt(count).write(writer)
261}
262
263fn read_count(data: &mut Cursor<&[u8]>) -> Result<usize> {
264    let count = VarInt::read(data)?.0;
265    usize::try_from(count).map_err(|_| Error::other(format!("Negative effect count: {count}")))
266}
267
268fn write_network_string(value: &str, writer: &mut impl Write) -> Result<()> {
269    if value.encode_utf16().count() > PotionContents::MAX_NETWORK_STRING_LENGTH
270        || value.len() > PotionContents::MAX_NETWORK_STRING_LENGTH * 3
271    {
272        return Err(Error::other("Potion custom name exceeds the network limit"));
273    }
274    value.write_prefixed::<VarInt>(writer)
275}
276
277fn read_network_string(data: &mut Cursor<&[u8]>) -> Result<String> {
278    let value =
279        String::read_prefixed_bound::<VarInt>(data, PotionContents::MAX_NETWORK_STRING_LENGTH * 3)?;
280    if value.encode_utf16().count() > PotionContents::MAX_NETWORK_STRING_LENGTH {
281        return Err(Error::other("Potion custom name exceeds the network limit"));
282    }
283    Ok(value)
284}
285
286fn push_hash_entry<T: HashComponent + ?Sized>(entries: &mut Vec<HashEntry>, key: &str, value: &T) {
287    let mut key_hasher = ComponentHasher::new();
288    key_hasher.put_string(key);
289    let mut value_hasher = ComponentHasher::new();
290    value.hash_component(&mut value_hasher);
291    entries.push(HashEntry::new(key_hasher, value_hasher));
292}
293
294#[cfg(test)]
295mod tests {
296    use std::io::Cursor;
297
298    use simdnbt::owned::NbtTag;
299    use simdnbt::{FromNbtTag as _, ToNbtTag as _};
300    use steel_utils::hash::HashComponent as _;
301    use steel_utils::serial::{ReadFrom as _, WriteTo as _};
302
303    use super::PotionContents;
304    use crate::data_components::vanilla_components::POTION_CONTENTS;
305    use crate::init_vanilla_registry;
306    use crate::{REGISTRY, RegistryExt, RegistryReference, vanilla_mob_effects, vanilla_potions};
307
308    fn parse(tag: NbtTag) -> Option<PotionContents> {
309        let mut bytes = Vec::new();
310        tag.write(&mut bytes);
311        let borrowed = simdnbt::borrow::read_tag(&mut Cursor::new(bytes.as_slice())).ok()?;
312        PotionContents::from_nbt_tag(borrowed.as_tag())
313    }
314
315    #[test]
316    fn full_and_alternative_potion_codecs_round_trip() {
317        init_vanilla_registry();
318        let value = PotionContents::new(
319            Some(RegistryReference::new(&vanilla_potions::SWIFTNESS)),
320            Some(0x12_34_56),
321            vec![crate::MobEffectInstance::simple(
322                vanilla_mob_effects::LUCK,
323                200,
324                1,
325            )],
326            Some("custom".to_owned()),
327        );
328        let nbt = value.clone().to_nbt_tag();
329        assert_eq!(parse(nbt.clone()), Some(value.clone()));
330        // MobEffectInstance contains Codec.BOOL fields while NbtOps represents
331        // those values as bytes.
332        assert_ne!(value.compute_hash(), nbt.compute_hash());
333
334        let mut network = Vec::new();
335        value
336            .write(&mut network)
337            .expect("potion contents should encode");
338        assert_eq!(
339            PotionContents::read(&mut Cursor::new(network.as_slice()))
340                .expect("potion contents should decode"),
341            value
342        );
343
344        assert_eq!(
345            parse(NbtTag::String("minecraft:water".into())),
346            Some(PotionContents::new(
347                Some(RegistryReference::new(&vanilla_potions::WATER)),
348                None,
349                Vec::new(),
350                None,
351            ))
352        );
353    }
354
355    #[test]
356    fn all_effects_combines_base_potion_then_custom_effects() {
357        init_vanilla_registry();
358        let custom = crate::MobEffectInstance::simple(vanilla_mob_effects::LUCK, 200, 1);
359        let contents = PotionContents::new(
360            Some(RegistryReference::new(&vanilla_potions::POISON)),
361            None,
362            vec![custom.clone()],
363            None,
364        );
365
366        let base_effects = vanilla_potions::POISON.effects;
367        let expected: Vec<_> = base_effects
368            .iter()
369            .map(|effect| {
370                crate::MobEffectInstance::simple(effect.effect, effect.duration, effect.amplifier)
371            })
372            .chain(std::iter::once(custom))
373            .collect();
374
375        assert_eq!(contents.all_effects(), expected);
376    }
377
378    #[test]
379    fn all_effects_is_empty_without_a_base_potion_or_custom_effects() {
380        init_vanilla_registry();
381        assert_eq!(PotionContents::empty().all_effects(), Vec::new());
382    }
383
384    #[test]
385    fn extracted_potion_items_have_empty_contents() {
386        init_vanilla_registry();
387        for name in [
388            "potion",
389            "splash_potion",
390            "tipped_arrow",
391            "lingering_potion",
392        ] {
393            let item = REGISTRY
394                .items
395                .by_key(&steel_utils::Identifier::vanilla(name.to_owned()))
396                .expect("potion item should be registered");
397            assert_eq!(
398                item.components.get(POTION_CONTENTS),
399                Some(PotionContents::empty())
400            );
401        }
402    }
403}
steel_registry/data_components/components/potion_contents.rs

steel_registry/data_components/components/
potion_contents.rs
