Address MR feedback

sharksforarms · Apr 5, 2024 · c1ac6ed · c1ac6ed
1 parent 8f41f13
commit c1ac6ed
Show file tree

Hide file tree

Showing 15 changed files with 63 additions and 82 deletions.
diff --git a/src/impls/bool.rs b/src/impls/bool.rs
@@ -76,19 +76,16 @@ mod tests {
 
     #[test]
     fn test_writer() {
-        let mut out_buf = vec![];
-        let mut writer = Writer::new(&mut out_buf);
+        let mut writer = Writer::new(vec![]);
         true.to_writer(&mut writer, BitSize(1)).unwrap();
         assert_eq!(vec![true], writer.rest());
 
-        let mut out_buf = vec![];
-        let mut writer = Writer::new(&mut out_buf);
+        let mut writer = Writer::new(vec![]);
         true.to_writer(&mut writer, ()).unwrap();
-        assert_eq!(vec![1], out_buf);
+        assert_eq!(vec![1], writer.inner);
 
-        let mut out_buf = vec![];
-        let mut writer = Writer::new(&mut out_buf);
+        let mut writer = Writer::new(vec![]);
         false.to_writer(&mut writer, ()).unwrap();
-        assert_eq!(vec![0], out_buf);
+        assert_eq!(vec![0], writer.inner);
     }
 }
diff --git a/src/impls/boxed.rs b/src/impls/boxed.rs
@@ -87,10 +87,9 @@ mod tests {
         let res_read = <Box<u16>>::from_reader_with_ctx(&mut reader, ()).unwrap();
         assert_eq!(expected, res_read);
 
-        let mut out_buf = vec![];
-        let mut writer = Writer::new(&mut out_buf);
+        let mut writer = Writer::new(vec![]);
         res_read.to_writer(&mut writer, ()).unwrap();
-        assert_eq!(input.to_vec(), out_buf.to_vec());
+        assert_eq!(input.to_vec(), writer.inner);
     }
 
     // Note: Copied tests from vec.rs impl
@@ -131,12 +130,11 @@ mod tests {
 
         assert_eq!(input[..expected_write.len()].to_vec(), expected_write);
 
-        let mut out_buf = vec![];
-        let mut writer = Writer::new(&mut out_buf);
+        let mut writer = Writer::new(vec![]);
         res_read
             .to_writer(&mut writer, (endian, BitSize(bit_size)))
             .unwrap();
-        assert_eq!(expected_write, out_buf.to_vec());
+        assert_eq!(expected_write, writer.inner);
 
         assert_eq!(input[..expected_write.len()].to_vec(), expected_write);
     }

diff --git a/src/impls/cow.rs b/src/impls/cow.rs
@@ -51,9 +51,8 @@ mod tests {
         let res_read = <Cow<u16>>::from_reader_with_ctx(&mut reader, ()).unwrap();
         assert_eq!(expected, res_read);
 
-        let mut out_buf = vec![];
-        let mut writer = Writer::new(&mut out_buf);
+        let mut writer = Writer::new(vec![]);
         res_read.to_writer(&mut writer, ()).unwrap();
-        assert_eq!(input.to_vec(), out_buf.to_vec());
+        assert_eq!(input.to_vec(), writer.inner);
     }
 }
diff --git a/src/impls/cstring.rs b/src/impls/cstring.rs
@@ -67,9 +67,8 @@ mod tests {
         cursor.read_to_end(&mut buf).unwrap();
         assert_eq!(expected_rest, buf);
 
-        let mut out_buf = vec![];
-        let mut writer = Writer::new(&mut out_buf);
+        let mut writer = Writer::new(vec![]);
         res_read.to_writer(&mut writer, ()).unwrap();
-        assert_eq!(vec![b't', b'e', b's', b't', b'\0'], out_buf.to_vec());
+        assert_eq!(vec![b't', b'e', b's', b't', b'\0'], writer.inner);
     }
 }
diff --git a/src/impls/hashmap.rs b/src/impls/hashmap.rs
@@ -302,9 +302,8 @@ mod tests {
         case::normal(fxhashmap!{0x11u8 => 0xAABBu16, 0x23u8 => 0xCCDDu16}, Endian::Little, vec![0x11, 0xBB, 0xAA, 0x23, 0xDD, 0xCC]),
     )]
     fn test_hashmap_write(input: FxHashMap<u8, u16>, endian: Endian, expected: Vec<u8>) {
-        let mut out_buf = vec![];
-        let mut writer = Writer::new(&mut out_buf);
+        let mut writer = Writer::new(vec![]);
         input.to_writer(&mut writer, endian).unwrap();
-        assert_eq!(expected, out_buf);
+        assert_eq!(expected, writer.inner);
     }
 }
