valence_nbt/binary/
encode.rs

1use std::borrow::Cow;
2use std::hash::Hash;
3use std::io::Write;
4
5use byteorder::{BigEndian, WriteBytesExt};
6
7use super::modified_utf8;
8use crate::conv::i8_slice_as_u8_slice;
9use crate::tag::Tag;
10use crate::{Compound, Error, List, Result, Value};
11
12/// Encodes uncompressed NBT binary data to the provided writer.
13///
14/// Only compounds are permitted at the top level. This is why the function
15/// accepts a [`Compound`] reference rather than a [`Value`].
16///
17/// Additionally, the root compound can be given a name. Typically the empty
18/// string `""` is used.
19pub fn to_binary<W, S, R>(comp: &Compound<S>, writer: W, root_name: &R) -> Result<()>
20where
21    W: Write,
22    S: ToModifiedUtf8 + Hash + Ord,
23    R: ToModifiedUtf8 + ?Sized,
24{
25    let mut state = EncodeState { writer };
26
27    state.write_tag(Tag::Compound)?;
28    state.write_string(root_name)?;
29    state.write_compound(comp)?;
30
31    Ok(())
32}
33
34/// Returns the number of bytes that will be written when
35/// [`to_binary`] is called with this compound and root name.
36///
37/// If `to_binary` results in `Ok`, the exact number of bytes
38/// reported by this function will have been written. If the result is
39/// `Err`, then the reported count will be greater than or equal to the
40/// number of bytes that have actually been written.
41pub fn written_size<S, R>(comp: &Compound<S>, root_name: &R) -> usize
42where
43    S: ToModifiedUtf8 + Hash + Ord,
44    R: ToModifiedUtf8 + ?Sized,
45{
46    fn value_size<S>(val: &Value<S>) -> usize
47    where
48        S: ToModifiedUtf8 + Hash + Ord,
49    {
50        match val {
51            Value::Byte(_) => 1,
52            Value::Short(_) => 2,
53            Value::Int(_) => 4,
54            Value::Long(_) => 8,
55            Value::Float(_) => 4,
56            Value::Double(_) => 8,
57            Value::ByteArray(v) => 4 + v.len(),
58            Value::String(v) => string_size(v),
59            Value::List(v) => list_size(v),
60            Value::Compound(v) => compound_size(v),
61            Value::IntArray(v) => 4 + v.len() * 4,
62            Value::LongArray(v) => 4 + v.len() * 8,
63        }
64    }
65
66    fn list_size<S>(l: &List<S>) -> usize
67    where
68        S: ToModifiedUtf8 + Hash + Ord,
69    {
70        let elems_size = match l {
71            List::End => 0,
72            List::Byte(v) => v.len(),
73            List::Short(v) => v.len() * 2,
74            List::Int(v) => v.len() * 4,
75            List::Long(v) => v.len() * 8,
76            List::Float(v) => v.len() * 4,
77            List::Double(v) => v.len() * 8,
78            List::ByteArray(v) => v.iter().map(|b| 4 + b.len()).sum(),
79            List::String(v) => v.iter().map(|s| string_size(s)).sum(),
80            List::List(v) => v.iter().map(list_size).sum(),
81            List::Compound(v) => v.iter().map(compound_size).sum(),
82            List::IntArray(v) => v.iter().map(|i| 4 + i.len() * 4).sum(),
83            List::LongArray(v) => v.iter().map(|l| 4 + l.len() * 8).sum(),
84        };
85
86        1 + 4 + elems_size
87    }
88
89    fn string_size<S: ToModifiedUtf8 + ?Sized>(s: &S) -> usize {
90        2 + s.modified_uf8_len()
91    }
92
93    fn compound_size<S>(c: &Compound<S>) -> usize
94    where
95        S: ToModifiedUtf8 + Hash + Ord,
96    {
97        c.iter()
98            .map(|(k, v)| 1 + string_size(k) + value_size(v))
99            .sum::<usize>()
100            + 1
101    }
102
103    1 + string_size(root_name) + compound_size(comp)
104}
105
106struct EncodeState<W> {
107    writer: W,
108}
109
110impl<W: Write> EncodeState<W> {
111    fn write_tag(&mut self, tag: Tag) -> Result<()> {
112        Ok(self.writer.write_u8(tag as u8)?)
113    }
114
115    fn write_value<S>(&mut self, v: &Value<S>) -> Result<()>
116    where
117        S: ToModifiedUtf8 + Hash + Ord,
118    {
119        match v {
120            Value::Byte(v) => self.write_byte(*v),
121            Value::Short(v) => self.write_short(*v),
122            Value::Int(v) => self.write_int(*v),
123            Value::Long(v) => self.write_long(*v),
124            Value::Float(v) => self.write_float(*v),
125            Value::Double(v) => self.write_double(*v),
126            Value::ByteArray(v) => self.write_byte_array(v),
127            Value::String(v) => self.write_string(v),
128            Value::List(v) => self.write_any_list(v),
129            Value::Compound(v) => self.write_compound(v),
130            Value::IntArray(v) => self.write_int_array(v),
131            Value::LongArray(v) => self.write_long_array(v),
132        }
133    }
134
135    fn write_byte(&mut self, byte: i8) -> Result<()> {
136        Ok(self.writer.write_i8(byte)?)
137    }
138
139    fn write_short(&mut self, short: i16) -> Result<()> {
140        Ok(self.writer.write_i16::<BigEndian>(short)?)
141    }
142
143    fn write_int(&mut self, int: i32) -> Result<()> {
144        Ok(self.writer.write_i32::<BigEndian>(int)?)
145    }
146
147    fn write_long(&mut self, long: i64) -> Result<()> {
148        Ok(self.writer.write_i64::<BigEndian>(long)?)
149    }
150
151    fn write_float(&mut self, float: f32) -> Result<()> {
152        Ok(self.writer.write_f32::<BigEndian>(float)?)
153    }
154
155    fn write_double(&mut self, double: f64) -> Result<()> {
156        Ok(self.writer.write_f64::<BigEndian>(double)?)
157    }
158
159    fn write_byte_array(&mut self, bytes: &[i8]) -> Result<()> {
160        match bytes.len().try_into() {
161            Ok(len) => self.write_int(len)?,
162            Err(_) => {
163                return Err(Error::new_owned(format!(
164                    "byte array of length {} exceeds maximum of i32::MAX",
165                    bytes.len(),
166                )))
167            }
168        }
169
170        Ok(self.writer.write_all(i8_slice_as_u8_slice(bytes))?)
171    }
172
173    fn write_string<S: ToModifiedUtf8 + ?Sized>(&mut self, s: &S) -> Result<()> {
174        let len = s.modified_uf8_len();
175
176        match len.try_into() {
177            Ok(n) => self.writer.write_u16::<BigEndian>(n)?,
178            Err(_) => {
179                return Err(Error::new_owned(format!(
180                    "string of length {len} exceeds maximum of u16::MAX"
181                )))
182            }
183        }
184
185        s.to_modified_utf8(len, &mut self.writer)?;
186
187        Ok(())
188    }
189
190    fn write_any_list<S>(&mut self, list: &List<S>) -> Result<()>
191    where
192        S: ToModifiedUtf8 + Hash + Ord,
193    {
194        match list {
195            List::End => {
196                self.write_tag(Tag::End)?;
197                // Length
198                self.writer.write_i32::<BigEndian>(0)?;
199                Ok(())
200            }
201            List::Byte(v) => {
202                self.write_tag(Tag::Byte)?;
203
204                match v.len().try_into() {
205                    Ok(len) => self.write_int(len)?,
206                    Err(_) => {
207                        return Err(Error::new_owned(format!(
208                            "byte list of length {} exceeds maximum of i32::MAX",
209                            v.len(),
210                        )))
211                    }
212                }
213
214                Ok(self.writer.write_all(i8_slice_as_u8_slice(v))?)
215            }
216            List::Short(sl) => self.write_list(sl, Tag::Short, |st, v| st.write_short(*v)),
217            List::Int(il) => self.write_list(il, Tag::Int, |st, v| st.write_int(*v)),
218            List::Long(ll) => self.write_list(ll, Tag::Long, |st, v| st.write_long(*v)),
219            List::Float(fl) => self.write_list(fl, Tag::Float, |st, v| st.write_float(*v)),
220            List::Double(dl) => self.write_list(dl, Tag::Double, |st, v| st.write_double(*v)),
221            List::ByteArray(v) => {
222                self.write_list(v, Tag::ByteArray, |st, v| st.write_byte_array(v))
223            }
224            List::String(v) => self.write_list(v, Tag::String, |st, v| st.write_string(v)),
225            List::List(v) => self.write_list(v, Tag::List, |st, v| st.write_any_list(v)),
226            List::Compound(v) => self.write_list(v, Tag::Compound, |st, v| st.write_compound(v)),
227            List::IntArray(v) => self.write_list(v, Tag::IntArray, |st, v| st.write_int_array(v)),
228            List::LongArray(v) => {
229                self.write_list(v, Tag::LongArray, |st, v| st.write_long_array(v))
230            }
231        }
232    }
233
234    fn write_list<T, F>(&mut self, list: &[T], elem_type: Tag, mut write_elem: F) -> Result<()>
235    where
236        F: FnMut(&mut Self, &T) -> Result<()>,
237    {
238        self.write_tag(elem_type)?;
239
240        match list.len().try_into() {
241            Ok(len) => self.writer.write_i32::<BigEndian>(len)?,
242            Err(_) => {
243                return Err(Error::new_owned(format!(
244                    "{} list of length {} exceeds maximum of i32::MAX",
245                    list.len(),
246                    elem_type.name()
247                )))
248            }
249        }
250
251        for elem in list {
252            write_elem(self, elem)?;
253        }
254
255        Ok(())
256    }
257
258    fn write_compound<S>(&mut self, c: &Compound<S>) -> Result<()>
259    where
260        S: ToModifiedUtf8 + Hash + Ord,
261    {
262        for (k, v) in c {
263            self.write_tag(v.tag())?;
264            self.write_string(k)?;
265            self.write_value(v)?;
266        }
267        self.write_tag(Tag::End)?;
268
269        Ok(())
270    }
271
272    fn write_int_array(&mut self, ia: &[i32]) -> Result<()> {
273        match ia.len().try_into() {
274            Ok(len) => self.write_int(len)?,
275            Err(_) => {
276                return Err(Error::new_owned(format!(
277                    "int array of length {} exceeds maximum of i32::MAX",
278                    ia.len(),
279                )))
280            }
281        }
282
283        for i in ia {
284            self.write_int(*i)?;
285        }
286
287        Ok(())
288    }
289
290    fn write_long_array(&mut self, la: &[i64]) -> Result<()> {
291        match la.len().try_into() {
292            Ok(len) => self.write_int(len)?,
293            Err(_) => {
294                return Err(Error::new_owned(format!(
295                    "long array of length {} exceeds maximum of i32::MAX",
296                    la.len(),
297                )))
298            }
299        }
300
301        for l in la {
302            self.write_long(*l)?;
303        }
304
305        Ok(())
306    }
307}
308
309/// A string type which can be encoded into Java's [modified UTF-8](https://docs.oracle.com/javase/8/docs/api/java/io/DataInput.html#modified-utf-8).
310pub trait ToModifiedUtf8 {
311    fn modified_uf8_len(&self) -> usize;
312    fn to_modified_utf8<W: Write>(&self, encoded_len: usize, writer: W) -> std::io::Result<()>;
313}
314
315impl ToModifiedUtf8 for str {
316    fn modified_uf8_len(&self) -> usize {
317        modified_utf8::encoded_len(self.as_bytes())
318    }
319
320    fn to_modified_utf8<W: Write>(&self, encoded_len: usize, mut writer: W) -> std::io::Result<()> {
321        // Conversion to modified UTF-8 always increases the size of the string.
322        // If the new len is equal to the original len, we know it doesn't need
323        // to be re-encoded.
324        if self.len() == encoded_len {
325            writer.write_all(self.as_bytes())
326        } else {
327            modified_utf8::write_modified_utf8(writer, self)
328        }
329    }
330}
331
332impl ToModifiedUtf8 for Cow<'_, str> {
333    #[inline]
334    fn modified_uf8_len(&self) -> usize {
335        str::modified_uf8_len(self)
336    }
337
338    fn to_modified_utf8<W: Write>(&self, encoded_len: usize, writer: W) -> std::io::Result<()> {
339        str::to_modified_utf8(self, encoded_len, writer)
340    }
341}
342
343impl ToModifiedUtf8 for String {
344    #[inline]
345    fn modified_uf8_len(&self) -> usize {
346        str::modified_uf8_len(self)
347    }
348
349    fn to_modified_utf8<W: Write>(&self, encoded_len: usize, writer: W) -> std::io::Result<()> {
350        str::to_modified_utf8(self, encoded_len, writer)
351    }
352}
353
354#[cfg(feature = "java_string")]
355impl ToModifiedUtf8 for java_string::JavaStr {
356    fn modified_uf8_len(&self) -> usize {
357        modified_utf8::encoded_len(self.as_bytes())
358    }
359
360    fn to_modified_utf8<W: Write>(
361        &self,
362        _encoded_len: usize,
363        mut writer: W,
364    ) -> std::io::Result<()> {
365        writer.write_all(&self.to_modified_utf8())
366    }
367}
368
369#[cfg(feature = "java_string")]
370impl ToModifiedUtf8 for Cow<'_, java_string::JavaStr> {
371    #[inline]
372    fn modified_uf8_len(&self) -> usize {
373        java_string::JavaStr::modified_uf8_len(self)
374    }
375
376    fn to_modified_utf8<W: Write>(&self, encoded_len: usize, writer: W) -> std::io::Result<()> {
377        <java_string::JavaStr as ToModifiedUtf8>::to_modified_utf8(self, encoded_len, writer)
378    }
379}
380
381#[cfg(feature = "java_string")]
382impl ToModifiedUtf8 for java_string::JavaString {
383    #[inline]
384    fn modified_uf8_len(&self) -> usize {
385        java_string::JavaStr::modified_uf8_len(self)
386    }
387
388    fn to_modified_utf8<W: Write>(&self, encoded_len: usize, writer: W) -> std::io::Result<()> {
389        <java_string::JavaStr as ToModifiedUtf8>::to_modified_utf8(self, encoded_len, writer)
390    }
391}