Skip to main content

steel_registry/data_components/components/
container_loot.rs

1//! Vanilla `minecraft:container_loot` item component.
2
3use std::io::{Cursor, Error, Result, Write};
4use std::str::FromStr;
5
6use simdnbt::owned::{NbtCompound, NbtTag, read_tag};
7use simdnbt::{FromNbtTag, ToNbtTag};
8use steel_utils::Identifier;
9use steel_utils::hash::{ComponentHasher, HashComponent, HashEntry, sort_map_entries};
10use steel_utils::nbt::vanilla_nbt_heap_size;
11use steel_utils::serial::{ReadFrom, WriteTo};
12
13const DEFAULT_NBT_QUOTA: u64 = 2_097_152;
14
15/// Loot-table resource key and optional deterministic seed for a container item.
16#[derive(Debug, Clone, PartialEq, Eq)]
17pub struct SeededContainerLoot {
18    loot_table: Identifier,
19    seed: i64,
20}
21
22impl SeededContainerLoot {
23    #[must_use]
24    pub const fn new(loot_table: Identifier, seed: i64) -> Self {
25        Self { loot_table, seed }
26    }
27
28    #[must_use]
29    pub const fn loot_table(&self) -> &Identifier {
30        &self.loot_table
31    }
32
33    #[must_use]
34    pub const fn seed(&self) -> i64 {
35        self.seed
36    }
37
38    fn to_nbt_tag_ref(&self) -> NbtTag {
39        let mut compound = NbtCompound::new();
40        compound.insert("loot_table", self.loot_table.to_string());
41        if self.seed != 0 {
42            compound.insert("seed", self.seed);
43        }
44        NbtTag::Compound(compound)
45    }
46
47    fn from_owned_nbt(tag: &NbtTag) -> Option<Self> {
48        let compound = tag.compound()?;
49        let loot_table =
50            Identifier::from_str(&compound.get("loot_table")?.string()?.to_string()).ok()?;
51        let seed = match compound.get("seed") {
52            Some(tag) => codec_i64(tag)?,
53            None => 0,
54        };
55        Some(Self::new(loot_table, seed))
56    }
57}
58
59impl WriteTo for SeededContainerLoot {
60    fn write(&self, writer: &mut impl Write) -> Result<()> {
61        let mut encoded = Vec::new();
62        self.to_nbt_tag_ref().write(&mut encoded);
63        writer.write_all(&encoded)
64    }
65}
66
67impl ReadFrom for SeededContainerLoot {
68    fn read(data: &mut Cursor<&[u8]>) -> Result<Self> {
69        let tag =
70            read_tag(data).map_err(|error| Error::other(format!("Invalid NBT: {error:?}")))?;
71        let Some(heap_size) = vanilla_nbt_heap_size(&tag) else {
72            return Err(Error::other("NBT contains malformed modified UTF-8"));
73        };
74        if heap_size > DEFAULT_NBT_QUOTA {
75            return Err(Error::other(format!(
76                "NBT exceeds Vanilla's {DEFAULT_NBT_QUOTA}-byte heap quota"
77            )));
78        }
79        Self::from_owned_nbt(&tag)
80            .ok_or_else(|| Error::other("Container loot network value is malformed"))
81    }
82}
83
84impl ToNbtTag for SeededContainerLoot {
85    fn to_nbt_tag(self) -> NbtTag {
86        self.to_nbt_tag_ref()
87    }
88}
89
90impl FromNbtTag for SeededContainerLoot {
91    fn from_nbt_tag(tag: simdnbt::borrow::NbtTag) -> Option<Self> {
92        Self::from_owned_nbt(&tag.to_owned())
93    }
94}
95
96impl HashComponent for SeededContainerLoot {
97    fn hash_component(&self, hasher: &mut ComponentHasher) {
98        let mut entries = Vec::with_capacity(2);
99        push_hash_entry(&mut entries, "loot_table", &self.loot_table);
100        if self.seed != 0 {
101            push_hash_entry(&mut entries, "seed", &self.seed);
102        }
103        sort_map_entries(&mut entries);
104        hasher.start_map();
105        for entry in &entries {
106            hasher.put_raw_bytes(&entry.key_bytes);
107            hasher.put_raw_bytes(&entry.value_bytes);
108        }
109        hasher.end_map();
110    }
111}
112
113fn codec_i64(tag: &NbtTag) -> Option<i64> {
114    match tag {
115        NbtTag::Byte(value) => Some(i64::from(*value)),
116        NbtTag::Short(value) => Some(i64::from(*value)),
117        NbtTag::Int(value) => Some(i64::from(*value)),
118        NbtTag::Long(value) => Some(*value),
119        NbtTag::Float(value) => Some(*value as i64),
120        NbtTag::Double(value) => Some(*value as i64),
121        _ => None,
122    }
123}
124
125fn push_hash_entry<T: HashComponent + ?Sized>(entries: &mut Vec<HashEntry>, key: &str, value: &T) {
126    let mut key_hasher = ComponentHasher::new();
127    key_hasher.put_string(key);
128    let mut value_hasher = ComponentHasher::new();
129    value.hash_component(&mut value_hasher);
130    entries.push(HashEntry::new(key_hasher, value_hasher));
131}
132
133#[cfg(test)]
134mod tests {
135    use std::io::Cursor;
136
137    use simdnbt::owned::{NbtCompound, NbtTag};
138    use simdnbt::{FromNbtTag as _, ToNbtTag as _};
139    use steel_utils::Identifier;
140    use steel_utils::hash::HashComponent as _;
141    use steel_utils::serial::{ReadFrom as _, WriteTo as _};
142
143    use super::SeededContainerLoot;
144
145    fn parse(tag: NbtTag) -> Option<SeededContainerLoot> {
146        let mut bytes = Vec::new();
147        tag.write(&mut bytes);
148        let borrowed = simdnbt::borrow::read_tag(&mut Cursor::new(bytes.as_slice())).ok()?;
149        SeededContainerLoot::from_nbt_tag(borrowed.as_tag())
150    }
151
152    #[test]
153    fn loot_table_resource_keys_and_optional_seed_round_trip() {
154        let loot =
155            SeededContainerLoot::new(Identifier::vanilla_static("chests/simple_dungeon"), 42);
156        let mut compound = NbtCompound::new();
157        compound.insert("loot_table", "minecraft:chests/simple_dungeon");
158        compound.insert("seed", 42_i64);
159        let nbt = NbtTag::Compound(compound);
160        assert_eq!(loot.clone().to_nbt_tag(), nbt);
161        assert_eq!(parse(nbt.clone()), Some(loot.clone()));
162        assert_eq!(loot.compute_hash(), nbt.compute_hash());
163
164        let mut network = Vec::new();
165        loot.write(&mut network)
166            .expect("container loot should encode");
167        assert_eq!(
168            SeededContainerLoot::read(&mut Cursor::new(network.as_slice()))
169                .expect("container loot should decode"),
170            loot
171        );
172    }
173
174    #[test]
175    fn omitted_seed_defaults_to_zero_without_registry_membership_validation() {
176        let mut compound = NbtCompound::new();
177        compound.insert("loot_table", "steel:unknown_table");
178        assert_eq!(
179            parse(NbtTag::Compound(compound)),
180            Some(SeededContainerLoot::new(
181                Identifier::new_static("steel", "unknown_table"),
182                0,
183            ))
184        );
185    }
186}