fix: address missing functionality in binary matching/construction

Author: bitwalker

compiler/binary/src/bitvec.rs

@@ -2,6 +2,7 @@ use alloc::alloc::{Allocator, Global};
 use alloc::vec::Vec;
 use core::fmt;
 use core::hash::{Hash, Hasher};
+use core::iter::Extend;
 use core::mem;
 
 use num_bigint::BigInt;
@@ -106,6 +107,12 @@ impl<A: Allocator> BitVec<A> {
     pub fn matcher<'a>(&'a self) -> Matcher<'a> {
         Matcher::new(self.select())
     }
+
+    /// Returns the available capacity in bytes of the underlying buffer
+    #[inline]
+    fn bytes_available(&self) -> usize {
+        self.data.capacity() - self.byte_size()
+    }
 }
 
 impl<A: Allocator> Bitstring for BitVec<A> {
@@ -431,11 +438,15 @@ impl<A: Allocator> BitVec<A> {
 
     #[inline]
     fn reserve(&mut self, cap: usize) {
-        if (self.pos + cap + 1) >= self.data.capacity() {
-            self.data.reserve(cap + mem::size_of::<usize>());
-            unsafe {
-                self.data.set_len(self.data.capacity());
-            }
+        let available = self.bytes_available();
+        if available >= cap {
+            return;
+        }
+        // Buffer some additional capacity above and beyond
+        let cap = mem::size_of::<usize>() + cap;
+        self.data.reserve(cap - available);
+        unsafe {
+            self.data.set_len(self.data.capacity());
         }
     }
 
@@ -581,6 +592,14 @@ impl<A: Allocator> BitVec<A> {
             self.bit_offset = size - offset_shift;
         }
     }
+
+    /// This is the fallback implementation for `extend` for cases where we don't
+    /// know the length of the iterator
+    fn default_extend<I: Iterator<Item = u8>>(&mut self, iter: I) {
+        for byte in iter {
+            self.push_byte(byte);
+        }
+    }
 }
 impl<A: Allocator> Eq for BitVec<A> {}
 impl<A: Allocator, T: ?Sized + Bitstring> PartialEq<T> for BitVec<A> {
@@ -692,6 +711,92 @@ impl<A: Allocator> std::io::Write for BitVec<A> {
         Ok(())
     }
 }
+impl Extend<u8> for BitVec {
+    fn extend<T: IntoIterator<Item = u8>>(&mut self, iter: T) {
+        <Self as SpecExtend<T::IntoIter>>::spec_extend(self, iter.into_iter())
+    }
+
+    #[inline]
+    fn extend_one(&mut self, byte: u8) {
+        self.push_byte(byte);
+    }
+
+    #[inline]
+    fn extend_reserve(&mut self, additional: usize) {
+        self.reserve(additional)
+    }
+}
+
+trait SpecExtend<I> {
+    fn spec_extend(&mut self, iter: I);
+}
+impl<I: Iterator<Item = u8>> SpecExtend<I> for BitVec {
+    default fn spec_extend(&mut self, iter: I) {
+        self.default_extend(iter)
+    }
+}
+impl<'a> SpecExtend<ByteIter<'a>> for BitVec {
+    fn spec_extend(&mut self, iter: ByteIter<'a>) {
+        match iter.as_slice() {
+            Some(bytes) => self.push_bytes(bytes),
+            None => {
+                self.reserve(iter.len());
+                if self.bit_offset == 0 {
+                    for byte in iter {
+                        unsafe {
+                            self.push_byte_fast(byte);
+                        }
+                    }
+                } else {
+                    for byte in iter {
+                        unsafe {
+                            self.push_byte_slow(byte);
+                        }
+                    }
+                }
+            }
+        }
+    }
+}
+impl<'a> SpecExtend<BitsIter<'a>> for BitVec {
+    fn spec_extend(&mut self, mut iter: BitsIter<'a>) {
+        match iter.as_slice() {
+            Some(bytes) => self.push_bytes(bytes),
+            None => {
+                self.reserve(iter.byte_size());
+                if self.bit_offset == 0 {
+                    loop {
+                        match iter.next() {
+                            Some(byte) => unsafe { self.push_byte_fast(byte) },
+                            None => {
+                                if let Some(b) = iter.consume() {
+                                    unsafe {
+                                        self.push_partial_byte(b.byte(), b.size);
+                                    }
+                                }
+                                break;
+                            }
+                        }
+                    }
+                } else {
+                    loop {
+                        match iter.next() {
+                            Some(byte) => unsafe { self.push_byte_slow(byte) },
+                            None => {
+                                if let Some(b) = iter.consume() {
+                                    unsafe {
+                                        self.push_partial_byte(b.byte(), b.size);
+                                    }
+                                }
+                                break;
+                            }
+                        }
+                    }
+                }
+            }
+        }
+    }
+}
 
 #[cfg(test)]
 mod test {
@@ -899,9 +1004,9 @@ mod test {
         vec.clear();
 
         vec.push_utf16('〦', Endianness::Native);
-        assert_eq!(vec.byte_size(), 4);
-        assert_eq!(vec.bit_size(), 32);
-        assert_eq!(vec.pos, 4);
+        assert_eq!(vec.byte_size(), 2);
+        assert_eq!(vec.bit_size(), 16);
+        assert_eq!(vec.pos, 2);
         assert_eq!(vec.bit_offset, 0);
 
         vec.clear();

compiler/binary/src/iter.rs

@@ -1,12 +1,13 @@
 use core::iter;
 
-use crate::{Bitstring, Selection};
+use crate::{Bitstring, MaybePartialByte, Selection};
 
 /// Represents iteration over the bytes in a selection, which may constitute
 /// either a binary or bitstring.
 ///
 /// Iteration may produce a trailing partial byte, of which all unused bits will
 /// be zeroed.
+#[derive(Debug)]
 pub struct ByteIter<'a> {
     selection: Selection<'a>,
 }
@@ -28,6 +29,13 @@ impl<'a> ByteIter<'a> {
             selection: Selection::all(data),
         }
     }
+
+    /// In some cases, the underlying bytes can be directly accessed allowing for
+    /// more optimal access patterns, see SpecExtend impl for BitVec for an example
+    #[inline]
+    pub fn as_slice(&self) -> Option<&[u8]> {
+        self.selection.as_bytes()
+    }
 }
 impl<'a> Iterator for ByteIter<'a> {
     type Item = u8;
@@ -57,6 +65,89 @@ impl<'a> iter::ExactSizeIterator for ByteIter<'a> {
 impl<'a> iter::FusedIterator for ByteIter<'a> {}
 unsafe impl<'a> iter::TrustedLen for ByteIter<'a> {}
 
+/// Like `ByteIter`, but intended for cases where special care must be
+/// taken around a trailing partial byte, if one is present. This iterator
+/// works like `ByteIter` until it encounters a trailing partial byte, in which
+/// case it will NOT emit the final byte at all, and instead it must be requested
+/// explicitly from the iterator and handled manually.
+#[derive(Debug)]
+pub struct BitsIter<'a> {
+    selection: Selection<'a>,
+}
+impl<'a> Clone for BitsIter<'a> {
+    #[inline]
+    fn clone(&self) -> Self {
+        Self {
+            selection: self.selection,
+        }
+    }
+}
+impl<'a> BitsIter<'a> {
+    pub fn new(selection: Selection<'a>) -> Self {
+        Self { selection }
+    }
+
+    /// Returns the size in bytes (including trailing partial byte) of the underlying selection
+    #[inline]
+    pub fn byte_size(&self) -> usize {
+        self.selection.byte_size()
+    }
+
+    /// Takes the selection from this iterator, consuming it
+    ///
+    /// This is how users of this iterator must consume the final trailing byte.
+    #[inline]
+    pub fn consume(self) -> Option<MaybePartialByte> {
+        match self.selection {
+            Selection::Byte(b) if b.is_partial() => Some(b),
+            _ => None,
+        }
+    }
+
+    /// In some cases, the underlying bytes can be directly accessed allowing for
+    /// more optimal access patterns, see SpecExtend impl for BitVec for an example
+    #[inline]
+    pub fn as_slice(&self) -> Option<&[u8]> {
+        self.selection.as_bytes()
+    }
+}
+impl<'a> Iterator for BitsIter<'a> {
+    type Item = u8;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        match &self.selection {
+            Selection::Empty => None,
+            Selection::Byte(b) if b.is_partial() => None,
+            _ => self.selection.pop(),
+        }
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let len = self.len();
+        (len, Some(len))
+    }
+
+    fn count(self) -> usize {
+        self.len()
+    }
+}
+impl<'a> iter::ExactSizeIterator for BitsIter<'a> {
+    fn len(&self) -> usize {
+        let size = self.selection.byte_size();
+        match self.selection.trailing_bits() {
+            0 => size,
+            _ if size == 0 => 0,
+            _ => size - 1,
+        }
+    }
+
+    fn is_empty(&self) -> bool {
+        self.len() == 0
+    }
+}
+impl<'a> iter::FusedIterator for BitsIter<'a> {}
+unsafe impl<'a> iter::TrustedLen for BitsIter<'a> {}
+
 #[cfg(test)]
 mod tests {
     use super::*;

compiler/binary/src/lib.rs

@@ -8,6 +8,8 @@
 #![feature(const_option_ext)]
 #![feature(slice_take)]
 #![feature(arbitrary_enum_discriminant)]
+#![feature(min_specialization)]
+#![feature(extend_one)]
 
 extern crate alloc;
 #[cfg(any(test, feature = "std"))]
@@ -28,7 +30,7 @@ mod traits;
 
 pub use self::bitvec::BitVec;
 pub use self::flags::{BinaryFlags, Encoding};
-pub use self::iter::ByteIter;
+pub use self::iter::{BitsIter, ByteIter};
 pub use self::matcher::Matcher;
 pub use self::select::{MaybePartialByte, Selection};
 pub use self::spec::BinaryEntrySpecifier;
@@ -53,10 +55,10 @@ pub use self::traits::{Aligned, Binary, Bitstring, FromEndianBytes, ToEndianByte
 ///
 /// So lets apply that to an example, a 16-bit integer 64542, as viewed on a little-endian machine:
 ///
-///     * 0xfc1e (little-endian hex)
-///     * 0x1efc (big-endian hex)
-///     * 0b1111110000011110 (little-endian binary)
-///     * 0b0001111011111100 (big-endian binary)
+/// * 0xfc1e (little-endian hex)
+/// * 0x1efc (big-endian hex)
+/// * 0b1111110000011110 (little-endian binary)
+/// * 0b0001111011111100 (big-endian binary)
 ///
 /// Well that's confusing, The bytes appear to be backwards! The little-endian version has the most-significant bits in the least-significant
 /// byte, and the big-endian version has the least-significant bits in the most-significant byte. What's going on here?