From 547d83cbe52af484fe2d78a50ff11ab03eb754f9 Mon Sep 17 00:00:00 2001
From: Korrat <korrat@protonmail.com>
Date: Sat, 16 Oct 2021 13:34:03 +0200
Subject: [PATCH] feat: sign extend signed integer types

---
 src/impls/primitive.rs | 140 +++++++++++++++++++++++++++++++++++++++--
 1 file changed, 134 insertions(+), 6 deletions(-)

diff --git a/src/impls/primitive.rs b/src/impls/primitive.rs
index 25e74d35..e6a4a4ab 100644
--- a/src/impls/primitive.rs
+++ b/src/impls/primitive.rs
@@ -97,6 +97,101 @@ macro_rules! ImplDekuRead {
     };
 }
 
+macro_rules! ImplDekuReadSignExtend {
+    ($typ:ty) => {
+        impl DekuRead<'_, (Endian, Size)> for $typ {
+            fn read(
+                input: &BitSlice<Msb0, u8>,
+                (endian, size): (Endian, Size),
+            ) -> Result<(&BitSlice<Msb0, u8>, Self), DekuError> {
+                let max_type_bits: usize = Size::of::<$typ>().bit_size();
+                let bit_size: usize = size.bit_size();
+
+                let input_is_le = endian.is_le();
+
+                if bit_size > max_type_bits {
+                    return Err(DekuError::Parse(format!(
+                        "too much data: container of {} bits cannot hold {} bits",
+                        max_type_bits, bit_size
+                    )));
+                }
+
+                if input.len() < bit_size {
+                    return Err(DekuError::Incomplete(crate::error::NeedSize::new(bit_size)));
+                }
+
+                let (bit_slice, rest) = input.split_at(bit_size);
+
+                let pad = 8 * ((bit_slice.len() + 7) / 8) - bit_slice.len();
+
+                let value = if pad == 0
+                    && bit_slice.len() == max_type_bits
+                    && bit_slice.as_raw_slice().len() * 8 == max_type_bits
+                {
+                    // if everything is aligned, just read the value
+
+                    let bytes: &[u8] = bit_slice.as_raw_slice();
+
+                    // Read value
+                    if input_is_le {
+                        <$typ>::from_le_bytes(bytes.try_into()?)
+                    } else {
+                        <$typ>::from_be_bytes(bytes.try_into()?)
+                    }
+                } else {
+                    // Create a new BitVec from the slice and pad un-aligned chunks
+                    // i.e. [10010110, 1110] -> [10010110, 00001110]
+                    let bits: BitVec<Msb0, u8> = {
+                        let mut bits = BitVec::with_capacity(bit_slice.len() + pad);
+
+                        // Copy bits to new BitVec
+                        bits.extend_from_bitslice(bit_slice);
+
+                        // Force align
+                        //i.e. [1110, 10010110] -> [11101001, 0110]
+                        bits.force_align();
+
+                        // Some padding to next byte
+                        let index = if input_is_le {
+                            bits.len() - (8 - pad)
+                        } else {
+                            0
+                        };
+                        for _ in 0..pad {
+                            bits.insert(index, false);
+                        }
+
+                        // Pad up-to size of type
+                        for _ in 0..(max_type_bits - bits.len()) {
+                            if input_is_le {
+                                bits.push(false);
+                            } else {
+                                bits.insert(0, false);
+                            }
+                        }
+
+                        bits
+                    };
+
+                    let bytes: &[u8] = bits.as_raw_slice();
+
+                    // Read value
+                    let value = if input_is_le {
+                        <$typ>::from_le_bytes(bytes.try_into()?)
+                    } else {
+                        <$typ>::from_be_bytes(bytes.try_into()?)
+                    };
+
+                    let shift = max_type_bits - bit_size;
+                    (value << shift) >> shift
+                };
+
+                Ok((rest, value))
+            }
+        }
+    };
+}
+
 macro_rules! ForwardDekuRead {
     ($typ:ty) => {
         // Only have `endian`, set `bit_size` to `Size::of::<Type>()`
@@ -230,18 +325,30 @@ macro_rules! ImplDekuTraits {
     };
 }
 
+macro_rules! ImplDekuTraitsSignExtend {
+    ($typ:ty) => {
+        ImplDekuReadSignExtend!($typ);
+        ForwardDekuRead!($typ);
+
+        ImplDekuWrite!($typ);
+        ForwardDekuWrite!($typ);
+    };
+}
+
 ImplDekuTraits!(u8);
 ImplDekuTraits!(u16);
 ImplDekuTraits!(u32);
 ImplDekuTraits!(u64);
 ImplDekuTraits!(u128);
 ImplDekuTraits!(usize);
-ImplDekuTraits!(i8);
-ImplDekuTraits!(i16);
-ImplDekuTraits!(i32);
-ImplDekuTraits!(i64);
-ImplDekuTraits!(i128);
-ImplDekuTraits!(isize);
+
+ImplDekuTraitsSignExtend!(i8);
+ImplDekuTraitsSignExtend!(i16);
+ImplDekuTraitsSignExtend!(i32);
+ImplDekuTraitsSignExtend!(i64);
+ImplDekuTraitsSignExtend!(i128);
+ImplDekuTraitsSignExtend!(isize);
+
 ImplDekuTraits!(f32);
 ImplDekuTraits!(f64);
 
@@ -432,4 +539,25 @@ mod tests {
 
         assert_eq!(expected_write, res_write.into_vec());
     }
+
+    macro_rules! TestSignExtending {
+        ($test_name:ident, $typ:ty) => {
+            #[test]
+            fn $test_name() {
+                let bit_slice = [0b10101_000].view_bits::<Msb0>();
+
+                let (rest, res_read) = <$typ>::read(bit_slice, (Endian::Little, Size::Bits(5))).unwrap();
+
+                assert_eq!(-11, res_read);
+                assert_eq!(bits![Msb0, u8; 0, 0, 0], rest);
+            }
+        };
+    }
+
+    TestSignExtending!(test_sign_extend_i8, i8);
+    TestSignExtending!(test_sign_extend_i16, i16);
+    TestSignExtending!(test_sign_extend_i32, i32);
+    TestSignExtending!(test_sign_extend_i64, i64);
+    TestSignExtending!(test_sign_extend_i128, i128);
+    TestSignExtending!(test_sign_extend_isize, isize);
 }