diff --git a/src/impls/hashset.rs b/src/impls/hashset.rs
@@ -271,10 +271,9 @@ mod tests {
         case::normal(vec![0xAABB, 0xCCDD].into_iter().collect(), Endian::Little, vec![0xDD, 0xCC, 0xBB, 0xAA]),
     )]
     fn test_hashset_write(input: FxHashSet<u16>, endian: Endian, expected: Vec<u8>) {
-        let mut out_buf = vec![];
-        let mut writer = Writer::new(&mut out_buf);
+        let mut writer = Writer::new(vec![]);
         input.to_writer(&mut writer, endian).unwrap();
-        assert_eq!(expected, out_buf);
+        assert_eq!(expected, writer.inner);
     }
 
     // Note: These tests also exist in boxed.rs
@@ -315,11 +314,10 @@ mod tests {
         cursor.read_to_end(&mut buf).unwrap();
         assert_eq!(expected_rest_bytes, buf);
 
-        let mut out_buf = vec![];
-        let mut writer = Writer::new(&mut out_buf);
+        let mut writer = Writer::new(vec![]);
         res_read
             .to_writer(&mut writer, (endian, BitSize(bit_size)))
             .unwrap();
-        assert_eq!(expected_write, out_buf);
+        assert_eq!(expected_write, writer.inner);
     }
 }
diff --git a/src/impls/ipaddr.rs b/src/impls/ipaddr.rs
@@ -83,10 +83,9 @@ mod tests {
         let res_read = Ipv4Addr::from_reader_with_ctx(&mut reader, endian).unwrap();
         assert_eq!(expected, res_read);
 
-        let mut out_buf = vec![];
-        let mut writer = Writer::new(&mut out_buf);
+        let mut writer = Writer::new(vec![]);
         res_read.to_writer(&mut writer, endian).unwrap();
-        assert_eq!(input.to_vec(), out_buf.to_vec());
+        assert_eq!(input.to_vec(), writer.inner);
     }
 
     #[rstest(input, endian, expected,
@@ -99,31 +98,28 @@ mod tests {
         let res_read = Ipv6Addr::from_reader_with_ctx(&mut reader, endian).unwrap();
         assert_eq!(expected, res_read);
 
-        let mut out_buf = vec![];
-        let mut writer = Writer::new(&mut out_buf);
+        let mut writer = Writer::new(vec![]);
         res_read.to_writer(&mut writer, endian).unwrap();
-        assert_eq!(input.to_vec(), out_buf.to_vec());
+        assert_eq!(input.to_vec(), writer.inner);
     }
 
     #[test]
     fn test_ip_addr_write() {
         let ip_addr = IpAddr::V4(Ipv4Addr::new(145, 254, 160, 237));
 
-        let mut out_buf = vec![];
-        let mut writer = Writer::new(&mut out_buf);
+        let mut writer = Writer::new(vec![]);
         ip_addr.to_writer(&mut writer, Endian::Little).unwrap();
-        assert_eq!(vec![237, 160, 254, 145], out_buf.to_vec());
+        assert_eq!(vec![237, 160, 254, 145], writer.inner);
 
         let ip_addr = IpAddr::V6(Ipv6Addr::new(0, 0, 0, 0, 0, 0xffff, 0xc00a, 0x02ff));
-        let mut out_buf = vec![];
-        let mut writer = Writer::new(&mut out_buf);
+        let mut writer = Writer::new(vec![]);
         ip_addr.to_writer(&mut writer, Endian::Little).unwrap();
         assert_eq!(
             vec![
                 0xff, 0x02, 0x0a, 0xc0, 0xff, 0xff, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
                 0x00, 0x00
             ],
-            out_buf.to_vec()
+            writer.inner
         );
     }
 }
diff --git a/src/impls/nonzero.rs b/src/impls/nonzero.rs
@@ -83,9 +83,8 @@ mod tests {
         let res_read = NonZeroU8::from_reader_with_ctx(&mut reader, ()).unwrap();
         assert_eq!(expected, res_read);
 
-        let mut out_buf = vec![];
-        let mut writer = Writer::new(&mut out_buf);
+        let mut writer = Writer::new(vec![]);
         res_read.to_writer(&mut writer, ()).unwrap();
-        assert_eq!(input.to_vec(), out_buf.to_vec());
+        assert_eq!(input.to_vec(), writer.inner);
     }
 }
