steel_registry/data_components/components/
container_loot.rs1use 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#[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}