1use core::fmt;
8use core::mem::take;
9
10use crate::alloc::{Allocator, Vec};
11use crate::Context;
12
13pub trait Writer {
15 type Mut<'this>: Writer
26 where
27 Self: 'this;
28
29 fn borrow_mut(&mut self) -> Self::Mut<'_>;
31
32 fn extend<C>(&mut self, cx: &C, buffer: Vec<'_, u8, C::Allocator>) -> Result<(), C::Error>
34 where
35 C: ?Sized + Context;
36
37 fn write_bytes<C>(&mut self, cx: &C, bytes: &[u8]) -> Result<(), C::Error>
39 where
40 C: ?Sized + Context;
41
42 #[inline]
44 fn write_byte<C>(&mut self, cx: &C, b: u8) -> Result<(), C::Error>
45 where
46 C: ?Sized + Context,
47 {
48 self.write_bytes(cx, &[b])
49 }
50}
51
52impl<W> Writer for &mut W
53where
54 W: ?Sized + Writer,
55{
56 type Mut<'this> = &'this mut W where Self: 'this;
57
58 #[inline]
59 fn borrow_mut(&mut self) -> Self::Mut<'_> {
60 self
61 }
62
63 #[inline]
64 fn extend<C>(&mut self, cx: &C, buffer: Vec<'_, u8, C::Allocator>) -> Result<(), C::Error>
65 where
66 C: ?Sized + Context,
67 {
68 (*self).extend(cx, buffer)
69 }
70
71 #[inline]
72 fn write_bytes<C>(&mut self, cx: &C, bytes: &[u8]) -> Result<(), C::Error>
73 where
74 C: ?Sized + Context,
75 {
76 (*self).write_bytes(cx, bytes)
77 }
78
79 #[inline]
80 fn write_byte<C>(&mut self, cx: &C, b: u8) -> Result<(), C::Error>
81 where
82 C: ?Sized + Context,
83 {
84 (*self).write_byte(cx, b)
85 }
86}
87
88#[cfg(feature = "alloc")]
89impl Writer for rust_alloc::vec::Vec<u8> {
90 type Mut<'this> = &'this mut Self where Self: 'this;
91
92 #[inline]
93 fn borrow_mut(&mut self) -> Self::Mut<'_> {
94 self
95 }
96
97 #[inline]
98 fn extend<C>(&mut self, cx: &C, buffer: Vec<'_, u8, C::Allocator>) -> Result<(), C::Error>
99 where
100 C: ?Sized + Context,
101 {
102 self.write_bytes(cx, buffer.as_slice())
104 }
105
106 #[inline]
107 fn write_bytes<C>(&mut self, cx: &C, bytes: &[u8]) -> Result<(), C::Error>
108 where
109 C: ?Sized + Context,
110 {
111 self.extend_from_slice(bytes);
112 cx.advance(bytes.len());
113 Ok(())
114 }
115
116 #[inline]
117 fn write_byte<C>(&mut self, cx: &C, b: u8) -> Result<(), C::Error>
118 where
119 C: ?Sized + Context,
120 {
121 self.push(b);
122 cx.advance(1);
123 Ok(())
124 }
125}
126
127impl Writer for &mut [u8] {
128 type Mut<'this> = &'this mut Self where Self: 'this;
129
130 #[inline]
131 fn borrow_mut(&mut self) -> Self::Mut<'_> {
132 self
133 }
134
135 #[inline]
136 fn extend<C>(&mut self, cx: &C, buffer: Vec<'_, u8, C::Allocator>) -> Result<(), C::Error>
137 where
138 C: ?Sized + Context,
139 {
140 self.write_bytes(cx, buffer.as_slice())
142 }
143
144 #[inline]
145 fn write_bytes<C>(&mut self, cx: &C, bytes: &[u8]) -> Result<(), C::Error>
146 where
147 C: ?Sized + Context,
148 {
149 if self.len() < bytes.len() {
150 return Err(cx.message(SliceOverflow {
151 n: bytes.len(),
152 capacity: self.len(),
153 }));
154 }
155
156 let next = take(self);
157 let (this, next) = next.split_at_mut(bytes.len());
158 this.copy_from_slice(bytes);
159 *self = next;
160 Ok(())
161 }
162
163 #[inline]
164 fn write_byte<C>(&mut self, cx: &C, b: u8) -> Result<(), C::Error>
165 where
166 C: ?Sized + Context,
167 {
168 if self.is_empty() {
169 return Err(cx.message(format_args!(
170 "Buffer overflow, remaining is {} while tried to write 1",
171 self.len()
172 )));
173 }
174
175 self[0] = b;
176 *self = &mut take(self)[1..];
177 Ok(())
178 }
179}
180
181pub struct BufWriter<'a, A>
183where
184 A: 'a + ?Sized + Allocator,
185{
186 buf: Vec<'a, u8, A>,
187}
188
189impl<'a, A> BufWriter<'a, A>
190where
191 A: 'a + ?Sized + Allocator,
192{
193 pub fn new(alloc: &'a A) -> Self {
195 Self {
196 buf: Vec::new_in(alloc),
197 }
198 }
199
200 pub fn into_inner(self) -> Vec<'a, u8, A> {
202 self.buf
203 }
204}
205
206impl<'a, A> Writer for BufWriter<'a, A>
207where
208 A: 'a + ?Sized + Allocator,
209{
210 type Mut<'this> = &'this mut Self
211 where
212 Self: 'this;
213
214 #[inline(always)]
215 fn borrow_mut(&mut self) -> Self::Mut<'_> {
216 self
217 }
218
219 #[inline(always)]
220 fn extend<C>(&mut self, cx: &C, buffer: Vec<'_, u8, C::Allocator>) -> Result<(), C::Error>
221 where
222 C: ?Sized + Context,
223 {
224 if !self.buf.write(buffer.as_slice()) {
225 return Err(cx.message("Buffer overflow"));
226 }
227
228 Ok(())
229 }
230
231 #[inline(always)]
232 fn write_bytes<C>(&mut self, cx: &C, bytes: &[u8]) -> Result<(), C::Error>
233 where
234 C: ?Sized + Context,
235 {
236 if !self.buf.write(bytes) {
237 return Err(cx.message("Buffer overflow"));
238 }
239
240 Ok(())
241 }
242}
243
244#[derive(Debug)]
246struct SliceOverflow {
247 n: usize,
248 capacity: usize,
249}
250
251impl fmt::Display for SliceOverflow {
252 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
253 let SliceOverflow { n, capacity } = self;
254
255 write!(
256 f,
257 "Tried to write {n} bytes to slice, with a remaining capacity of {capacity}"
258 )
259 }
260}