diff --git a/src/impls/option.rs b/src/impls/option.rs
@@ -27,12 +27,19 @@ mod tests {
     use crate::reader::Reader;
 
     #[test]
-    fn test_option() {
+    fn test_option_read() {
         use crate::ctx::*;
         let input = &[1u8, 2, 3, 4];
         let mut cursor = Cursor::new(input);
         let mut reader = Reader::new(&mut cursor);
         let v = Option::<u32>::from_reader_with_ctx(&mut reader, Endian::Little).unwrap();
         assert_eq!(v, Some(0x04030201))
     }
+
+    #[test]
+    fn test_option_write() {
+        let mut writer = Writer::new(vec![]);
+        Some(true).to_writer(&mut writer, ()).unwrap();
+        assert_eq!(vec![1], writer.inner);
+    }
 }
diff --git a/src/impls/primitive.rs b/src/impls/primitive.rs
@@ -624,10 +624,9 @@ mod tests {
                 let res_read = <$typ>::from_reader_with_ctx(&mut reader, ENDIAN).unwrap();
                 assert_eq!($expected, res_read);
 
-                let mut out_buf = vec![];
-                let mut writer = Writer::new(&mut out_buf);
+                let mut writer = Writer::new(vec![]);
                 res_read.to_writer(&mut writer, ENDIAN).unwrap();
-                assert_eq!($input, out_buf);
+                assert_eq!($input, writer.inner);
             }
         };
     }
@@ -816,16 +815,15 @@ mod tests {
         expected: Vec<u8>,
         expected_leftover: Vec<bool>,
     ) {
-        let mut out_buf = vec![];
-        let mut writer = Writer::new(&mut out_buf);
+        let mut writer = Writer::new(vec![]);
         match bit_size {
             Some(bit_size) => input
                 .to_writer(&mut writer, (endian, BitSize(bit_size)))
                 .unwrap(),
             None => input.to_writer(&mut writer, endian).unwrap(),
         };
         assert_eq!(expected_leftover, writer.rest());
-        assert_eq!(expected, out_buf);
+        assert_eq!(expected, writer.inner);
     }
 
     #[rstest(input, endian, byte_size, expected,
@@ -837,15 +835,14 @@ mod tests {
         case::byte_size_le_bigger(0x03AB, Endian::Little, Some(10), vec![0xAB, 0b11_000000]),
     )]
     fn test_byte_writer(input: u32, endian: Endian, byte_size: Option<usize>, expected: Vec<u8>) {
-        let mut out_buf = vec![];
-        let mut writer = Writer::new(&mut out_buf);
+        let mut writer = Writer::new(vec![]);
         match byte_size {
             Some(byte_size) => input
                 .to_writer(&mut writer, (endian, ByteSize(byte_size)))
                 .unwrap(),
             None => input.to_writer(&mut writer, endian).unwrap(),
         };
-        assert_hex::assert_eq_hex!(expected, out_buf);
+        assert_hex::assert_eq_hex!(expected, writer.inner);
     }
 
     #[rstest(input, endian, bit_size, expected, expected_write,
@@ -869,15 +866,14 @@ mod tests {
         };
         assert_eq!(expected, res_read);
 
-        let mut out_buf = vec![];
-        let mut writer = Writer::new(&mut out_buf);
+        let mut writer = Writer::new(vec![]);
         match bit_size {
             Some(bit_size) => res_read
                 .to_writer(&mut writer, (endian, BitSize(bit_size)))
                 .unwrap(),
             None => res_read.to_writer(&mut writer, endian).unwrap(),
         };
-        assert_hex::assert_eq_hex!(expected_write, out_buf);
+        assert_hex::assert_eq_hex!(expected_write, writer.inner);
     }
 
     macro_rules! TestSignExtending {

diff --git a/src/impls/slice.rs b/src/impls/slice.rs
@@ -106,11 +106,9 @@ mod tests {
     )]
     fn test_bit_write(input: [u16; 2], endian: Endian, expected: Vec<u8>) {
         // test writer
-        let mut out_buf = vec![];
-        let mut writer = Writer::new(&mut out_buf);
+        let mut writer = Writer::new(vec![]);
         input.to_writer(&mut writer, endian).unwrap();
-        assert_eq!(expected, out_buf.to_vec());
-    }
+        assert_eq!(expected, writer.inner);
 
     #[rstest(input,endian,expected,expected_rest,
         case::normal_le(
@@ -158,10 +156,10 @@ mod tests {
     )]
     fn test_nested_array_bit_write(input: [[u16; 2]; 2], endian: Endian, expected: Vec<u8>) {
         // test &slice
-        let mut out_buf = vec![];
-        let mut writer = Writer::new(&mut out_buf);
+        let input = input.as_ref();
+        let mut writer = Writer::new(vec![]);
         input.to_writer(&mut writer, endian).unwrap();
-        assert_eq!(expected, out_buf.to_vec());
+        assert_eq!(expected, writer.inner);
     }
 
     #[test]

