Skip to main content

azalea_buf/impls/
extra.rs

1use std::{
2    collections::HashMap,
3    hash::Hash,
4    io::{self, Cursor, Write},
5    sync::Arc,
6};
7
8use indexmap::IndexMap;
9
10use crate::{
11    AzBuf, AzBufLimited, AzBufVar, BufReadError, MAX_STRING_LENGTH, UnsizedByteArray, read_bytes,
12    read_utf_with_len, write_utf_with_len,
13};
14
15impl AzBuf for UnsizedByteArray {
16    fn azalea_read(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
17        // read to end of the buffer
18        let data = buf.get_ref()[buf.position() as usize..].to_vec();
19        buf.set_position((buf.position()) + data.len() as u64);
20        Ok(UnsizedByteArray(data))
21    }
22    fn azalea_write(&self, buf: &mut impl Write) -> io::Result<()> {
23        buf.write_all(self)
24    }
25}
26
27macro_rules! impl_for_map_type {
28    ($ty: ident) => {
29        impl<K: AzBuf + Eq + Hash, V: AzBuf> AzBuf for $ty<K, V> {
30            fn azalea_read(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
31                let length = i32::azalea_read_var(buf)? as usize;
32                let mut contents = Self::with_capacity(usize::min(length, 65536));
33                for _ in 0..length {
34                    contents.insert(K::azalea_read(buf)?, V::azalea_read(buf)?);
35                }
36                Ok(contents)
37            }
38            fn azalea_write(&self, buf: &mut impl Write) -> io::Result<()> {
39                u32::azalea_write_var(&(self.len() as u32), buf)?;
40                for (key, value) in self {
41                    key.azalea_write(buf)?;
42                    value.azalea_write(buf)?;
43                }
44
45                Ok(())
46            }
47        }
48        impl<K: AzBuf + Eq + Hash, V: AzBufVar> AzBufVar for $ty<K, V> {
49            fn azalea_read_var(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
50                let length = i32::azalea_read_var(buf)? as usize;
51                let mut contents = Self::with_capacity(usize::min(length, 65536));
52                for _ in 0..length {
53                    contents.insert(K::azalea_read(buf)?, V::azalea_read_var(buf)?);
54                }
55                Ok(contents)
56            }
57            fn azalea_write_var(&self, buf: &mut impl Write) -> io::Result<()> {
58                u32::azalea_write_var(&(self.len() as u32), buf)?;
59                for (key, value) in self {
60                    key.azalea_write(buf)?;
61                    value.azalea_write_var(buf)?;
62                }
63
64                Ok(())
65            }
66        }
67    };
68}
69
70impl_for_map_type!(HashMap);
71impl_for_map_type!(IndexMap);
72
73macro_rules! impl_for_list_type {
74    ($ty: ty) => {
75        impl<T: AzBuf> AzBuf for $ty {
76            default fn azalea_read(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
77                let length = u32::azalea_read_var(buf)? as usize;
78                // we limit the capacity to not get exploited into allocating a bunch
79                let mut contents = Vec::with_capacity(usize::min(length, 65536));
80                for _ in 0..length {
81                    contents.push(T::azalea_read(buf)?);
82                }
83                Ok(contents.into())
84            }
85            default fn azalea_write(&self, buf: &mut impl Write) -> io::Result<()> {
86                (self.len() as u32).azalea_write_var(buf)?;
87                for item in self.iter() {
88                    T::azalea_write(item, buf)?;
89                }
90                Ok(())
91            }
92        }
93        impl<T: AzBufVar> AzBufVar for $ty {
94            fn azalea_read_var(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
95                let length = i32::azalea_read_var(buf)? as usize;
96                let mut contents = Vec::with_capacity(usize::min(length, 65536));
97                for _ in 0..length {
98                    contents.push(T::azalea_read_var(buf)?);
99                }
100                Ok(contents.into())
101            }
102            fn azalea_write_var(&self, buf: &mut impl Write) -> io::Result<()> {
103                (self.len() as u32).azalea_write_var(buf)?;
104                for item in self.iter() {
105                    T::azalea_write_var(item, buf)?;
106                }
107                Ok(())
108            }
109        }
110        impl<T: AzBuf> AzBufLimited for $ty {
111            fn azalea_read_limited(
112                buf: &mut Cursor<&[u8]>,
113                limit: u32,
114            ) -> Result<Self, BufReadError> {
115                let length = u32::azalea_read_var(buf)?;
116                if length > limit {
117                    return Err(BufReadError::VecLengthTooLong {
118                        length: length as u32,
119                        max_length: limit as u32,
120                    });
121                }
122
123                let mut contents = Vec::with_capacity(u32::min(length, 65536) as usize);
124                for _ in 0..length {
125                    contents.push(T::azalea_read(buf)?);
126                }
127                Ok(contents.into())
128            }
129        }
130    };
131}
132
133impl_for_list_type!(Vec<T>);
134impl_for_list_type!(Box<[T]>);
135// `Arc<[T]>` is deliberately not implemented here, because converting a
136// `Vec<T>` to `Arc<[T]>` results in allocations (since `Arc` stores the
137// counters on the heap) and we'd like to avoid that. for us it's typically
138// better to do `Arc<Box<[T]>>`.
139
140impl AzBuf for Vec<u8> {
141    fn azalea_read(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
142        let length = i32::azalea_read_var(buf)? as usize;
143        read_bytes(buf, length).map(|b| b.to_vec())
144    }
145    fn azalea_write(&self, buf: &mut impl Write) -> io::Result<()> {
146        (self.len() as u32).azalea_write_var(buf)?;
147        buf.write_all(self)
148    }
149}
150
151impl AzBuf for String {
152    fn azalea_read(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
153        read_utf_with_len(buf, MAX_STRING_LENGTH).map(Into::into)
154    }
155    fn azalea_write(&self, buf: &mut impl Write) -> io::Result<()> {
156        write_utf_with_len(buf, self, MAX_STRING_LENGTH)
157    }
158}
159impl AzBufLimited for String {
160    fn azalea_read_limited(buf: &mut Cursor<&[u8]>, limit: u32) -> Result<Self, BufReadError> {
161        read_utf_with_len(buf, limit).map(Into::into)
162    }
163}
164
165impl AzBuf for Box<str> {
166    fn azalea_read(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
167        read_utf_with_len(buf, MAX_STRING_LENGTH).map(Into::into)
168    }
169    fn azalea_write(&self, buf: &mut impl Write) -> io::Result<()> {
170        write_utf_with_len(buf, self, MAX_STRING_LENGTH)
171    }
172}
173
174impl AzBufLimited for Box<str> {
175    fn azalea_read_limited(buf: &mut Cursor<&[u8]>, limit: u32) -> Result<Self, BufReadError> {
176        String::azalea_read_limited(buf, limit).map(Into::into)
177    }
178}
179
180impl<T: AzBuf> AzBuf for Option<T> {
181    fn azalea_read(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
182        let present = bool::azalea_read(buf)?;
183        Ok(if present {
184            Some(T::azalea_read(buf)?)
185        } else {
186            None
187        })
188    }
189    fn azalea_write(&self, buf: &mut impl Write) -> io::Result<()> {
190        if let Some(s) = self {
191            true.azalea_write(buf)?;
192            s.azalea_write(buf)?;
193        } else {
194            false.azalea_write(buf)?;
195        };
196        Ok(())
197    }
198}
199
200impl<T: AzBufVar> AzBufVar for Option<T> {
201    fn azalea_read_var(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
202        let present = bool::azalea_read(buf)?;
203        Ok(if present {
204            Some(T::azalea_read_var(buf)?)
205        } else {
206            None
207        })
208    }
209    fn azalea_write_var(&self, buf: &mut impl Write) -> io::Result<()> {
210        if let Some(s) = self {
211            true.azalea_write(buf)?;
212            s.azalea_write_var(buf)?;
213        } else {
214            false.azalea_write(buf)?;
215        };
216        Ok(())
217    }
218}
219impl<T: AzBufLimited> AzBufLimited for Option<T> {
220    fn azalea_read_limited(buf: &mut Cursor<&[u8]>, limit: u32) -> Result<Self, BufReadError> {
221        let present = bool::azalea_read(buf)?;
222        Ok(if present {
223            Some(T::azalea_read_limited(buf, limit)?)
224        } else {
225            None
226        })
227    }
228}
229
230impl<T: AzBuf, const N: usize> AzBuf for [T; N] {
231    fn azalea_read(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
232        let mut contents = Vec::with_capacity(N);
233        for _ in 0..N {
234            contents.push(T::azalea_read(buf)?);
235        }
236        Ok(contents
237            .try_into()
238            .unwrap_or_else(|_| unreachable!("The vec is the same size as the array")))
239    }
240    fn azalea_write(&self, buf: &mut impl Write) -> io::Result<()> {
241        for i in self {
242            i.azalea_write(buf)?;
243        }
244        Ok(())
245    }
246}
247
248impl AzBuf for simdnbt::owned::NbtTag {
249    fn azalea_read(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
250        Ok(simdnbt::owned::read_tag(buf).map_err(simdnbt::Error::from)?)
251    }
252    fn azalea_write(&self, buf: &mut impl Write) -> io::Result<()> {
253        let mut data = Vec::new();
254        self.write(&mut data);
255        buf.write_all(&data)
256    }
257}
258
259impl AzBuf for simdnbt::owned::NbtCompound {
260    fn azalea_read(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
261        match simdnbt::owned::read_tag(buf).map_err(simdnbt::Error::from)? {
262            simdnbt::owned::NbtTag::Compound(compound) => Ok(compound),
263            _ => Err(BufReadError::Custom("Expected compound tag".to_owned())),
264        }
265    }
266    fn azalea_write(&self, buf: &mut impl Write) -> io::Result<()> {
267        let mut data = Vec::new();
268        simdnbt::owned::NbtTag::Compound(self.clone()).write(&mut data);
269        buf.write_all(&data)
270    }
271}
272
273impl AzBuf for simdnbt::owned::Nbt {
274    fn azalea_read(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
275        Ok(simdnbt::owned::read_unnamed(buf)?)
276    }
277    fn azalea_write(&self, buf: &mut impl Write) -> io::Result<()> {
278        let mut data = Vec::new();
279        self.write_unnamed(&mut data);
280        buf.write_all(&data)
281    }
282}
283
284impl<T> AzBuf for Box<T>
285where
286    T: AzBuf,
287{
288    fn azalea_read(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
289        T::azalea_read(buf).map(Box::new)
290    }
291    fn azalea_write(&self, buf: &mut impl Write) -> io::Result<()> {
292        T::azalea_write(&**self, buf)
293    }
294}
295
296impl<A: AzBuf, B: AzBuf> AzBuf for (A, B) {
297    fn azalea_read(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
298        Ok((A::azalea_read(buf)?, B::azalea_read(buf)?))
299    }
300    fn azalea_write(&self, buf: &mut impl Write) -> io::Result<()> {
301        self.0.azalea_write(buf)?;
302        self.1.azalea_write(buf)
303    }
304}
305
306impl<T: AzBuf> AzBuf for Arc<T> {
307    fn azalea_read(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
308        Ok(Arc::new(T::azalea_read(buf)?))
309    }
310    fn azalea_write(&self, buf: &mut impl Write) -> io::Result<()> {
311        T::azalea_write(&**self, buf)
312    }
313}