steel_registry/data_components/components/
potion_contents.rs

1//! Vanilla `minecraft:potion_contents` item component.
2
3use std::io::{Cursor, Error, Result, Write};
4use std::num::Wrapping;
5
6use simdnbt::owned::{NbtCompound, NbtList, NbtTag};
7use simdnbt::{FromNbtTag, ToNbtTag};
8use steel_utils::Identifier;
9use steel_utils::codec::VarInt;
10use steel_utils::hash::{ComponentHasher, HashComponent, HashEntry, sort_map_entries};
11use steel_utils::nbt::NbtNumeric as _;
12use steel_utils::serial::{PrefixedRead as _, PrefixedWrite as _, ReadFrom, WriteTo};
13
14use crate::RegistryReference;
15use crate::data_components::vanilla_components::POTION_CONTENTS;
16use crate::item_stack::ItemStack;
17use crate::items::ItemRef;
18use crate::mob_effect_instance::MobEffectInstance;
19use crate::potion::{Potion, PotionRef};
20
21/// Color used when there is neither a custom color nor a visible effect.
22const BASE_POTION_COLOR: i32 = -13_083_194;
23/// Fully opaque alpha channel, matching vanilla `ARGB.color(r, g, b)`'s implicit
24/// `ARGB.color(255, r, g, b)`.
25const OPAQUE_ALPHA: i32 = 0xFF00_0000_u32 as i32;
26
27/// A registered base potion plus optional custom display and effect data.
28#[derive(Debug, Default, Clone, PartialEq)]
29pub struct PotionContents {
30    potion: Option<RegistryReference<Potion>>,
31    custom_color: Option<i32>,
32    custom_effects: Vec<MobEffectInstance>,
33    custom_name: Option<String>,
34}
35
36impl PotionContents {
37    const MAX_NETWORK_STRING_LENGTH: usize = 32_767;
38
39    #[must_use]
40    pub const fn empty() -> Self {
41        Self {
42            potion: None,
43            custom_color: None,
44            custom_effects: Vec::new(),
45            custom_name: None,
46        }
47    }
48
49    #[must_use]
50    pub const fn new(
51        potion: Option<RegistryReference<Potion>>,
52        custom_color: Option<i32>,
53        custom_effects: Vec<MobEffectInstance>,
54        custom_name: Option<String>,
55    ) -> Self {
56        Self {
57            potion,
58            custom_color,
59            custom_effects,
60            custom_name,
61        }
62    }
63
64    #[must_use]
65    pub const fn of(potion: PotionRef) -> Self {
66        Self::new(Some(RegistryReference::new(potion)), None, Vec::new(), None)
67    }
68
69    #[must_use]
70    pub fn create_item_stack(item: ItemRef, potion: PotionRef) -> ItemStack {
71        let mut stack = ItemStack::new(item);
72        stack.set(POTION_CONTENTS, Self::of(potion));
73        stack
74    }
75
76    #[must_use]
77    pub const fn potion(&self) -> Option<RegistryReference<Potion>> {
78        self.potion
79    }
80
81    #[must_use]
82    pub const fn custom_color(&self) -> Option<i32> {
83        self.custom_color
84    }
85
86    #[must_use]
87    pub fn custom_effects(&self) -> &[MobEffectInstance] {
88        &self.custom_effects
89    }
90
91    #[must_use]
92    pub fn custom_name(&self) -> Option<&str> {
93        self.custom_name.as_deref()
94    }
95
96    #[must_use]
97    pub fn is(&self, potion: &Potion) -> bool {
98        self.potion.is_some_and(|p| p.value().key == potion.key) && self.custom_effects.is_empty()
99    }
100
101    /// Whether the base potion is in `tag` and there are no custom effects.
102    #[must_use]
103    pub fn is_in_tag(&self, tag: &Identifier) -> bool {
104        self.custom_effects.is_empty() && self.potion.is_some_and(|p| p.value().has_tag(tag))
105    }
106
107    /// Returns vanilla `PotionContents.getAllEffects()`: the base potion's
108    /// effects followed by the custom effects.
109    #[must_use]
110    pub fn all_effects(&self) -> Vec<MobEffectInstance> {
111        let mut effects = Vec::with_capacity(self.custom_effects.len() + 1);
112        if let Some(potion) = self.potion {
113            effects.extend(potion.value().effects.iter().map(|effect| {
114                MobEffectInstance::simple(effect.effect, effect.duration, effect.amplifier)
115            }));
116        }
117        effects.extend(self.custom_effects.iter().cloned());
118        effects
119    }
120
121    /// Whether the base potion or the custom effects contain any effect.
122    #[must_use]
123    pub fn has_effects(&self) -> bool {
124        !self.custom_effects.is_empty()
125            || self
126                .potion
127                .is_some_and(|potion| !potion.value().effects.is_empty())
128    }
129
130    /// Returns the display color, falling back to `BASE_POTION_COLOR`.
131    #[must_use]
132    pub fn get_color(&self) -> i32 {
133        self.get_color_or(BASE_POTION_COLOR)
134    }
135
136    /// Returns the custom color, else the color blended from the effects,
137    /// else `default_color`.
138    #[must_use]
139    pub fn get_color_or(&self, default_color: i32) -> i32 {
140        self.custom_color
141            .or_else(|| Self::color_from_effects(&self.all_effects()))
142            .unwrap_or(default_color)
143    }
144
145    /// Blends visible effect colors weighted by amplifier + 1 and returns `None` when no effect is visible
146    fn color_from_effects(effects: &[MobEffectInstance]) -> Option<i32> {
147        // Vanilla accumulates in `int`, so huge effect lists wrap instead of
148        // widening; `ARGB.color` then masks each channel to its low byte.
149        let mut red = Wrapping(0_i32);
150        let mut green = Wrapping(0_i32);
151        let mut blue = Wrapping(0_i32);
152        let mut total_weight = Wrapping(0_i32);
153
154        for effect in effects {
155            if !effect.show_particles() {
156                continue;
157            }
158            let color = effect.effect().color;
159            let weight = Wrapping(effect.amplifier()) + Wrapping(1);
160            red += weight * Wrapping(i32::from(color.red()));
161            green += weight * Wrapping(i32::from(color.green()));
162            blue += weight * Wrapping(i32::from(color.blue()));
163            total_weight += weight;
164        }
165
166        if total_weight.0 == 0 {
167            return None;
168        }
169        let r = (red / total_weight).0 & 0xFF;
170        let g = (green / total_weight).0 & 0xFF;
171        let b = (blue / total_weight).0 & 0xFF;
172        Some(OPAQUE_ALPHA | (r << 16) | (g << 8) | b)
173    }
174
175    fn to_nbt_tag_ref(&self) -> NbtTag {
176        let mut compound = NbtCompound::new();
177        if let Some(potion) = self.potion {
178            compound.insert("potion", potion.to_nbt_tag());
179        }
180        if let Some(custom_color) = self.custom_color {
181            compound.insert("custom_color", custom_color);
182        }
183        if !self.custom_effects.is_empty() {
184            compound.insert(
185                "custom_effects",
186                NbtList::Compound(
187                    self.custom_effects
188                        .iter()
189                        .map(|effect| match effect.to_nbt_tag_ref() {
190                            NbtTag::Compound(compound) => compound,
191                            _ => unreachable!("mob effect codec always produces a compound"),
192                        })
193                        .collect(),
194                ),
195            );
196        }
197        if let Some(custom_name) = &self.custom_name {
198            compound.insert("custom_name", custom_name.clone());
199        }
200        NbtTag::Compound(compound)
201    }
202
203    fn from_owned_nbt(tag: &NbtTag) -> Option<Self> {
204        if tag.string().is_some() {
205            return registry_reference_from_owned_nbt(tag)
206                .map(|potion| Self::new(Some(potion), None, Vec::new(), None));
207        }
208
209        let compound = tag.compound()?;
210        let potion = match compound.get("potion") {
211            Some(tag) => Some(registry_reference_from_owned_nbt(tag)?),
212            None => None,
213        };
214        let custom_color = match compound.get("custom_color") {
215            Some(tag) => Some(tag.codec_i32()?),
216            None => None,
217        };
218        let custom_effects = match compound.get("custom_effects") {
219            Some(tag) => tag
220                .list()?
221                .as_nbt_tags()
222                .iter()
223                .map(MobEffectInstance::from_owned_nbt)
224                .collect::<Option<Vec<_>>>()?,
225            None => Vec::new(),
226        };
227        let custom_name = match compound.get("custom_name") {
228            Some(tag) => Some(tag.string()?.to_string()),
229            None => None,
230        };
231        Some(Self::new(potion, custom_color, custom_effects, custom_name))
232    }
233}
234
235impl WriteTo for PotionContents {
236    fn write(&self, writer: &mut impl Write) -> Result<()> {
237        self.potion.is_some().write(writer)?;
238        if let Some(potion) = self.potion {
239            potion.write(writer)?;
240        }
241        self.custom_color.is_some().write(writer)?;
242        if let Some(custom_color) = self.custom_color {
243            custom_color.write(writer)?;
244        }
245        write_count(self.custom_effects.len(), writer)?;
246        for effect in &self.custom_effects {
247            effect.write(writer)?;
248        }
249        self.custom_name.is_some().write(writer)?;
250        if let Some(custom_name) = &self.custom_name {
251            write_network_string(custom_name, writer)?;
252        }
253        Ok(())
254    }
255}
256
257impl ReadFrom for PotionContents {
258    fn read(data: &mut Cursor<&[u8]>) -> Result<Self> {
259        let potion = if bool::read(data)? {
260            Some(RegistryReference::read(data)?)
261        } else {
262            None
263        };
264        let custom_color = if bool::read(data)? {
265            Some(i32::read(data)?)
266        } else {
267            None
268        };
269        let count = read_count(data)?;
270        let mut custom_effects = Vec::with_capacity(count.min(65_536));
271        for _ in 0..count {
272            custom_effects.push(MobEffectInstance::read(data)?);
273        }
274        let custom_name = if bool::read(data)? {
275            Some(read_network_string(data)?)
276        } else {
277            None
278        };
279        Ok(Self::new(potion, custom_color, custom_effects, custom_name))
280    }
281}
282
283impl ToNbtTag for PotionContents {
284    fn to_nbt_tag(self) -> NbtTag {
285        self.to_nbt_tag_ref()
286    }
287}
288
289impl FromNbtTag for PotionContents {
290    fn from_nbt_tag(tag: simdnbt::borrow::NbtTag) -> Option<Self> {
291        Self::from_owned_nbt(&tag.to_owned())
292    }
293}
294
295impl HashComponent for PotionContents {
296    fn hash_component(&self, hasher: &mut ComponentHasher) {
297        let mut entries = Vec::with_capacity(4);
298        if let Some(potion) = &self.potion {
299            push_hash_entry(&mut entries, "potion", potion);
300        }
301        if let Some(custom_color) = self.custom_color {
302            push_hash_entry(&mut entries, "custom_color", &custom_color);
303        }
304        if !self.custom_effects.is_empty() {
305            push_hash_entry(
306                &mut entries,
307                "custom_effects",
308                &MobEffectList(&self.custom_effects),
309            );
310        }
311        if let Some(custom_name) = &self.custom_name {
312            push_hash_entry(&mut entries, "custom_name", custom_name);
313        }
314        sort_map_entries(&mut entries);
315        hasher.start_map();
316        for entry in entries {
317            hasher.put_raw_bytes(&entry.key_bytes);
318            hasher.put_raw_bytes(&entry.value_bytes);
319        }
320        hasher.end_map();
321    }
322}
323
324struct MobEffectList<'a>(&'a [MobEffectInstance]);
325
326impl HashComponent for MobEffectList<'_> {
327    fn hash_component(&self, hasher: &mut ComponentHasher) {
328        hasher.start_list();
329        for effect in self.0 {
330            hasher.put_component_hash(effect);
331        }
332        hasher.end_list();
333    }
334}
335
336fn registry_reference_from_owned_nbt(tag: &NbtTag) -> Option<RegistryReference<Potion>> {
337    let key = tag.string()?.to_string().parse().ok()?;
338    <Potion as crate::RegistryReferenceEntry>::reference_by_key(&key).map(RegistryReference::new)
339}
340
341fn write_count(count: usize, writer: &mut impl Write) -> Result<()> {
342    let count = i32::try_from(count).map_err(|_| Error::other("Effect list is too large"))?;
343    VarInt(count).write(writer)
344}
345
346fn read_count(data: &mut Cursor<&[u8]>) -> Result<usize> {
347    let count = VarInt::read(data)?.0;
348    usize::try_from(count).map_err(|_| Error::other(format!("Negative effect count: {count}")))
349}
350
351fn write_network_string(value: &str, writer: &mut impl Write) -> Result<()> {
352    if value.encode_utf16().count() > PotionContents::MAX_NETWORK_STRING_LENGTH
353        || value.len() > PotionContents::MAX_NETWORK_STRING_LENGTH * 3
354    {
355        return Err(Error::other("Potion custom name exceeds the network limit"));
356    }
357    value.write_prefixed::<VarInt>(writer)
358}
359
360fn read_network_string(data: &mut Cursor<&[u8]>) -> Result<String> {
361    let value =
362        String::read_prefixed_bound::<VarInt>(data, PotionContents::MAX_NETWORK_STRING_LENGTH * 3)?;
363    if value.encode_utf16().count() > PotionContents::MAX_NETWORK_STRING_LENGTH {
364        return Err(Error::other("Potion custom name exceeds the network limit"));
365    }
366    Ok(value)
367}
368
369fn push_hash_entry<T: HashComponent + ?Sized>(entries: &mut Vec<HashEntry>, key: &str, value: &T) {
370    let mut key_hasher = ComponentHasher::new();
371    key_hasher.put_string(key);
372    let mut value_hasher = ComponentHasher::new();
373    value.hash_component(&mut value_hasher);
374    entries.push(HashEntry::new(key_hasher, value_hasher));
375}
376
377#[cfg(test)]
378mod tests {
379    use std::io::Cursor;
380
381    use simdnbt::owned::NbtTag;
382    use simdnbt::{FromNbtTag as _, ToNbtTag as _};
383    use steel_utils::hash::HashComponent as _;
384    use steel_utils::serial::{ReadFrom as _, WriteTo as _};
385
386    use super::PotionContents;
387    use crate::data_components::vanilla_components::POTION_CONTENTS;
388    use crate::init_vanilla_registry;
389    use crate::{REGISTRY, RegistryExt, RegistryReference, vanilla_mob_effects, vanilla_potions};
390
391    fn parse(tag: NbtTag) -> Option<PotionContents> {
392        let mut bytes = Vec::new();
393        tag.write(&mut bytes);
394        let borrowed = simdnbt::borrow::read_tag(&mut Cursor::new(bytes.as_slice())).ok()?;
395        PotionContents::from_nbt_tag(borrowed.as_tag())
396    }
397
398    #[test]
399    fn full_and_alternative_potion_codecs_round_trip() {
400        init_vanilla_registry();
401        let value = PotionContents::new(
402            Some(RegistryReference::new(&vanilla_potions::SWIFTNESS)),
403            Some(0x12_34_56),
404            vec![crate::MobEffectInstance::simple(
405                vanilla_mob_effects::LUCK,
406                200,
407                1,
408            )],
409            Some("custom".to_owned()),
410        );
411        let nbt = value.clone().to_nbt_tag();
412        assert_eq!(parse(nbt.clone()), Some(value.clone()));
413        // MobEffectInstance contains Codec.BOOL fields while NbtOps represents
414        // those values as bytes.
415        assert_ne!(value.compute_hash(), nbt.compute_hash());
416
417        let mut network = Vec::new();
418        value
419            .write(&mut network)
420            .expect("potion contents should encode");
421        assert_eq!(
422            PotionContents::read(&mut Cursor::new(network.as_slice()))
423                .expect("potion contents should decode"),
424            value
425        );
426
427        assert_eq!(
428            parse(NbtTag::String("minecraft:water".into())),
429            Some(PotionContents::new(
430                Some(RegistryReference::new(&vanilla_potions::WATER)),
431                None,
432                Vec::new(),
433                None,
434            ))
435        );
436    }
437
438    #[test]
439    fn all_effects_combines_base_potion_then_custom_effects() {
440        init_vanilla_registry();
441        let custom = crate::MobEffectInstance::simple(vanilla_mob_effects::LUCK, 200, 1);
442        let contents = PotionContents::new(
443            Some(RegistryReference::new(&vanilla_potions::POISON)),
444            None,
445            vec![custom.clone()],
446            None,
447        );
448
449        let base_effects = vanilla_potions::POISON.effects;
450        let expected: Vec<_> = base_effects
451            .iter()
452            .map(|effect| {
453                crate::MobEffectInstance::simple(effect.effect, effect.duration, effect.amplifier)
454            })
455            .chain(std::iter::once(custom))
456            .collect();
457
458        assert_eq!(contents.all_effects(), expected);
459    }
460
461    #[test]
462    fn all_effects_is_empty_without_a_base_potion_or_custom_effects() {
463        init_vanilla_registry();
464        assert_eq!(PotionContents::empty().all_effects(), Vec::new());
465    }
466
467    #[test]
468    fn has_effects_checks_base_potion_and_custom_effects_independently() {
469        init_vanilla_registry();
470        assert!(!PotionContents::empty().has_effects());
471        assert!(
472            !PotionContents::new(
473                Some(RegistryReference::new(&vanilla_potions::WATER)),
474                None,
475                Vec::new(),
476                None
477            )
478            .has_effects()
479        );
480        assert!(
481            PotionContents::new(
482                Some(RegistryReference::new(&vanilla_potions::POISON)),
483                None,
484                Vec::new(),
485                None
486            )
487            .has_effects()
488        );
489        assert!(
490            PotionContents::new(
491                None,
492                None,
493                vec![crate::MobEffectInstance::simple(
494                    vanilla_mob_effects::LUCK,
495                    200,
496                    1
497                )],
498                None,
499            )
500            .has_effects()
501        );
502    }
503
504    #[test]
505    fn get_color_prefers_custom_color_over_effect_blend() {
506        init_vanilla_registry();
507        let contents = PotionContents::new(
508            Some(RegistryReference::new(&vanilla_potions::POISON)),
509            Some(0x00_ff_00),
510            Vec::new(),
511            None,
512        );
513        assert_eq!(contents.get_color(), 0x00_ff_00);
514    }
515
516    #[test]
517    fn get_color_falls_back_to_default_without_visible_effects() {
518        init_vanilla_registry();
519        assert_eq!(
520            PotionContents::empty().get_color(),
521            super::BASE_POTION_COLOR
522        );
523        assert_eq!(PotionContents::empty().get_color_or(0x11_22_33), 0x11_22_33);
524    }
525
526    #[test]
527    fn get_color_blends_visible_effect_colors_by_amplifier_weight() {
528        init_vanilla_registry();
529
530        // Poison's single effect is `RgbColor::new(0x87_A3_63)` at amplifier 0, so
531        // the weighted average is that color verbatim, returned as opaque ARGB
532        // exactly like vanilla's `ARGB.color(r, g, b)`.
533        let poison = PotionContents::new(
534            Some(RegistryReference::new(&vanilla_potions::POISON)),
535            None,
536            Vec::new(),
537            None,
538        );
539        assert_eq!(poison.get_color(), 0xFF_87_A3_63_u32 as i32);
540
541        // Two custom effects with different amplifiers blend weighted by
542        // `amplifier + 1`: red at weight 1 and blue at weight 3 average to
543        // `(255 / 4, 0, 3 * 255 / 4)` per channel.
544        let blended = PotionContents::new(
545            None,
546            None,
547            vec![
548                crate::MobEffectInstance::simple(vanilla_mob_effects::LUCK, 200, 0),
549                crate::MobEffectInstance::simple(vanilla_mob_effects::UNLUCK, 200, 2),
550            ],
551            None,
552        );
553        let luck = vanilla_mob_effects::LUCK.color;
554        let unluck = vanilla_mob_effects::UNLUCK.color;
555        let expected_channel = |luck_channel: u8, unluck_channel: u8| -> i32 {
556            (i32::from(luck_channel) + 3 * i32::from(unluck_channel)) / 4
557        };
558        assert_eq!(
559            blended.get_color(),
560            0xFF00_0000_u32 as i32
561                | (expected_channel(luck.red(), unluck.red()) << 16)
562                | (expected_channel(luck.green(), unluck.green()) << 8)
563                | expected_channel(luck.blue(), unluck.blue())
564        );
565    }
566
567    #[test]
568    fn extracted_potion_items_have_empty_contents() {
569        init_vanilla_registry();
570        for name in [
571            "potion",
572            "splash_potion",
573            "tipped_arrow",
574            "lingering_potion",
575        ] {
576            let item = REGISTRY
577                .items
578                .by_key(&steel_utils::Identifier::vanilla(name.to_owned()))
579                .expect("potion item should be registered");
580            assert_eq!(
581                item.components.get(POTION_CONTENTS),
582                Some(PotionContents::empty())
583            );
584        }
585    }
586}
steel_registry/data_components/components/potion_contents.rs

steel_registry/data_components/components/
potion_contents.rs