diff --git a/src/impls/tuple.rs b/src/impls/tuple.rs
@@ -97,9 +97,8 @@ mod tests {
     where
         T: DekuWriter,
     {
-        let mut out_buf = vec![];
-        let mut writer = Writer::new(&mut out_buf);
+        let mut writer = Writer::new(vec![]);
         input.to_writer(&mut writer, ()).unwrap();
-        assert_eq!(expected, out_buf);
+        assert_eq!(expected, writer.inner);
     }
 }
diff --git a/src/impls/unit.rs b/src/impls/unit.rs
@@ -40,9 +40,8 @@ mod tests {
         let res_read = <()>::from_reader_with_ctx(&mut reader, ()).unwrap();
         assert_eq!((), res_read);
 
-        let mut out_buf = vec![];
-        let mut writer = Writer::new(&mut out_buf);
+        let mut writer = Writer::new(vec![]);
         res_read.to_writer(&mut writer, ()).unwrap();
-        assert_eq!(0, out_buf.len());
+        assert!(writer.inner.is_empty());
     }
 }
diff --git a/src/impls/vec.rs b/src/impls/vec.rs
@@ -256,10 +256,9 @@ mod tests {
         case::normal(vec![0xAABB, 0xCCDD], Endian::Little, vec![0xBB, 0xAA, 0xDD, 0xCC]),
     )]
     fn test_vec_write(input: Vec<u16>, endian: Endian, expected: Vec<u8>) {
-        let mut out_buf = vec![];
-        let mut writer = Writer::new(&mut out_buf);
+        let mut writer = Writer::new(vec![]);
         input.to_writer(&mut writer, endian).unwrap();
-        assert_eq!(expected, out_buf.to_vec());
+        assert_eq!(expected, writer.inner);
     }
 
     // Note: These tests also exist in boxed.rs
@@ -298,12 +297,11 @@ mod tests {
         input.read_to_end(&mut buf).unwrap();
         assert_eq!(expected_rest_bytes, buf);
 
-        let mut out_buf = vec![];
-        let mut writer = Writer::new(&mut out_buf);
+        let mut writer = Writer::new(vec![]);
         res_read
             .to_writer(&mut writer, (endian, BitSize(bit_size)))
             .unwrap();
-        assert_eq!(expected_write, out_buf.to_vec());
+        assert_eq!(expected_write, writer.inner);
 
         assert_eq!(input_clone[..expected_write.len()].to_vec(), expected_write);
     }

diff --git a/src/writer.rs b/src/writer.rs
@@ -75,12 +75,11 @@ impl<W: Write> Writer<W> {
                 count += 1;
             });
 
-        // SAFETY: This does noto have a safety comment in bitvec. But this is safe
+        // SAFETY: This does not have a safety comment in bitvec. But this is safe
         // because of `count` here will always still be within the bounds
         // of `bits`
         bits = unsafe { bits.get_unchecked(count * bits_of::<u8>()..) };
 
-        // TODO: with_capacity?
         self.leftover = bits.to_bitvec();
         if self.inner.write_all(&buf).is_err() {
             return Err(DekuError::Write);
@@ -107,8 +106,8 @@ impl<W: Write> Writer<W> {
             #[cfg(feature = "logging")]
             log::trace!("leftover exists");
 
-            // TODO: we could check here and only send the required bits to finish the byte?
-            // (instead of sending the entire thing)
+            // TODO(perf): we could check here and only send the required bits to finish the byte,
+            // instead of sending the entire thing. The rest would be through self.inner.write.
             self.write_bits(&BitVec::from_slice(buf))?;
         } else {
             if self.inner.write_all(buf).is_err() {
@@ -133,8 +132,8 @@ impl<W: Write> Writer<W> {
                 .extend_from_bitslice(&bitvec![u8, Msb0; 0; 8 - self.leftover.len()]);
             let mut buf = alloc::vec![0x00; self.leftover.len() / 8];
 
-            // write as many leftover to the buffer (as we can, can't write bits just bytes)
-            // TODO: error if bits are leftover? (not bytes aligned)
+            // write as many leftover to the buffer. Because of the previous extend,
+            // this will include all the bits in self.leftover
             self.leftover
                 .chunks_exact(bits_of::<u8>())
                 .zip(buf.iter_mut())