bp3d_proto/codec/
bits.rs

1// Copyright (c) 2024, BlockProject 3D
2//
3// All rights reserved.
4//
5// Redistribution and use in source and binary forms, with or without modification,
6// are permitted provided that the following conditions are met:
7//
8//     * Redistributions of source code must retain the above copyright notice,
9//       this list of conditions and the following disclaimer.
10//     * Redistributions in binary form must reproduce the above copyright notice,
11//       this list of conditions and the following disclaimer in the documentation
12//       and/or other materials provided with the distribution.
13//     * Neither the name of BlockProject 3D nor the names of its contributors
14//       may be used to endorse or promote products derived from this software
15//       without specific prior written permission.
16//
17// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
18// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
19// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
20// A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR
21// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
22// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
23// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
24// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
25// LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
26// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
27// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
28
29use crate::util::ToUsize;
30use bytesutil::{ReadBytes, WriteBytes};
31use std::ops::{BitAnd, BitOr, Shl, Shr};
32
33pub trait BitCodec {
34    /// Reads a value of type T from the buffer argument with a custom bit offset and size assuming
35    /// the buffer size is greater or equal to the size of T.
36    ///
37    /// # Arguments
38    ///
39    /// * `buffer`: the buffer to read from.
40    ///
41    /// returns: T
42    ///
43    /// # Safety
44    ///
45    /// This function assumes that the length of the buffer passed in as argument is at least as
46    /// large as the size of T. Currently, this relies on bytesutil which does not apply any
47    /// optimization and as such passing a too small buffer will only panic, however a future
48    /// optimization might remove the panic check from release builds, essentially causing UB in
49    /// such build.
50    unsafe fn read_aligned<
51        T: ToUsize + ReadBytes + Shr<Output = T> + BitAnd<Output = T>,
52        const BIT_OFFSET: usize,
53        const BIT_SIZE: usize,
54    >(
55        buffer: &[u8],
56    ) -> T;
57
58    /// Reads a value of type T from the buffer argument with a custom bit offset and size assuming
59    /// the buffer size is always less than the size of T. This is not unsafe as will always cause
60    /// a copy into an 8 bytes buffer (the maximum size for T is 8).
61    ///
62    /// # Arguments
63    ///
64    /// * `buffer`: the buffer to read from.
65    ///
66    /// returns: T
67    fn read_unaligned<
68        T: ToUsize + ReadBytes + Shr<Output = T> + BitAnd<Output = T>,
69        const BIT_OFFSET: usize,
70        const BIT_SIZE: usize,
71    >(
72        buffer: &[u8],
73    ) -> T;
74
75    fn read<
76        T: ToUsize + ReadBytes + Shr<Output = T> + BitAnd<Output = T>,
77        const BIT_OFFSET: usize,
78        const BIT_SIZE: usize,
79    >(
80        buffer: &[u8],
81    ) -> T {
82        if size_of::<T>() != buffer.len() {
83            Self::read_unaligned::<T, BIT_OFFSET, BIT_SIZE>(buffer)
84        } else {
85            unsafe { Self::read_aligned::<T, BIT_OFFSET, BIT_SIZE>(buffer) }
86        }
87    }
88
89    /// Writes a value of type T in the buffer argument with a custom bit offset and size assuming
90    /// the buffer size is greater or equal to the size of T.
91    ///
92    /// # Arguments
93    ///
94    /// * `buffer`: the buffer to write to.
95    /// * `value`: the value to write.
96    ///
97    /// # Safety
98    ///
99    /// This function assumes that the length of the buffer passed in as argument is at least as
100    /// large as the size of T. Currently, this relies on bytesutil which does not apply any
101    /// optimization and as such passing a too small buffer will only panic, however a future
102    /// optimization might remove the panic check from release builds, essentially causing UB in
103    /// such build.
104    unsafe fn write_aligned<
105        T: ToUsize + ReadBytes + WriteBytes + Shl<Output = T> + Shr<Output = T> + BitAnd<Output = T> + BitOr<Output = T>,
106        const BIT_OFFSET: usize,
107        const BIT_SIZE: usize,
108    >(
109        buffer: &mut [u8],
110        value: T,
111    );
112
113    fn write_unaligned<
114        T: ToUsize + ReadBytes + WriteBytes + Shl<Output = T> + Shr<Output = T> + BitAnd<Output = T> + BitOr<Output = T>,
115        const BIT_OFFSET: usize,
116        const BIT_SIZE: usize,
117    >(
118        buffer: &mut [u8],
119        value: T,
120    );
121
122    fn write<
123        T: ToUsize + ReadBytes + WriteBytes + Shl<Output = T> + Shr<Output = T> + BitAnd<Output = T> + BitOr<Output = T>,
124        const BIT_OFFSET: usize,
125        const BIT_SIZE: usize,
126    >(
127        buffer: &mut [u8],
128        value: T,
129    ) {
130        if size_of::<T>() != buffer.len() {
131            Self::write_unaligned::<T, BIT_OFFSET, BIT_SIZE>(buffer, value);
132        } else {
133            unsafe { Self::write_aligned::<T, BIT_OFFSET, BIT_SIZE>(buffer, value) };
134        }
135    }
136}
137
138pub struct BitCodecLE;
139pub struct BitCodecBE;
140
141impl BitCodec for BitCodecLE {
142    unsafe fn read_aligned<
143        T: ToUsize + ReadBytes + Shr<Output = T> + BitAnd<Output = T>,
144        const BIT_OFFSET: usize,
145        const BIT_SIZE: usize,
146    >(
147        buffer: &[u8],
148    ) -> T {
149        let mask: usize = (1 << BIT_SIZE) - 1;
150        let value = T::read_bytes_le(buffer);
151        (value >> T::from_usize(BIT_OFFSET)) & T::from_usize(mask)
152    }
153
154    fn read_unaligned<
155        T: ToUsize + ReadBytes + Shr<Output = T> + BitAnd<Output = T>,
156        const BIT_OFFSET: usize,
157        const BIT_SIZE: usize,
158    >(
159        buffer: &[u8],
160    ) -> T {
161        let mut data = [0; 8];
162        data[..buffer.len()].copy_from_slice(buffer);
163        unsafe { Self::read_aligned::<T, BIT_OFFSET, BIT_SIZE>(&data) }
164    }
165
166    unsafe fn write_aligned<
167        T: ToUsize + ReadBytes + WriteBytes + Shl<Output = T> + Shr<Output = T> + BitAnd<Output = T> + BitOr<Output = T>,
168        const BIT_OFFSET: usize,
169        const BIT_SIZE: usize,
170    >(
171        buffer: &mut [u8],
172        value: T,
173    ) {
174        let mask: usize = (1 << BIT_SIZE) - 1;
175        let reset_mask = !(mask << BIT_OFFSET);
176        let original = T::read_bytes_le(buffer);
177        let clean = original & T::from_usize(reset_mask);
178        let value = (value & T::from_usize(mask)) << T::from_usize(BIT_OFFSET);
179        (clean | value).write_bytes_le(buffer);
180    }
181
182    fn write_unaligned<
183        T: ToUsize + ReadBytes + WriteBytes + Shl<Output = T> + Shr<Output = T> + BitAnd<Output = T> + BitOr<Output = T>,
184        const BIT_OFFSET: usize,
185        const BIT_SIZE: usize,
186    >(
187        buffer: &mut [u8],
188        value: T,
189    ) {
190        let mut data = [0; 8];
191        data[..buffer.len()].copy_from_slice(buffer);
192        unsafe { Self::write_aligned::<T, BIT_OFFSET, BIT_SIZE>(&mut data, value) };
193        let motherfuckingrust = buffer.len();
194        buffer.copy_from_slice(&data[..motherfuckingrust]);
195    }
196}
197
198impl BitCodec for BitCodecBE {
199    unsafe fn read_aligned<
200        T: ToUsize + ReadBytes + Shr<Output = T> + BitAnd<Output = T>,
201        const BIT_OFFSET: usize,
202        const BIT_SIZE: usize,
203    >(
204        buffer: &[u8],
205    ) -> T {
206        let mask: usize = (1 << BIT_SIZE) - 1;
207        let value = T::read_bytes_be(buffer);
208        (value >> T::from_usize(8 - (BIT_SIZE % 8) - BIT_OFFSET)) & T::from_usize(mask)
209    }
210
211    fn read_unaligned<
212        T: ToUsize + ReadBytes + Shr<Output = T> + BitAnd<Output = T>,
213        const BIT_OFFSET: usize,
214        const BIT_SIZE: usize,
215    >(
216        buffer: &[u8],
217    ) -> T {
218        let offset = size_of::<T>() - buffer.len();
219        let mut data = [0; 8];
220        data[offset..buffer.len() + offset].copy_from_slice(buffer);
221        unsafe { Self::read_aligned::<T, BIT_OFFSET, BIT_SIZE>(&data) }
222    }
223
224    unsafe fn write_aligned<
225        T: ToUsize + ReadBytes + WriteBytes + Shl<Output = T> + Shr<Output = T> + BitAnd<Output = T> + BitOr<Output = T>,
226        const BIT_OFFSET: usize,
227        const BIT_SIZE: usize,
228    >(
229        buffer: &mut [u8],
230        value: T,
231    ) {
232        let mask: usize = (1 << BIT_SIZE) - 1;
233        let reset_mask = !(mask << (8 - (BIT_SIZE % 8) - BIT_OFFSET));
234        let original = T::read_bytes_be(buffer);
235        let clean = original & T::from_usize(reset_mask);
236        let value = (value & T::from_usize(mask)) << T::from_usize(8 - (BIT_SIZE % 8) - BIT_OFFSET);
237        (clean | value).write_bytes_be(buffer);
238    }
239
240    fn write_unaligned<
241        T: ToUsize + ReadBytes + WriteBytes + Shl<Output = T> + Shr<Output = T> + BitAnd<Output = T> + BitOr<Output = T>,
242        const BIT_OFFSET: usize,
243        const BIT_SIZE: usize,
244    >(
245        buffer: &mut [u8],
246        value: T,
247    ) {
248        let offset = size_of::<T>() - buffer.len();
249        let mut data = [0; 8];
250        data[offset..buffer.len() + offset].copy_from_slice(buffer);
251        unsafe { Self::write_aligned::<T, BIT_OFFSET, BIT_SIZE>(&mut data, value) };
252        let motherfuckingrust = buffer.len();
253        buffer.copy_from_slice(&data[offset..motherfuckingrust + offset]);
254    }
255}
256
257#[cfg(test)]
258mod tests {
259    use crate::codec::{BitCodec, BitCodecBE, BitCodecLE};
260
261    #[test]
262    fn little_endian() {
263        let buffer = [0xFF, 0xFF, 0xFF, 0xFF];
264        assert_eq!(BitCodecLE::read::<u32, 0, 32>(&buffer[0..4]), 0xFFFFFFFF);
265        assert_eq!(BitCodecLE::read::<u8, 0, 1>(&buffer[0..1]), 1);
266        assert_eq!(BitCodecLE::read::<u8, 0, 4>(&buffer[0..1]), 0xF);
267        assert_eq!(BitCodecLE::read::<u8, 4, 4>(&buffer[0..1]), 0xF);
268    }
269
270    #[test]
271    fn big_endian() {
272        let buffer = [0xAB, 0xF0];
273        assert_eq!(BitCodecBE::read::<u16, 0, 12>(&buffer[0..2]), 0xABF);
274        let mut buffer = [0x0, 0x0];
275        BitCodecBE::write::<u8, 0, 4>(&mut buffer[0..1], 0xF);
276        assert_eq!(BitCodecBE::read::<u8, 0, 4>(&buffer[0..1]), 0xF);
277        BitCodecBE::write::<u16, 0, 12>(&mut buffer[0..2], 0xABF);
278        assert_eq!(BitCodecBE::read::<u16, 0, 12>(&buffer[0..2]), 0xABF);
279        BitCodecBE::write::<u8, 1, 7>(&mut buffer[1..2], 127);
280        assert_eq!(BitCodecBE::read::<u8, 1, 7>(&buffer[1..2]), 127);
281    }
282}