Skip to main content

steel_utils/nbt/
mod.rs

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