Skip to main content

steel_utils/nbt/
mod.rs

1//! Vanilla-compatible NBT helpers.
2
3mod codec;
4mod path;
5mod snbt;
6mod value_input;
7
8use rustc_hash::FxHashSet;
9use simdnbt::owned::{NbtCompound, NbtList, NbtTag};
10
11pub use codec::NbtNumeric;
12pub use path::{
13    NbtPath, NbtPathError, NbtPathErrorKind, NbtPathMutationError, parse_nbt_path,
14    parse_nbt_path_argument,
15};
16pub use snbt::{
17    SnbtError, SnbtErrorKind, SnbtNumberType, parse_snbt, parse_snbt_argument, parse_snbt_compound,
18    parse_snbt_compound_argument, to_canonical_snbt,
19};
20pub use value_input::{NbtValueInput, NumericTag};
21
22/// Compares two NBT tags for equality, optionally allowing partial list matches.
23#[must_use]
24pub fn compare_nbt(
25    expected: Option<&NbtTag>,
26    actual: Option<&NbtTag>,
27    partial_list_matches: bool,
28) -> bool {
29    let Some(expected) = expected else {
30        return true;
31    };
32    let Some(actual) = actual else {
33        return false;
34    };
35
36    match (expected, actual) {
37        (NbtTag::Compound(expected), NbtTag::Compound(actual)) if partial_list_matches => {
38            compare_nbt_compounds(expected, actual, partial_list_matches)
39        }
40        (NbtTag::List(expected), NbtTag::List(actual)) if partial_list_matches => {
41            compare_lists_partially(expected, actual)
42        }
43        _ => nbt_tags_equal(expected, actual),
44    }
45}
46
47/// Compares NBT values with Vanilla's `Tag.equals` semantics.
48///
49/// Compounds are maps rather than ordered entry lists. Floating-point tags
50/// canonicalize NaNs while still distinguishing positive and negative zero.
51#[must_use]
52pub fn nbt_tags_equal(left: &NbtTag, right: &NbtTag) -> bool {
53    match (left, right) {
54        (NbtTag::Byte(left), NbtTag::Byte(right)) => left == right,
55        (NbtTag::Short(left), NbtTag::Short(right)) => left == right,
56        (NbtTag::Int(left), NbtTag::Int(right)) => left == right,
57        (NbtTag::Long(left), NbtTag::Long(right)) => left == right,
58        (NbtTag::Float(left), NbtTag::Float(right)) => float_tags_equal(*left, *right),
59        (NbtTag::Double(left), NbtTag::Double(right)) => double_tags_equal(*left, *right),
60        (NbtTag::ByteArray(left), NbtTag::ByteArray(right)) => left == right,
61        (NbtTag::String(left), NbtTag::String(right)) => left == right,
62        (NbtTag::List(left), NbtTag::List(right)) => {
63            let left = nbt_list_values(left);
64            let right = nbt_list_values(right);
65            left.len() == right.len()
66                && left
67                    .iter()
68                    .zip(&right)
69                    .all(|(left, right)| nbt_tags_equal(left, right))
70        }
71        (NbtTag::Compound(left), NbtTag::Compound(right)) => nbt_compounds_equal(left, right),
72        (NbtTag::IntArray(left), NbtTag::IntArray(right)) => left == right,
73        (NbtTag::LongArray(left), NbtTag::LongArray(right)) => left == right,
74        _ => false,
75    }
76}
77
78/// Compares compounds as Vanilla maps, independent of serialized entry order.
79#[must_use]
80pub fn nbt_compounds_equal(left: &NbtCompound, right: &NbtCompound) -> bool {
81    left.len() == right.len()
82        && left.iter().all(|(key, left_value)| {
83            right
84                .get(&key.to_str())
85                .is_some_and(|right_value| nbt_tags_equal(left_value, right_value))
86        })
87}
88
89const fn float_tags_equal(left: f32, right: f32) -> bool {
90    (left.is_nan() && right.is_nan()) || left.to_bits() == right.to_bits()
91}
92
93const fn double_tags_equal(left: f64, right: f64) -> bool {
94    (left.is_nan() && right.is_nan()) || left.to_bits() == right.to_bits()
95}
96
97/// Converts an NBT compound to Vanilla's map representation.
98///
99/// Vanilla rejects malformed modified UTF-8 and keeps the final occurrence of
100/// a duplicate compound key. `simdnbt` preserves raw strings and duplicates,
101/// so codecs crossing a trust boundary normalize them explicitly.
102#[must_use]
103pub fn normalize_nbt_compound(compound: NbtCompound) -> Option<NbtCompound> {
104    let mut normalized = NbtCompound::new();
105    for (key, value) in compound {
106        let key = key.try_into_string().ok()?;
107        let value = normalize_nbt_tag(value)?;
108        while normalized.remove(&key).is_some() {}
109        normalized.insert(key, value);
110    }
111    Some(normalized)
112}
113
114/// Converts an NBT value and all descendants to Vanilla's canonical in-memory
115/// representation.
116pub fn normalize_nbt_tag(tag: NbtTag) -> Option<NbtTag> {
117    match tag {
118        NbtTag::String(value) => Some(NbtTag::String(value.try_into_string().ok()?.into())),
119        NbtTag::List(list) => normalize_nbt_list(list).map(NbtTag::List),
120        NbtTag::Compound(compound) => normalize_nbt_compound(compound).map(NbtTag::Compound),
121        tag => Some(tag),
122    }
123}
124
125fn normalize_nbt_list(list: NbtList) -> Option<NbtList> {
126    match list {
127        NbtList::String(values) => values
128            .into_iter()
129            .map(|value| value.try_into_string().ok().map(Into::into))
130            .collect::<Option<Vec<_>>>()
131            .map(NbtList::String),
132        NbtList::List(values) => values
133            .into_iter()
134            .map(normalize_nbt_list)
135            .collect::<Option<Vec<_>>>()
136            .map(NbtList::List),
137        NbtList::Compound(values) => values
138            .into_iter()
139            .map(normalize_nbt_compound)
140            .collect::<Option<Vec<_>>>()
141            .map(NbtList::Compound),
142        list => Some(list),
143    }
144}
145
146/// Recursively applies Vanilla `CompoundTag.merge` semantics.
147pub fn merge_nbt_compounds(target: &mut NbtCompound, source: &NbtCompound) {
148    for (key, source_value) in source.iter() {
149        let key = key.to_str();
150        if let NbtTag::Compound(source_compound) = source_value
151            && let Some(NbtTag::Compound(target_compound)) = target.get_mut(&key)
152        {
153            merge_nbt_compounds(target_compound, source_compound);
154            continue;
155        }
156
157        while target.remove(&key).is_some() {}
158        target.insert(key.into_owned(), source_value.clone());
159    }
160}
161
162/// Returns the heap usage charged by Vanilla's `NbtAccounter` while decoding.
163///
164/// `None` indicates malformed modified UTF-8 or arithmetic overflow.
165#[must_use]
166pub fn vanilla_nbt_heap_size(tag: &NbtTag) -> Option<u64> {
167    match tag {
168        NbtTag::Byte(_) => Some(9),
169        NbtTag::Short(_) => Some(10),
170        NbtTag::Int(_) | NbtTag::Float(_) => Some(12),
171        NbtTag::Long(_) | NbtTag::Double(_) => Some(16),
172        NbtTag::ByteArray(values) => sized_array(24, 1, values.len()),
173        NbtTag::String(value) => sized_string(36, value),
174        NbtTag::List(values) => vanilla_nbt_list_heap_size(values),
175        NbtTag::Compound(values) => vanilla_nbt_compound_heap_size(values),
176        NbtTag::IntArray(values) => sized_array(24, 4, values.len()),
177        NbtTag::LongArray(values) => sized_array(24, 8, values.len()),
178    }
179}
180
181fn vanilla_nbt_list_heap_size(list: &NbtList) -> Option<u64> {
182    let values = list.as_nbt_tags();
183    let count = u64::try_from(values.len()).ok()?;
184    let mut size = 36_u64.checked_add(4_u64.checked_mul(count)?)?;
185    for value in &values {
186        size = size.checked_add(vanilla_nbt_heap_size(value)?)?;
187    }
188    Some(size)
189}
190
191fn vanilla_nbt_compound_heap_size(compound: &NbtCompound) -> Option<u64> {
192    let mut size = 48_u64;
193    let mut keys = FxHashSet::default();
194    for (key, value) in compound.iter() {
195        let key = key.to_owned().try_into_string().ok()?;
196        let key_units = u64::try_from(key.encode_utf16().count()).ok()?;
197        size = size.checked_add(28_u64.checked_add(2_u64.checked_mul(key_units)?)?)?;
198        size = size.checked_add(vanilla_nbt_heap_size(value)?)?;
199        if keys.insert(key) {
200            size = size.checked_add(36)?;
201        }
202    }
203    Some(size)
204}
205
206fn sized_string(base: u64, value: &simdnbt::Mutf8Str) -> Option<u64> {
207    let value = value.to_owned().try_into_string().ok()?;
208    let units = u64::try_from(value.encode_utf16().count()).ok()?;
209    base.checked_add(2_u64.checked_mul(units)?)
210}
211
212fn sized_array(base: u64, element_size: u64, len: usize) -> Option<u64> {
213    let len = u64::try_from(len).ok()?;
214    base.checked_add(element_size.checked_mul(len)?)
215}
216
217/// Compares two compounds with vanilla's partial compound semantics.
218#[must_use]
219pub fn compare_nbt_compounds(
220    expected: &NbtCompound,
221    actual: &NbtCompound,
222    partial_list_matches: bool,
223) -> bool {
224    if actual.len() < expected.len() {
225        return false;
226    }
227
228    expected.iter().all(|(key, expected_tag)| {
229        compare_nbt(
230            Some(expected_tag),
231            actual.get(&key.to_str()),
232            partial_list_matches,
233        )
234    })
235}
236
237fn compare_lists_partially(expected: &NbtList, actual: &NbtList) -> bool {
238    let expected = nbt_list_values(expected);
239    let actual = nbt_list_values(actual);
240    if expected.is_empty() {
241        return actual.is_empty();
242    }
243    if actual.len() < expected.len() {
244        return false;
245    }
246
247    expected.iter().all(|expected_tag| {
248        actual
249            .iter()
250            .any(|actual_tag| compare_nbt(Some(expected_tag), Some(actual_tag), true))
251    })
252}
253
254/// Returns the semantic values of a Vanilla list.
255///
256/// Vanilla permits heterogeneous lists and wraps non-compound elements
257/// when serializing them through the homogeneous binary NBT format.
258#[must_use]
259pub fn nbt_list_values(list: &NbtList) -> Vec<NbtTag> {
260    list.as_nbt_tags()
261        .into_iter()
262        .map(unwrap_list_wrapper)
263        .collect()
264}
265
266/// Returns values exposed by Vanilla's `NbtOps` collection interface.
267///
268/// Numeric arrays are collections as well as ordinary list tags, so codecs
269/// built with `Codec.listOf()` accept all four representations.
270#[must_use]
271pub fn nbt_collection_values(tag: &NbtTag) -> Option<Vec<NbtTag>> {
272    match tag {
273        NbtTag::ByteArray(values) => Some(
274            values
275                .iter()
276                .map(|value| NbtTag::Byte(i8::from_ne_bytes([*value])))
277                .collect(),
278        ),
279        NbtTag::List(values) => Some(nbt_list_values(values)),
280        NbtTag::IntArray(values) => Some(values.iter().copied().map(NbtTag::Int).collect()),
281        NbtTag::LongArray(values) => Some(values.iter().copied().map(NbtTag::Long).collect()),
282        _ => None,
283    }
284}
285
286fn unwrap_list_wrapper(tag: NbtTag) -> NbtTag {
287    match tag {
288        NbtTag::Compound(mut compound) if compound.len() == 1 && compound.contains("") => {
289            let Some(value) = compound.take("") else {
290                return NbtTag::Compound(compound);
291            };
292            value
293        }
294        tag => tag,
295    }
296}
297
298#[cfg(test)]
299mod tests {
300    use super::*;
301
302    fn compound(entries: impl IntoIterator<Item = (&'static str, NbtTag)>) -> NbtTag {
303        let mut compound = NbtCompound::new();
304        for (key, tag) in entries {
305            compound.insert(key, tag);
306        }
307        NbtTag::Compound(compound)
308    }
309
310    fn list(entries: impl IntoIterator<Item = NbtTag>) -> NbtTag {
311        NbtTag::List(NbtList::from(entries.into_iter().collect::<Vec<_>>()))
312    }
313
314    #[test]
315    fn compounds_and_lists_match_partially() {
316        let expected = compound([(
317            "values",
318            list([compound([("name", NbtTag::String("second".into()))])]),
319        )]);
320        let actual = compound([(
321            "values",
322            list([
323                compound([("name", NbtTag::String("first".into()))]),
324                compound([
325                    ("name", NbtTag::String("second".into())),
326                    ("extra", NbtTag::Byte(1)),
327                ]),
328            ]),
329        )]);
330
331        assert!(compare_nbt(Some(&expected), Some(&actual), true));
332        assert!(!compare_nbt(Some(&expected), Some(&actual), false));
333    }
334
335    #[test]
336    fn empty_partial_list_only_matches_an_empty_list() {
337        let empty = list([]);
338        let non_empty = list([NbtTag::Int(1)]);
339
340        assert!(compare_nbt(Some(&empty), Some(&empty), true));
341        assert!(!compare_nbt(Some(&empty), Some(&non_empty), true));
342    }
343
344    #[test]
345    fn partial_lists_match_heterogeneous_values() {
346        let expected = list([NbtTag::String("two".into())]);
347        let actual = list([NbtTag::Int(1), NbtTag::String("two".into())]);
348
349        assert!(compare_nbt(Some(&expected), Some(&actual), true));
350    }
351
352    #[test]
353    fn scalar_tags_require_the_same_nbt_type() {
354        assert!(compare_nbt(
355            Some(&NbtTag::Int(1)),
356            Some(&NbtTag::Int(1)),
357            true
358        ));
359        assert!(!compare_nbt(
360            Some(&NbtTag::Int(1)),
361            Some(&NbtTag::Long(1)),
362            true
363        ));
364    }
365
366    #[test]
367    fn exact_equality_matches_vanilla_map_and_float_rules() {
368        let left = compound([
369            ("first", NbtTag::Int(1)),
370            ("second", NbtTag::Float(f32::from_bits(0x7fc0_0001))),
371        ]);
372        let right = compound([
373            ("second", NbtTag::Float(f32::from_bits(0x7fc0_0002))),
374            ("first", NbtTag::Int(1)),
375        ]);
376
377        assert!(nbt_tags_equal(&left, &right));
378        assert!(!nbt_tags_equal(&NbtTag::Float(0.0), &NbtTag::Float(-0.0)));
379    }
380
381    #[test]
382    fn normalization_keeps_the_final_duplicate_key() {
383        let mut raw = NbtCompound::new();
384        raw.insert("value", 1);
385        raw.insert("value", 2);
386
387        let normalized = normalize_nbt_compound(raw).expect("valid compound should normalize");
388        assert_eq!(normalized.len(), 1);
389        assert_eq!(normalized.int("value"), Some(2));
390    }
391
392    #[test]
393    fn compound_merge_recurses_only_when_both_values_are_compounds() {
394        let mut nested = NbtCompound::new();
395        nested.insert("kept", 1);
396        nested.insert("replaced", 1);
397        let mut target = NbtCompound::new();
398        target.insert("nested", nested);
399        target.insert("scalar", 1);
400
401        let mut nested = NbtCompound::new();
402        nested.insert("replaced", 2);
403        let mut source = NbtCompound::new();
404        source.insert("nested", nested);
405        source.insert("scalar", NbtCompound::new());
406
407        merge_nbt_compounds(&mut target, &source);
408        let nested = target
409            .compound("nested")
410            .expect("nested compound should remain");
411        assert_eq!(nested.int("kept"), Some(1));
412        assert_eq!(nested.int("replaced"), Some(2));
413        assert!(target.compound("scalar").is_some());
414    }
415
416    #[test]
417    fn heap_size_matches_vanilla_nbt_accounter_formulas() {
418        let mut raw = NbtCompound::new();
419        raw.insert("a", 1);
420        raw.insert("a", 2);
421
422        assert_eq!(vanilla_nbt_heap_size(&NbtTag::Compound(raw)), Some(168));
423    }
424}