lz4_flex/
sink.rs

1#[allow(unused_imports)]
2use alloc::vec::Vec;
3
4use crate::fastcpy::slice_copy;
5
6/// Returns a Sink implementation appropriate for outputing up to `required_capacity`
7/// bytes at `vec[offset..offset+required_capacity]`.
8/// It can be either a `SliceSink` (pre-filling the vec with zeroes if necessary)
9/// when the `safe-decode` feature is enabled, or `VecSink` otherwise.
10/// The argument `pos` defines the initial output position in the Sink.
11#[inline]
12#[cfg(feature = "frame")]
13pub fn vec_sink_for_compression(
14    vec: &mut Vec<u8>,
15    offset: usize,
16    pos: usize,
17    required_capacity: usize,
18) -> SliceSink {
19    {
20        vec.resize(offset + required_capacity, 0);
21        SliceSink::new(&mut vec[offset..], pos)
22    }
23}
24
25/// Returns a Sink implementation appropriate for outputing up to `required_capacity`
26/// bytes at `vec[offset..offset+required_capacity]`.
27/// It can be either a `SliceSink` (pre-filling the vec with zeroes if necessary)
28/// when the `safe-decode` feature is enabled, or `VecSink` otherwise.
29/// The argument `pos` defines the initial output position in the Sink.
30#[cfg(feature = "frame")]
31#[inline]
32pub fn vec_sink_for_decompression(
33    vec: &mut Vec<u8>,
34    offset: usize,
35    pos: usize,
36    required_capacity: usize,
37) -> SliceSink {
38    {
39        vec.resize(offset + required_capacity, 0);
40        SliceSink::new(&mut vec[offset..], pos)
41    }
42}
43
44pub trait Sink {
45    /// Returns a raw ptr to the first unfilled byte of the Sink. Analogous to `[pos..].as_ptr()`.
46    #[cfg(not(all(feature = "safe-encode", feature = "safe-decode")))]
47    unsafe fn pos_mut_ptr(&mut self) -> *mut u8;
48
49    /// read byte at position
50    #[allow(dead_code)]
51    fn byte_at(&mut self, pos: usize) -> u8;
52
53    /// Pushes a byte to the end of the Sink.
54    #[cfg(feature = "safe-encode")]
55    fn push(&mut self, byte: u8);
56
57    #[cfg(not(all(feature = "safe-encode", feature = "safe-decode")))]
58    unsafe fn base_mut_ptr(&mut self) -> *mut u8;
59
60    fn pos(&self) -> usize;
61
62    fn capacity(&self) -> usize;
63
64    #[cfg(not(all(feature = "safe-encode", feature = "safe-decode")))]
65    unsafe fn set_pos(&mut self, new_pos: usize);
66
67    #[cfg(feature = "safe-decode")]
68    fn extend_with_fill(&mut self, byte: u8, len: usize);
69
70    /// Extends the Sink with `data`.
71    fn extend_from_slice(&mut self, data: &[u8]);
72
73    fn extend_from_slice_wild(&mut self, data: &[u8], copy_len: usize);
74
75    /// Copies `len` bytes starting from `start` to the end of the Sink.
76    /// # Panics
77    /// Panics if `start` >= `pos`.
78    #[cfg(feature = "safe-decode")]
79    fn extend_from_within(&mut self, start: usize, wild_len: usize, copy_len: usize);
80
81    #[cfg(feature = "safe-decode")]
82    fn extend_from_within_overlapping(&mut self, start: usize, num_bytes: usize);
83}
84
85/// SliceSink is used as target to de/compress data into a preallocated and possibly uninitialized
86/// `&[u8]`
87/// space.
88///
89/// # Handling of Capacity
90/// Extend methods will panic if there's insufficient capacity left in the Sink.
91///
92/// # Invariants
93///   - Bytes `[..pos()]` are always initialized.
94pub struct SliceSink<'a> {
95    /// The working slice, which may contain uninitialized bytes
96    output: &'a mut [u8],
97    /// Number of bytes in start of `output` guaranteed to be initialized
98    pos: usize,
99}
100
101impl<'a> SliceSink<'a> {
102    /// Creates a `Sink` backed by the given byte slice.
103    /// `pos` defines the initial output position in the Sink.
104    /// # Panics
105    /// Panics if `pos` is out of bounds.
106    #[inline]
107    pub fn new(output: &'a mut [u8], pos: usize) -> Self {
108        // SAFETY: Caller guarantees that all elements of `output[..pos]` are initialized.
109        let _ = &mut output[..pos]; // bounds check pos
110        SliceSink { output, pos }
111    }
112}
113
114impl<'a> Sink for SliceSink<'a> {
115    /// Returns a raw ptr to the first unfilled byte of the Sink. Analogous to `[pos..].as_ptr()`.
116    #[inline]
117    #[cfg(not(all(feature = "safe-encode", feature = "safe-decode")))]
118    unsafe fn pos_mut_ptr(&mut self) -> *mut u8 {
119        self.base_mut_ptr().add(self.pos()) as *mut u8
120    }
121
122    /// Pushes a byte to the end of the Sink.
123    #[inline]
124    fn byte_at(&mut self, pos: usize) -> u8 {
125        self.output[pos]
126    }
127
128    /// Pushes a byte to the end of the Sink.
129    #[inline]
130    #[cfg(feature = "safe-encode")]
131    fn push(&mut self, byte: u8) {
132        self.output[self.pos] = byte;
133        self.pos += 1;
134    }
135
136    #[cfg(not(all(feature = "safe-encode", feature = "safe-decode")))]
137    unsafe fn base_mut_ptr(&mut self) -> *mut u8 {
138        self.output.as_mut_ptr()
139    }
140
141    #[inline]
142    fn pos(&self) -> usize {
143        self.pos
144    }
145
146    #[inline]
147    fn capacity(&self) -> usize {
148        self.output.len()
149    }
150
151    #[cfg(not(all(feature = "safe-encode", feature = "safe-decode")))]
152    #[inline]
153    unsafe fn set_pos(&mut self, new_pos: usize) {
154        debug_assert!(new_pos <= self.capacity());
155        self.pos = new_pos;
156    }
157
158    #[inline]
159    #[cfg(feature = "safe-decode")]
160    fn extend_with_fill(&mut self, byte: u8, len: usize) {
161        self.output[self.pos..self.pos + len].fill(byte);
162        self.pos += len;
163    }
164
165    /// Extends the Sink with `data`.
166    #[inline]
167    fn extend_from_slice(&mut self, data: &[u8]) {
168        self.extend_from_slice_wild(data, data.len())
169    }
170
171    #[inline]
172    fn extend_from_slice_wild(&mut self, data: &[u8], copy_len: usize) {
173        assert!(copy_len <= data.len());
174        slice_copy(data, &mut self.output[self.pos..(self.pos) + data.len()]);
175        self.pos += copy_len;
176    }
177
178    /// Copies `len` bytes starting from `start` to the end of the Sink.
179    /// # Panics
180    /// Panics if `start` >= `pos`.
181    #[inline]
182    #[cfg(feature = "safe-decode")]
183    fn extend_from_within(&mut self, start: usize, wild_len: usize, copy_len: usize) {
184        self.output.copy_within(start..start + wild_len, self.pos);
185        self.pos += copy_len;
186    }
187
188    #[inline]
189    #[cfg(feature = "safe-decode")]
190    #[cfg_attr(nightly, optimize(size))] // to avoid loop unrolling
191    fn extend_from_within_overlapping(&mut self, start: usize, num_bytes: usize) {
192        let offset = self.pos - start;
193        for i in start + offset..start + offset + num_bytes {
194            self.output[i] = self.output[i - offset];
195        }
196        self.pos += num_bytes;
197    }
198}
199
200/// PtrSink is used as target to de/compress data into a preallocated and possibly uninitialized
201/// `&[u8]`
202/// space.
203///
204///
205#[cfg(not(all(feature = "safe-encode", feature = "safe-decode")))]
206pub struct PtrSink {
207    /// The working slice, which may contain uninitialized bytes
208    output: *mut u8,
209    /// Number of bytes in start of `output` guaranteed to be initialized
210    pos: usize,
211    /// Number of bytes in output available
212    cap: usize,
213}
214
215#[cfg(not(all(feature = "safe-encode", feature = "safe-decode")))]
216impl PtrSink {
217    /// Creates a `Sink` backed by the given byte slice.
218    /// `pos` defines the initial output position in the Sink.
219    /// # Panics
220    /// Panics if `pos` is out of bounds.
221    #[inline]
222    pub fn from_vec(output: &mut Vec<u8>, pos: usize) -> Self {
223        // SAFETY: Bytes behind pointer may be uninitialized.
224        Self {
225            output: output.as_mut_ptr(),
226            pos,
227            cap: output.capacity(),
228        }
229    }
230}
231
232#[cfg(not(all(feature = "safe-encode", feature = "safe-decode")))]
233impl Sink for PtrSink {
234    /// Returns a raw ptr to the first unfilled byte of the Sink. Analogous to `[pos..].as_ptr()`.
235    #[inline]
236    #[cfg(not(all(feature = "safe-encode", feature = "safe-decode")))]
237    unsafe fn pos_mut_ptr(&mut self) -> *mut u8 {
238        self.base_mut_ptr().add(self.pos()) as *mut u8
239    }
240
241    /// Pushes a byte to the end of the Sink.
242    #[inline]
243    fn byte_at(&mut self, pos: usize) -> u8 {
244        unsafe { self.output.add(pos).read() }
245    }
246
247    /// Pushes a byte to the end of the Sink.
248    #[inline]
249    #[cfg(feature = "safe-encode")]
250    fn push(&mut self, byte: u8) {
251        unsafe {
252            self.pos_mut_ptr().write(byte);
253        }
254        self.pos += 1;
255    }
256
257    #[cfg(not(all(feature = "safe-encode", feature = "safe-decode")))]
258    unsafe fn base_mut_ptr(&mut self) -> *mut u8 {
259        self.output
260    }
261
262    #[inline]
263    fn pos(&self) -> usize {
264        self.pos
265    }
266
267    #[inline]
268    fn capacity(&self) -> usize {
269        self.cap
270    }
271
272    #[cfg(not(all(feature = "safe-encode", feature = "safe-decode")))]
273    #[inline]
274    unsafe fn set_pos(&mut self, new_pos: usize) {
275        debug_assert!(new_pos <= self.capacity());
276        self.pos = new_pos;
277    }
278
279    #[inline]
280    #[cfg(feature = "safe-decode")]
281    fn extend_with_fill(&mut self, _byte: u8, _len: usize) {
282        unreachable!();
283    }
284
285    /// Extends the Sink with `data`.
286    #[inline]
287    fn extend_from_slice(&mut self, data: &[u8]) {
288        self.extend_from_slice_wild(data, data.len())
289    }
290
291    #[inline]
292    fn extend_from_slice_wild(&mut self, data: &[u8], copy_len: usize) {
293        assert!(copy_len <= data.len());
294        unsafe {
295            core::ptr::copy_nonoverlapping(data.as_ptr(), self.pos_mut_ptr(), copy_len);
296        }
297        self.pos += copy_len;
298    }
299
300    /// Copies `len` bytes starting from `start` to the end of the Sink.
301    /// # Panics
302    /// Panics if `start` >= `pos`.
303    #[inline]
304    #[cfg(feature = "safe-decode")]
305    fn extend_from_within(&mut self, _start: usize, _wild_len: usize, _copy_len: usize) {
306        unreachable!();
307    }
308
309    #[inline]
310    #[cfg(feature = "safe-decode")]
311    fn extend_from_within_overlapping(&mut self, _start: usize, _num_bytes: usize) {
312        unreachable!();
313    }
314}
315
316#[cfg(test)]
317mod tests {
318
319    #[test]
320    #[cfg(any(feature = "safe-encode", feature = "safe-decode"))]
321    fn test_sink_slice() {
322        use crate::sink::Sink;
323        use crate::sink::SliceSink;
324        use alloc::vec::Vec;
325        let mut data = Vec::new();
326        data.resize(5, 0);
327        let sink = SliceSink::new(&mut data, 1);
328        assert_eq!(sink.pos(), 1);
329        assert_eq!(sink.capacity(), 5);
330    }
331}