steel_registry/data_components/components/
custom_model_data.rs

1//! Vanilla `minecraft:custom_model_data` item component.
2
3use std::io::{Cursor, Error, Read, Result, Write};
4
5use simdnbt::owned::{NbtCompound, NbtList, NbtTag};
6use simdnbt::{FromNbtTag, ToNbtTag};
7use steel_utils::codec::VarInt;
8use steel_utils::hash::{ComponentHasher, HashComponent};
9use steel_utils::nbt::{NbtNumeric as _, nbt_collection_values};
10use steel_utils::serial::{ReadFrom, WriteTo};
11
12use super::rgb_color::decode_rgb_color;
13
14const MAX_NETWORK_STRING_LENGTH: usize = 32_767;
15const MAX_NETWORK_STRING_BYTES: usize = MAX_NETWORK_STRING_LENGTH * 3;
16const MAX_INITIAL_LIST_CAPACITY: usize = 65_536;
17
18/// Values exposed to item models and tint sources through custom model data.
19#[derive(Debug, Clone)]
20pub struct CustomModelData {
21    floats: Vec<f32>,
22    flags: Vec<bool>,
23    strings: Vec<String>,
24    colors: Vec<i32>,
25}
26
27impl CustomModelData {
28    pub const EMPTY: Self = Self {
29        floats: Vec::new(),
30        flags: Vec::new(),
31        strings: Vec::new(),
32        colors: Vec::new(),
33    };
34
35    #[must_use]
36    pub const fn new(
37        floats: Vec<f32>,
38        flags: Vec<bool>,
39        strings: Vec<String>,
40        colors: Vec<i32>,
41    ) -> Self {
42        Self {
43            floats,
44            flags,
45            strings,
46            colors,
47        }
48    }
49
50    #[must_use]
51    pub fn floats(&self) -> &[f32] {
52        &self.floats
53    }
54
55    #[must_use]
56    pub fn flags(&self) -> &[bool] {
57        &self.flags
58    }
59
60    #[must_use]
61    pub fn strings(&self) -> &[String] {
62        &self.strings
63    }
64
65    #[must_use]
66    pub fn colors(&self) -> &[i32] {
67        &self.colors
68    }
69
70    #[must_use]
71    pub fn get_float(&self, index: i32) -> Option<f32> {
72        usize::try_from(index)
73            .ok()
74            .and_then(|index| self.floats.get(index))
75            .copied()
76    }
77
78    #[must_use]
79    pub fn get_boolean(&self, index: i32) -> Option<bool> {
80        usize::try_from(index)
81            .ok()
82            .and_then(|index| self.flags.get(index))
83            .copied()
84    }
85
86    #[must_use]
87    pub fn get_string(&self, index: i32) -> Option<&str> {
88        usize::try_from(index)
89            .ok()
90            .and_then(|index| self.strings.get(index))
91            .map(String::as_str)
92    }
93
94    #[must_use]
95    pub fn get_color(&self, index: i32) -> Option<i32> {
96        usize::try_from(index)
97            .ok()
98            .and_then(|index| self.colors.get(index))
99            .copied()
100    }
101}
102
103impl Default for CustomModelData {
104    fn default() -> Self {
105        Self::EMPTY
106    }
107}
108
109impl PartialEq for CustomModelData {
110    fn eq(&self, other: &Self) -> bool {
111        self.floats.len() == other.floats.len()
112            && self
113                .floats
114                .iter()
115                .zip(&other.floats)
116                .all(|(left, right)| float_equals(*left, *right))
117            && self.flags == other.flags
118            && self.strings == other.strings
119            && self.colors == other.colors
120    }
121}
122
123const fn float_equals(left: f32, right: f32) -> bool {
124    (left.is_nan() && right.is_nan()) || left.to_bits() == right.to_bits()
125}
126
127impl WriteTo for CustomModelData {
128    fn write(&self, writer: &mut impl Write) -> Result<()> {
129        write_count(self.floats.len(), writer)?;
130        for value in &self.floats {
131            value.write(writer)?;
132        }
133
134        write_count(self.flags.len(), writer)?;
135        for value in &self.flags {
136            value.write(writer)?;
137        }
138
139        write_count(self.strings.len(), writer)?;
140        for value in &self.strings {
141            write_network_string(value, writer)?;
142        }
143
144        write_count(self.colors.len(), writer)?;
145        for value in &self.colors {
146            value.write(writer)?;
147        }
148        Ok(())
149    }
150}
151
152impl ReadFrom for CustomModelData {
153    fn read(data: &mut Cursor<&[u8]>) -> Result<Self> {
154        let floats = read_list(data, f32::read)?;
155        let flags = read_list(data, bool::read)?;
156        let strings = read_list(data, read_network_string)?;
157        let colors = read_list(data, i32::read)?;
158        Ok(Self::new(floats, flags, strings, colors))
159    }
160}
161
162fn write_count(count: usize, writer: &mut impl Write) -> Result<()> {
163    let count = i32::try_from(count).map_err(|_| Error::other("List is too long"))?;
164    VarInt(count).write(writer)
165}
166
167fn read_list<T>(
168    data: &mut Cursor<&[u8]>,
169    mut read_value: impl FnMut(&mut Cursor<&[u8]>) -> Result<T>,
170) -> Result<Vec<T>> {
171    let count = VarInt::read(data)?.0;
172    let count = usize::try_from(count).map_err(|_| Error::other("Negative list length"))?;
173    let mut values = Vec::with_capacity(count.min(MAX_INITIAL_LIST_CAPACITY));
174    for _ in 0..count {
175        values.push(read_value(data)?);
176    }
177    Ok(values)
178}
179
180fn write_network_string(value: &str, writer: &mut impl Write) -> Result<()> {
181    if value.encode_utf16().count() > MAX_NETWORK_STRING_LENGTH {
182        return Err(Error::other("String is longer than 32767 UTF-16 units"));
183    }
184    if value.len() > MAX_NETWORK_STRING_BYTES {
185        return Err(Error::other("Encoded string is longer than 98301 bytes"));
186    }
187    write_count(value.len(), writer)?;
188    writer.write_all(value.as_bytes())
189}
190
191fn read_network_string(data: &mut Cursor<&[u8]>) -> Result<String> {
192    let byte_count = VarInt::read(data)?.0;
193    let byte_count =
194        usize::try_from(byte_count).map_err(|_| Error::other("Negative encoded string length"))?;
195    if byte_count > MAX_NETWORK_STRING_BYTES {
196        return Err(Error::other("Encoded string is longer than 98301 bytes"));
197    }
198
199    let mut bytes = vec![0; byte_count];
200    data.read_exact(&mut bytes)?;
201    let value = String::from_utf8_lossy(&bytes).into_owned();
202    if value.encode_utf16().count() > MAX_NETWORK_STRING_LENGTH {
203        return Err(Error::other("String is longer than 32767 UTF-16 units"));
204    }
205    Ok(value)
206}
207
208impl ToNbtTag for CustomModelData {
209    fn to_nbt_tag(self) -> NbtTag {
210        let mut compound = NbtCompound::new();
211        if !self.floats.is_empty() {
212            compound.insert("floats", NbtList::Float(self.floats));
213        }
214        if !self.flags.is_empty() {
215            compound.insert(
216                "flags",
217                NbtList::Byte(self.flags.into_iter().map(i8::from).collect()),
218            );
219        }
220        if !self.strings.is_empty() {
221            compound.insert(
222                "strings",
223                NbtList::String(self.strings.into_iter().map(Into::into).collect()),
224            );
225        }
226        if !self.colors.is_empty() {
227            compound.insert("colors", NbtList::Int(self.colors));
228        }
229        NbtTag::Compound(compound)
230    }
231}
232
233impl FromNbtTag for CustomModelData {
234    fn from_nbt_tag(tag: simdnbt::borrow::NbtTag<'_, '_>) -> Option<Self> {
235        let compound = tag.compound()?;
236        let floats = decode_collection(compound.get("floats"), NbtTag::codec_f32)?;
237        let flags = decode_collection(compound.get("flags"), NbtTag::codec_bool)?;
238        let strings = decode_collection(compound.get("strings"), decode_string)?;
239        let colors = decode_collection(compound.get("colors"), decode_rgb_color)?;
240        Some(Self::new(floats, flags, strings, colors))
241    }
242}
243
244fn decode_collection<T>(
245    tag: Option<simdnbt::borrow::NbtTag<'_, '_>>,
246    decode_value: impl Fn(&NbtTag) -> Option<T>,
247) -> Option<Vec<T>> {
248    let Some(tag) = tag else {
249        return Some(Vec::new());
250    };
251    nbt_collection_values(&tag.to_owned())?
252        .iter()
253        .map(decode_value)
254        .collect()
255}
256
257fn decode_string(tag: &NbtTag) -> Option<String> {
258    let NbtTag::String(value) = tag else {
259        return None;
260    };
261    value.to_owned().try_into_string().ok()
262}
263
264impl HashComponent for CustomModelData {
265    fn hash_component(&self, hasher: &mut ComponentHasher) {
266        self.clone().to_nbt_tag().hash_component(hasher);
267    }
268}
269
270#[cfg(test)]
271mod tests {
272    use std::io::Cursor;
273
274    use simdnbt::borrow::read_tag;
275    use simdnbt::owned::{NbtCompound, NbtList, NbtTag};
276    use simdnbt::{FromNbtTag as _, ToNbtTag as _};
277    use steel_utils::codec::VarInt;
278    use steel_utils::hash::HashComponent as _;
279    use steel_utils::serial::{ReadFrom as _, WriteTo as _};
280
281    use super::CustomModelData;
282
283    fn parse(tag: NbtTag) -> Option<CustomModelData> {
284        let mut bytes = Vec::new();
285        tag.write(&mut bytes);
286        let borrowed = read_tag(&mut Cursor::new(bytes.as_slice())).ok()?;
287        CustomModelData::from_nbt_tag(borrowed.as_tag())
288    }
289
290    #[test]
291    fn persistent_codec_omits_default_lists() {
292        assert_eq!(
293            CustomModelData::default().to_nbt_tag(),
294            NbtTag::Compound(NbtCompound::new())
295        );
296        assert_eq!(
297            parse(NbtTag::Compound(NbtCompound::new())),
298            Some(CustomModelData::EMPTY)
299        );
300    }
301
302    #[test]
303    fn persistent_codec_accepts_numeric_collections_and_rgb_vectors() {
304        let mut compound = NbtCompound::new();
305        compound.insert("floats", NbtTag::IntArray(vec![1, 2]));
306        compound.insert("flags", NbtTag::ByteArray(vec![0, 2]));
307        compound.insert("strings", NbtList::String(vec!["steel".into()]));
308        compound.insert(
309            "colors",
310            NbtList::List(vec![NbtList::Float(vec![1.0, 0.5, 0.0])]),
311        );
312
313        let value = parse(NbtTag::Compound(compound)).expect("component should decode");
314        assert_eq!(value.floats(), &[1.0, 2.0]);
315        assert_eq!(value.flags(), &[false, true]);
316        assert_eq!(value.strings(), &["steel"]);
317        assert_eq!(value.colors(), &[0xffff_7f00_u32 as i32]);
318    }
319
320    #[test]
321    fn network_codec_round_trips_all_lists() {
322        let value = CustomModelData::new(
323            vec![1.25, -0.0],
324            vec![true, false],
325            vec!["steel".to_owned(), "🦀".to_owned()],
326            vec![0x123456, -1],
327        );
328        let mut encoded = Vec::new();
329        value.write(&mut encoded).expect("component should encode");
330        assert_eq!(
331            CustomModelData::read(&mut Cursor::new(encoded.as_slice()))
332                .expect("component should decode"),
333            value
334        );
335    }
336
337    #[test]
338    fn network_codec_rejects_negative_counts_and_long_strings() {
339        let mut negative_count = Vec::new();
340        VarInt(-1)
341            .write(&mut negative_count)
342            .expect("count should encode");
343        assert!(CustomModelData::read(&mut Cursor::new(negative_count.as_slice())).is_err());
344
345        let value =
346            CustomModelData::new(Vec::new(), Vec::new(), vec!["a".repeat(32_768)], Vec::new());
347        assert!(value.write(&mut Vec::new()).is_err());
348    }
349
350    #[test]
351    fn equality_matches_java_float_rules_and_getters_are_safe() {
352        let left = CustomModelData::new(
353            vec![f32::from_bits(0x7fc0_0001), 0.0],
354            vec![true],
355            vec!["value".to_owned()],
356            vec![7],
357        );
358        let same = CustomModelData::new(
359            vec![f32::from_bits(0x7fc0_0002), 0.0],
360            vec![true],
361            vec!["value".to_owned()],
362            vec![7],
363        );
364        let negative_zero = CustomModelData::new(
365            vec![f32::NAN, -0.0],
366            vec![true],
367            vec!["value".to_owned()],
368            vec![7],
369        );
370
371        assert_eq!(left, same);
372        assert_ne!(left, negative_zero);
373        assert_eq!(left.get_float(-1), None);
374        assert_eq!(left.get_boolean(0), Some(true));
375        assert_eq!(left.get_string(0), Some("value"));
376        assert_eq!(left.get_color(1), None);
377    }
378
379    #[test]
380    fn persistent_hash_uses_the_record_codec_shape() {
381        let value = CustomModelData::new(
382            vec![1.25],
383            vec![true],
384            vec!["steel".to_owned()],
385            vec![0x123456],
386        );
387
388        assert_eq!(
389            value.compute_hash(),
390            value.clone().to_nbt_tag().compute_hash()
391        );
392    }
393}
steel_registry/data_components/components/custom_model_data.rs

steel_registry/data_components/components/
custom_model_data.rs
