| // Licensed to the Apache Software Foundation (ASF) under one |
| // or more contributor license agreements. See the NOTICE file |
| // distributed with this work for additional information |
| // regarding copyright ownership. The ASF licenses this file |
| // to you under the Apache License, Version 2.0 (the |
| // "License"); you may not use this file except in compliance |
| // with the License. You may obtain a copy of the License at |
| // |
| // http://www.apache.org/licenses/LICENSE-2.0 |
| // |
| // Unless required by applicable law or agreed to in writing, |
| // software distributed under the License is distributed on an |
| // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| // KIND, either express or implied. See the License for the |
| // specific language governing permissions and limitations |
| // under the License. |
| |
| use std::{cmp::min, mem::size_of}; |
| |
| use crate::{ |
| errors::CometResult as Result, |
| parquet::{data_type::AsBytes, util::bit_packing::unpack32}, |
| }; |
| use arrow::buffer::Buffer; |
| use datafusion_comet_spark_expr::utils::{likely, unlikely}; |
| |
| #[inline] |
| pub fn from_ne_slice<T: FromBytes>(bs: &[u8]) -> T { |
| let mut b = T::Buffer::default(); |
| { |
| let b = b.as_mut(); |
| let bs = &bs[..b.len()]; |
| b.copy_from_slice(bs); |
| } |
| T::from_ne_bytes(b) |
| } |
| |
| pub trait FromBytes: Sized { |
| type Buffer: AsMut<[u8]> + Default; |
| fn from_le_bytes(bs: Self::Buffer) -> Self; |
| fn from_be_bytes(bs: Self::Buffer) -> Self; |
| fn from_ne_bytes(bs: Self::Buffer) -> Self; |
| fn from(v: u64) -> Self; |
| } |
| |
| macro_rules! from_le_bytes { |
| ($($ty: ty),*) => { |
| $( |
| impl FromBytes for $ty { |
| type Buffer = [u8; size_of::<Self>()]; |
| fn from_le_bytes(bs: Self::Buffer) -> Self { |
| <$ty>::from_le_bytes(bs) |
| } |
| fn from_be_bytes(bs: Self::Buffer) -> Self { |
| <$ty>::from_be_bytes(bs) |
| } |
| fn from_ne_bytes(bs: Self::Buffer) -> Self { |
| <$ty>::from_ne_bytes(bs) |
| } |
| fn from(v: u64) -> Self { |
| v as $ty |
| } |
| } |
| )* |
| }; |
| } |
| |
| impl FromBytes for bool { |
| type Buffer = [u8; 1]; |
| fn from_le_bytes(bs: Self::Buffer) -> Self { |
| Self::from_ne_bytes(bs) |
| } |
| fn from_be_bytes(bs: Self::Buffer) -> Self { |
| Self::from_ne_bytes(bs) |
| } |
| fn from_ne_bytes(bs: Self::Buffer) -> Self { |
| match bs[0] { |
| 0 => false, |
| 1 => true, |
| _ => panic!("Invalid byte when reading bool"), |
| } |
| } |
| fn from(v: u64) -> Self { |
| (v & 1) == 1 |
| } |
| } |
| |
| // TODO: support f32 and f64 in the future, but there is no use case right now |
| // f32/f64::from(v: u64) will be like `from_ne_slice(v.as_bytes()))` and that is |
| // expensive as it involves copying buffers |
| from_le_bytes! { u8, u16, u32, u64, i8, i16, i32, i64 } |
| |
| /// Reads `$size` of bytes from `$src`, and reinterprets them as type `$ty`, in |
| /// little-endian order. `$ty` must implement the `Default` trait. Otherwise this won't |
| /// compile. |
| /// This is copied and modified from byteorder crate. |
| macro_rules! read_num_bytes { |
| ($ty:ty, $size:expr, $src:expr) => {{ |
| debug_assert!($size <= $src.len()); |
| let mut buffer = <$ty as $crate::common::bit::FromBytes>::Buffer::default(); |
| buffer.as_mut()[..$size].copy_from_slice(&$src[..$size]); |
| <$ty>::from_ne_bytes(buffer) |
| }}; |
| } |
| |
| /// u64 specific version of read_num_bytes! |
| /// This is faster than read_num_bytes! because this method avoids buffer copies. |
| #[inline] |
| pub fn read_num_bytes_u64(size: usize, src: &[u8]) -> u64 { |
| debug_assert!(size <= src.len()); |
| if unlikely(src.len() < 8) { |
| return read_num_bytes!(u64, size, src); |
| } |
| let in_ptr = src as *const [u8] as *const u8 as *const u64; |
| let v = unsafe { in_ptr.read_unaligned() }; |
| trailing_bits(v, size * 8) |
| } |
| |
| /// u32 specific version of read_num_bytes! |
| /// This is faster than read_num_bytes! because this method avoids buffer copies. |
| #[inline] |
| pub fn read_num_bytes_u32(size: usize, src: &[u8]) -> u32 { |
| debug_assert!(size <= src.len()); |
| if unlikely(src.len() < 4) { |
| return read_num_bytes!(u32, size, src); |
| } |
| let in_ptr = src as *const [u8] as *const u8 as *const u32; |
| let v = unsafe { in_ptr.read_unaligned() }; |
| trailing_bits(v as u64, size * 8) as u32 |
| } |
| |
| #[inline] |
| pub fn read_u64(src: &[u8]) -> u64 { |
| let in_ptr = src.as_ptr() as *const u64; |
| unsafe { in_ptr.read_unaligned() } |
| } |
| |
| #[inline] |
| pub fn read_u32(src: &[u8]) -> u32 { |
| let in_ptr = src.as_ptr() as *const u32; |
| unsafe { in_ptr.read_unaligned() } |
| } |
| |
| /// Similar to the `read_num_bytes` but read nums from bytes in big-endian order |
| /// This is used to read bytes from Java's OutputStream which writes bytes in big-endian |
| macro_rules! read_num_be_bytes { |
| ($ty:ty, $size:expr, $src:expr) => {{ |
| debug_assert!($size <= $src.len()); |
| let mut buffer = <$ty as $crate::common::bit::FromBytes>::Buffer::default(); |
| buffer.as_mut()[..$size].copy_from_slice(&$src[..$size]); |
| <$ty>::from_be_bytes(buffer) |
| }}; |
| } |
| |
| #[inline] |
| pub fn memcpy(source: &[u8], target: &mut [u8]) { |
| debug_assert!(target.len() >= source.len(), "Copying from source to target is not possible. Source has {} bytes but target has {} bytes", source.len(), target.len()); |
| // Originally `target[..source.len()].copy_from_slice(source)` |
| // We use the unsafe copy method to avoid some expensive bounds checking/ |
| unsafe { std::ptr::copy_nonoverlapping(source.as_ptr(), target.as_mut_ptr(), source.len()) } |
| } |
| |
| #[inline] |
| pub fn memcpy_value<T>(source: &T, num_bytes: usize, target: &mut [u8]) |
| where |
| T: ?Sized + AsBytes, |
| { |
| debug_assert!( |
| target.len() >= num_bytes, |
| "Not enough space. Only had {} bytes but need to put {} bytes", |
| target.len(), |
| num_bytes |
| ); |
| memcpy(&source.as_bytes()[..num_bytes], target) |
| } |
| |
| /// Returns the ceil of value/divisor |
| #[inline] |
| pub fn ceil(value: usize, divisor: usize) -> usize { |
| value / divisor + ((value % divisor != 0) as usize) |
| } |
| |
| /// Returns ceil(log2(x)) |
| #[inline] |
| pub fn log2(mut x: u64) -> u32 { |
| if x == 1 { |
| return 0; |
| } |
| x -= 1; |
| 64u32 - x.leading_zeros() |
| } |
| |
| /// Returns the `num_bits` least-significant bits of `v` |
| #[inline] |
| pub fn trailing_bits(v: u64, num_bits: usize) -> u64 { |
| if unlikely(num_bits == 0) { |
| return 0; |
| } |
| if unlikely(num_bits >= 64) { |
| return v; |
| } |
| v & ((1 << num_bits) - 1) |
| } |
| |
| #[inline] |
| pub fn set_bit(bits: &mut [u8], i: usize) { |
| bits[i / 8] |= 1 << (i % 8); |
| } |
| |
| /// Set the bit value at index `i`, for buffer `bits`. |
| /// |
| /// # Safety |
| /// This doesn't check bounds, the caller must ensure that `i` is in (0, bits.len() * 8) |
| #[inline] |
| pub unsafe fn set_bit_raw(bits: *mut u8, i: usize) { |
| *bits.add(i / 8) |= 1 << (i % 8); |
| } |
| |
| #[inline] |
| pub fn unset_bit(bits: &mut [u8], i: usize) { |
| bits[i / 8] &= !(1 << (i % 8)); |
| } |
| |
| #[inline] |
| pub fn set_bits(bits: &mut [u8], offset: usize, length: usize) { |
| let mut byte_i = offset / 8; |
| let offset_r = offset % 8; |
| let end = offset + length; |
| let end_byte_i = end / 8; |
| let end_r = end % 8; |
| |
| // if the offset starts in the middle of a byte, update the byte first |
| if offset_r != 0 { |
| let num_bits = min(length, 7); |
| bits[byte_i] |= ((1u8 << num_bits) - 1) << offset_r; |
| byte_i += 1; |
| } |
| |
| // See if there is an opportunity to do a bulk byte write |
| if byte_i < end_byte_i { |
| unsafe { |
| bits.as_mut_ptr() |
| .add(byte_i) |
| .write_bytes(255, end_byte_i - byte_i); |
| } |
| byte_i = end_byte_i; |
| } |
| |
| // take care of the last byte |
| if end_r > 0 && (byte_i == end_byte_i) { |
| bits[byte_i] |= (1u8 << end_r) - 1; |
| } |
| } |
| |
| #[inline(always)] |
| pub fn mix_hash(lower: u64, upper: u64) -> u64 { |
| let hash = (17 * 37u64).wrapping_add(lower); |
| hash.wrapping_mul(37).wrapping_add(upper) |
| } |
| |
| static BIT_MASK: [u8; 8] = [1, 2, 4, 8, 16, 32, 64, 128]; |
| |
| /// Returns whether bit at position `i` in `data` is set or not |
| #[inline] |
| pub fn get_bit(data: &[u8], i: usize) -> bool { |
| (data[i >> 3] & BIT_MASK[i & 7]) != 0 |
| } |
| |
| /// Returns the boolean value at index `i`. |
| /// |
| /// # Safety |
| /// This doesn't check bounds, the caller must ensure that index < self.len() |
| #[inline] |
| pub unsafe fn get_bit_raw(ptr: *const u8, i: usize) -> bool { |
| (*ptr.add(i >> 3) & BIT_MASK[i & 7]) != 0 |
| } |
| |
| /// Utility class for writing bit/byte streams. This class can write data in either |
| /// bit packed or byte aligned fashion. |
| pub struct BitWriter { |
| buffer: Vec<u8>, |
| max_bytes: usize, |
| buffered_values: u64, |
| byte_offset: usize, |
| bit_offset: usize, |
| start: usize, |
| } |
| |
| impl BitWriter { |
| pub fn new(max_bytes: usize) -> Self { |
| Self { |
| buffer: vec![0; max_bytes], |
| max_bytes, |
| buffered_values: 0, |
| byte_offset: 0, |
| bit_offset: 0, |
| start: 0, |
| } |
| } |
| |
| /// Initializes the writer from the existing buffer `buffer` and starting |
| /// offset `start`. |
| pub fn new_from_buf(buffer: Vec<u8>, start: usize) -> Self { |
| debug_assert!(start < buffer.len()); |
| let len = buffer.len(); |
| Self { |
| buffer, |
| max_bytes: len, |
| buffered_values: 0, |
| byte_offset: start, |
| bit_offset: 0, |
| start, |
| } |
| } |
| |
| /// Extend buffer size by `increment` bytes |
| #[inline] |
| pub fn extend(&mut self, increment: usize) { |
| self.max_bytes += increment; |
| let extra = vec![0; increment]; |
| self.buffer.extend(extra); |
| } |
| |
| /// Report buffer size, in bytes |
| #[inline] |
| pub fn capacity(&mut self) -> usize { |
| self.max_bytes |
| } |
| |
| /// Consumes and returns the current buffer. |
| #[inline] |
| pub fn consume(mut self) -> Vec<u8> { |
| self.flush(); |
| self.buffer.truncate(self.byte_offset); |
| self.buffer |
| } |
| |
| /// Flushes the internal buffered bits and returns the buffer's content. |
| /// This is a borrow equivalent of `consume` method. |
| #[inline] |
| pub fn flush_buffer(&mut self) -> &[u8] { |
| self.flush(); |
| &self.buffer()[0..self.byte_offset] |
| } |
| |
| /// Clears the internal state so the buffer can be reused. |
| #[inline] |
| pub fn clear(&mut self) { |
| self.buffered_values = 0; |
| self.byte_offset = self.start; |
| self.bit_offset = 0; |
| } |
| |
| /// Flushes the internal buffered bits and the align the buffer to the next byte. |
| #[inline] |
| pub fn flush(&mut self) { |
| let num_bytes = ceil(self.bit_offset, 8); |
| debug_assert!(self.byte_offset + num_bytes <= self.max_bytes); |
| memcpy_value( |
| &self.buffered_values, |
| num_bytes, |
| &mut self.buffer[self.byte_offset..], |
| ); |
| self.buffered_values = 0; |
| self.bit_offset = 0; |
| self.byte_offset += num_bytes; |
| } |
| |
| /// Advances the current offset by skipping `num_bytes`, flushing the internal bit |
| /// buffer first. |
| /// This is useful when you want to jump over `num_bytes` bytes and come back later |
| /// to fill these bytes. |
| /// |
| /// Returns error if `num_bytes` is beyond the boundary of the internal buffer. |
| /// Otherwise, returns the old offset. |
| #[inline] |
| pub fn skip(&mut self, num_bytes: usize) -> Result<usize> { |
| self.flush(); |
| debug_assert!(self.byte_offset <= self.max_bytes); |
| if unlikely(self.byte_offset + num_bytes > self.max_bytes) { |
| return Err(general_err!( |
| "Not enough bytes left in BitWriter. Need {} but only have {}", |
| self.byte_offset + num_bytes, |
| self.max_bytes |
| )); |
| } |
| let result = self.byte_offset; |
| self.byte_offset += num_bytes; |
| Ok(result) |
| } |
| |
| /// Returns a slice containing the next `num_bytes` bytes starting from the current |
| /// offset, and advances the underlying buffer by `num_bytes`. |
| /// This is useful when you want to jump over `num_bytes` bytes and come back later |
| /// to fill these bytes. |
| #[inline] |
| pub fn get_next_byte_ptr(&mut self, num_bytes: usize) -> Result<&mut [u8]> { |
| let offset = self.skip(num_bytes)?; |
| Ok(&mut self.buffer[offset..offset + num_bytes]) |
| } |
| |
| #[inline] |
| pub fn bytes_written(&self) -> usize { |
| self.byte_offset - self.start + ceil(self.bit_offset, 8) |
| } |
| |
| #[inline] |
| pub fn buffer(&self) -> &[u8] { |
| &self.buffer[self.start..] |
| } |
| |
| #[inline] |
| pub fn byte_offset(&self) -> usize { |
| self.byte_offset |
| } |
| |
| /// Returns the internal buffer length. This is the maximum number of bytes that this |
| /// writer can write. User needs to call `consume` to consume the current buffer |
| /// before more data can be written. |
| #[inline] |
| pub fn buffer_len(&self) -> usize { |
| self.max_bytes |
| } |
| |
| /// Writes the entire byte `value` at the byte `offset` |
| pub fn write_at(&mut self, offset: usize, value: u8) { |
| self.buffer[offset] = value; |
| } |
| |
| /// Writes the `num_bits` LSB of value `v` to the internal buffer of this writer. |
| /// The `num_bits` must not be greater than 64. This is bit packed. |
| /// |
| /// Returns false if there's not enough room left. True otherwise. |
| #[inline] |
| #[allow(clippy::unnecessary_cast)] |
| pub fn put_value(&mut self, v: u64, num_bits: usize) -> bool { |
| debug_assert!(num_bits <= 64); |
| debug_assert_eq!(v.checked_shr(num_bits as u32).unwrap_or(0), 0); // covers case v >> 64 |
| |
| let num_bytes = self.byte_offset * 8 + self.bit_offset + num_bits; |
| if unlikely(num_bytes > self.max_bytes as usize * 8) { |
| return false; |
| } |
| |
| self.buffered_values |= v << self.bit_offset; |
| self.bit_offset += num_bits; |
| if self.bit_offset >= 64 { |
| memcpy_value( |
| &self.buffered_values, |
| 8, |
| &mut self.buffer[self.byte_offset..], |
| ); |
| self.byte_offset += 8; |
| self.bit_offset -= 64; |
| self.buffered_values = 0; |
| // Perform checked right shift: v >> offset, where offset < 64, otherwise we |
| // shift all bits |
| self.buffered_values = v |
| .checked_shr((num_bits - self.bit_offset) as u32) |
| .unwrap_or(0); |
| } |
| debug_assert!(self.bit_offset < 64); |
| true |
| } |
| |
| /// Writes `val` of `num_bytes` bytes to the next aligned byte. If size of `T` is |
| /// larger than `num_bytes`, extra higher ordered bytes will be ignored. |
| /// |
| /// Returns false if there's not enough room left. True otherwise. |
| #[inline] |
| pub fn put_aligned<T: AsBytes>(&mut self, val: T, num_bytes: usize) -> bool { |
| let result = self.get_next_byte_ptr(num_bytes); |
| if unlikely(result.is_err()) { |
| // TODO: should we return `Result` for this func? |
| return false; |
| } |
| let ptr = result.unwrap(); |
| memcpy_value(&val, num_bytes, ptr); |
| true |
| } |
| |
| /// Writes `val` of `num_bytes` bytes at the designated `offset`. The `offset` is the |
| /// offset starting from the beginning of the internal buffer that this writer |
| /// maintains. Note that this will overwrite any existing data between `offset` and |
| /// `offset + num_bytes`. Also that if size of `T` is larger than `num_bytes`, extra |
| /// higher ordered bytes will be ignored. |
| /// |
| /// Returns false if there's not enough room left, or the `pos` is not valid. |
| /// True otherwise. |
| #[inline] |
| pub fn put_aligned_offset<T: AsBytes>( |
| &mut self, |
| val: T, |
| num_bytes: usize, |
| offset: usize, |
| ) -> bool { |
| if unlikely(num_bytes + offset > self.max_bytes) { |
| return false; |
| } |
| memcpy_value( |
| &val, |
| num_bytes, |
| &mut self.buffer[offset..offset + num_bytes], |
| ); |
| true |
| } |
| |
| /// Writes a VLQ encoded integer `v` to this buffer. The value is byte aligned. |
| /// |
| /// Returns false if there's not enough room left. True otherwise. |
| #[inline] |
| pub fn put_vlq_int(&mut self, mut v: u64) -> bool { |
| let mut result = true; |
| while v & 0xFFFFFFFFFFFFFF80 != 0 { |
| result &= self.put_aligned::<u8>(((v & 0x7F) | 0x80) as u8, 1); |
| v >>= 7; |
| } |
| result &= self.put_aligned::<u8>((v & 0x7F) as u8, 1); |
| result |
| } |
| |
| /// Writes a zigzag-VLQ encoded (in little endian order) int `v` to this buffer. |
| /// Zigzag-VLQ is a variant of VLQ encoding where negative and positive |
| /// numbers are encoded in a zigzag fashion. |
| /// See: https://developers.google.com/protocol-buffers/docs/encoding |
| /// |
| /// Returns false if there's not enough room left. True otherwise. |
| #[inline] |
| pub fn put_zigzag_vlq_int(&mut self, v: i64) -> bool { |
| let u: u64 = ((v << 1) ^ (v >> 63)) as u64; |
| self.put_vlq_int(u) |
| } |
| } |
| |
| /// Maximum byte length for a VLQ encoded integer |
| /// MAX_VLQ_BYTE_LEN = 5 for i32, and MAX_VLQ_BYTE_LEN = 10 for i64 |
| pub const MAX_VLQ_BYTE_LEN: usize = 10; |
| |
| pub struct BitReader { |
| /// The byte buffer to read from, passed in by client |
| buffer: Buffer, // TODO: generalize this |
| |
| /// Bytes are memcpy'd from `buffer` and values are read from this variable. |
| /// This is faster than reading values byte by byte directly from `buffer` |
| buffered_values: u64, |
| |
| /// |
| /// End Start |
| /// |............|B|B|B|B|B|B|B|B|..............| |
| /// ^ ^ |
| /// bit_offset byte_offset |
| /// |
| /// Current byte offset in `buffer` |
| byte_offset: usize, |
| |
| /// Current bit offset in `buffered_values` |
| bit_offset: usize, |
| |
| /// Total number of bytes in `buffer` |
| total_bytes: usize, |
| } |
| |
| /// Utility class to read bit/byte stream. This class can read bits or bytes that are |
| /// either byte aligned or not. |
| impl BitReader { |
| pub fn new(buf: Buffer, len: usize) -> Self { |
| let buffered_values = if size_of::<u64>() > len { |
| read_num_bytes_u64(len, buf.as_slice()) |
| } else { |
| read_u64(buf.as_slice()) |
| }; |
| BitReader { |
| buffer: buf, |
| buffered_values, |
| byte_offset: 0, |
| bit_offset: 0, |
| total_bytes: len, |
| } |
| } |
| |
| pub fn new_all(buf: Buffer) -> Self { |
| let len = buf.len(); |
| Self::new(buf, len) |
| } |
| |
| pub fn reset(&mut self, buf: Buffer) { |
| self.buffer = buf; |
| self.total_bytes = self.buffer.len(); |
| self.buffered_values = if size_of::<u64>() > self.total_bytes { |
| read_num_bytes_u64(self.total_bytes, self.buffer.as_slice()) |
| } else { |
| read_u64(self.buffer.as_slice()) |
| }; |
| self.byte_offset = 0; |
| self.bit_offset = 0; |
| } |
| |
| /// Gets the current byte offset |
| #[inline] |
| pub fn get_byte_offset(&self) -> usize { |
| self.byte_offset + ceil(self.bit_offset, 8) |
| } |
| |
| /// Reads a value of type `T` and of size `num_bits`. |
| /// |
| /// Returns `None` if there's not enough data available. `Some` otherwise. |
| pub fn get_value<T: FromBytes>(&mut self, num_bits: usize) -> Option<T> { |
| debug_assert!(num_bits <= 64); |
| debug_assert!(num_bits <= size_of::<T>() * 8); |
| |
| if unlikely(self.byte_offset * 8 + self.bit_offset + num_bits > self.total_bytes * 8) { |
| return None; |
| } |
| |
| let v = self.get_u64_value(num_bits); |
| Some(T::from(v)) |
| } |
| |
| /// Reads a `u32` value encoded using `num_bits` of bits. |
| /// |
| /// # Safety |
| /// |
| /// This method asusumes the following: |
| /// |
| /// - the `num_bits` is <= 64 |
| /// - the remaining number of bits to read in this reader is >= `num_bits`. |
| /// |
| /// Undefined behavior will happen if any of the above assumptions is violated. |
| #[inline] |
| pub fn get_u32_value(&mut self, num_bits: usize) -> u32 { |
| self.get_u64_value(num_bits) as u32 |
| } |
| |
| #[inline(always)] |
| fn get_u64_value(&mut self, num_bits: usize) -> u64 { |
| if unlikely(num_bits == 0) { |
| 0 |
| } else { |
| let v = self.buffered_values >> self.bit_offset; |
| let mask = u64::MAX >> (64 - num_bits); |
| self.bit_offset += num_bits; |
| if self.bit_offset < 64 { |
| v & mask |
| } else { |
| self.byte_offset += 8; |
| self.bit_offset -= 64; |
| self.reload_buffer_values(); |
| ((self.buffered_values << (num_bits - self.bit_offset)) | v) & mask |
| } |
| } |
| } |
| |
| /// Gets at most `num` bits from this reader, and append them to the `dst` byte slice, starting |
| /// at bit offset `offset`. |
| /// |
| /// Returns the actual number of bits appended. In case either the `dst` slice doesn't have |
| /// enough space or the current reader doesn't have enough bits to consume, the returned value |
| /// will be less than the input `num_bits`. |
| /// |
| /// # Preconditions |
| /// * `offset` MUST < dst.len() * 8 |
| pub fn get_bits(&mut self, dst: &mut [u8], offset: usize, num_bits: usize) -> usize { |
| debug_assert!(offset < dst.len() * 8); |
| |
| let remaining_bits = (self.total_bytes - self.byte_offset) * 8 - self.bit_offset; |
| let num_bits_to_read = min(remaining_bits, min(num_bits, dst.len() * 8 - offset)); |
| let mut i = 0; |
| |
| // First consume all the remaining bits from the `buffered_values` if there're any. |
| if likely(self.bit_offset != 0) { |
| i += self.get_bits_buffered(dst, offset, num_bits_to_read); |
| } |
| |
| debug_assert!(self.bit_offset == 0 || i == num_bits_to_read); |
| |
| // Check if there's opportunity to directly copy bytes using `memcpy`. |
| if (offset + i) % 8 == 0 && i < num_bits_to_read { |
| let num_bytes = (num_bits_to_read - i) / 8; |
| let dst_byte_offset = (offset + i) / 8; |
| if num_bytes > 0 { |
| memcpy( |
| &self.buffer[self.byte_offset..self.byte_offset + num_bytes], |
| &mut dst[dst_byte_offset..], |
| ); |
| i += num_bytes * 8; |
| self.byte_offset += num_bytes; |
| self.reload_buffer_values(); |
| } |
| } |
| |
| debug_assert!((offset + i) % 8 != 0 || num_bits_to_read - i < 8); |
| |
| // Now copy the remaining bits if there's any. |
| while i < num_bits_to_read { |
| i += self.get_bits_buffered(dst, offset + i, num_bits_to_read - i); |
| } |
| |
| num_bits_to_read |
| } |
| |
| /// Consume at most `n` bits from `buffered_values`. Returns the actual number of bits consumed. |
| /// |
| /// # Postcondition |
| /// - either bits from `buffered_values` are completely drained (i.e., `bit_offset` == 0) |
| /// - OR the `num_bits` is < the number of remaining bits in `buffered_values` and thus the |
| /// returned value is < `num_bits`. |
| /// |
| /// Either way, the returned value is in range [0, 64]. |
| #[inline] |
| fn get_bits_buffered(&mut self, dst: &mut [u8], offset: usize, num_bits: usize) -> usize { |
| if unlikely(num_bits == 0) { |
| return 0; |
| } |
| |
| let n = min(num_bits, 64 - self.bit_offset); |
| let offset_i = offset / 8; |
| let offset_r = offset % 8; |
| |
| // Extract the value to read out of the buffer |
| let mut v = trailing_bits(self.buffered_values >> self.bit_offset, n); |
| |
| // Read the first byte always because n > 0 |
| dst[offset_i] |= (v << offset_r) as u8; |
| v >>= 8 - offset_r; |
| |
| // Read the rest of the bytes |
| ((offset_i + 1)..(offset_i + ceil(n + offset_r, 8))).for_each(|i| { |
| dst[i] |= v as u8; |
| v >>= 8; |
| }); |
| |
| self.bit_offset += n; |
| if self.bit_offset == 64 { |
| self.byte_offset += 8; |
| self.bit_offset -= 64; |
| self.reload_buffer_values(); |
| } |
| |
| n |
| } |
| |
| /// Skips at most `num` bits from this reader. |
| /// |
| /// Returns the actual number of bits skipped. |
| pub fn skip_bits(&mut self, num_bits: usize) -> usize { |
| let remaining_bits = (self.total_bytes - self.byte_offset) * 8 - self.bit_offset; |
| let num_bits_to_read = min(remaining_bits, num_bits); |
| let mut i = 0; |
| |
| // First skip all the remaining bits by updating the offsets of `buffered_values`. |
| if likely(self.bit_offset != 0) { |
| let n = 64 - self.bit_offset; |
| if num_bits_to_read < n { |
| self.bit_offset += num_bits_to_read; |
| i = num_bits_to_read; |
| } else { |
| self.byte_offset += 8; |
| self.bit_offset = 0; |
| i = n; |
| } |
| } |
| |
| // Check if there's opportunity to skip by byte |
| if i + 7 < num_bits_to_read { |
| let num_bytes = (num_bits_to_read - i) / 8; |
| i += num_bytes * 8; |
| self.byte_offset += num_bytes; |
| } |
| |
| if self.bit_offset == 0 { |
| self.reload_buffer_values(); |
| } |
| |
| // Now skip the remaining bits if there's any. |
| if i < num_bits_to_read { |
| self.bit_offset += num_bits_to_read - i; |
| } |
| |
| num_bits_to_read |
| } |
| |
| /// Reads a batch of `u32` values encoded using `num_bits` of bits, into `dst`. |
| /// |
| /// # Safety |
| /// |
| /// This method asusumes the following: |
| /// |
| /// - the `num_bits` is <= 64 |
| /// - the remaining number of bits to read in this reader is >= `total * num_bits`. |
| /// |
| /// Undefined behavior will happen if any of the above assumptions is violated. |
| /// |
| /// Unlike `[get_batch]`, this method removes a few checks such as checking the remaining number |
| /// of bits as well as checking the bit width for the element type in `dst`. Therefore, it is |
| /// more efficient. |
| pub unsafe fn get_u32_batch(&mut self, mut dst: *mut u32, total: usize, num_bits: usize) { |
| let mut i = 0; |
| |
| // First align bit offset to byte offset |
| if likely(self.bit_offset != 0) { |
| while i < total && self.bit_offset != 0 { |
| *dst = self.get_u32_value(num_bits); |
| dst = dst.offset(1); |
| i += 1; |
| } |
| } |
| |
| let in_buf = &self.buffer.as_slice()[self.byte_offset..]; |
| let mut in_ptr = in_buf as *const [u8] as *const u8 as *const u32; |
| while total - i >= 32 { |
| in_ptr = unpack32(in_ptr, dst, num_bits); |
| self.byte_offset += 4 * num_bits; |
| dst = dst.offset(32); |
| i += 32; |
| } |
| |
| self.reload_buffer_values(); |
| while i < total { |
| *dst = self.get_u32_value(num_bits); |
| dst = dst.offset(1); |
| i += 1; |
| } |
| } |
| |
| pub fn get_batch<T: FromBytes>(&mut self, batch: &mut [T], num_bits: usize) -> usize { |
| debug_assert!(num_bits <= 32); |
| debug_assert!(num_bits <= size_of::<T>() * 8); |
| |
| let mut values_to_read = batch.len(); |
| let needed_bits = num_bits * values_to_read; |
| let remaining_bits = (self.total_bytes - self.byte_offset) * 8 - self.bit_offset; |
| if remaining_bits < needed_bits { |
| values_to_read = remaining_bits / num_bits; |
| } |
| |
| let mut i = 0; |
| |
| // First align bit offset to byte offset |
| if likely(self.bit_offset != 0) { |
| while i < values_to_read && self.bit_offset != 0 { |
| batch[i] = self |
| .get_value(num_bits) |
| .expect("expected to have more data"); |
| i += 1; |
| } |
| } |
| |
| unsafe { |
| let in_buf = &self.buffer.as_slice()[self.byte_offset..]; |
| let mut in_ptr = in_buf as *const [u8] as *const u8 as *const u32; |
| // FIXME assert!(memory::is_ptr_aligned(in_ptr)); |
| if size_of::<T>() == 4 { |
| while values_to_read - i >= 32 { |
| let out_ptr = &mut batch[i..] as *mut [T] as *mut T as *mut u32; |
| in_ptr = unpack32(in_ptr, out_ptr, num_bits); |
| self.byte_offset += 4 * num_bits; |
| i += 32; |
| } |
| } else { |
| let mut out_buf = [0u32; 32]; |
| let out_ptr = &mut out_buf as &mut [u32] as *mut [u32] as *mut u32; |
| while values_to_read - i >= 32 { |
| in_ptr = unpack32(in_ptr, out_ptr, num_bits); |
| self.byte_offset += 4 * num_bits; |
| for n in 0..32 { |
| // We need to copy from smaller size to bigger size to avoid |
| // overwriting other memory regions. |
| if size_of::<T>() > size_of::<u32>() { |
| std::ptr::copy_nonoverlapping( |
| out_buf[n..].as_ptr(), |
| &mut batch[i] as *mut T as *mut u32, |
| 1, |
| ); |
| } else { |
| std::ptr::copy_nonoverlapping( |
| out_buf[n..].as_ptr() as *const T, |
| &mut batch[i] as *mut T, |
| 1, |
| ); |
| } |
| i += 1; |
| } |
| } |
| } |
| } |
| |
| debug_assert!(values_to_read - i < 32); |
| |
| self.reload_buffer_values(); |
| while i < values_to_read { |
| batch[i] = self |
| .get_value(num_bits) |
| .expect("expected to have more data"); |
| i += 1; |
| } |
| |
| values_to_read |
| } |
| |
| /// Reads a `num_bytes`-sized value from this buffer and return it. |
| /// `T` needs to be a little-endian native type. The value is assumed to be byte |
| /// aligned so the bit reader will be advanced to the start of the next byte before |
| /// reading the value. |
| |
| /// Returns `Some` if there's enough bytes left to form a value of `T`. |
| /// Otherwise `None`. |
| pub fn get_aligned<T: FromBytes>(&mut self, num_bytes: usize) -> Option<T> { |
| debug_assert!(8 >= size_of::<T>()); |
| debug_assert!(num_bytes <= size_of::<T>()); |
| |
| let bytes_read = ceil(self.bit_offset, 8); |
| if unlikely(self.byte_offset + bytes_read + num_bytes > self.total_bytes) { |
| return None; |
| } |
| |
| if bytes_read + num_bytes > 8 { |
| // There may be still unread bytes in buffered_values; however, just reloading seems to |
| // be faster than stitching the buffer with the next buffer based on micro benchmarks |
| // because reloading logic can be simpler |
| |
| // Advance byte_offset to next unread byte |
| self.byte_offset += bytes_read; |
| // Reset buffered_values |
| self.reload_buffer_values(); |
| self.bit_offset = 0 |
| } else { |
| // Advance bit_offset to next unread byte |
| self.bit_offset = bytes_read * 8; |
| } |
| |
| let v = T::from(trailing_bits( |
| self.buffered_values >> self.bit_offset, |
| num_bytes * 8, |
| )); |
| self.bit_offset += num_bytes * 8; |
| |
| if self.bit_offset == 64 { |
| self.byte_offset += 8; |
| self.bit_offset -= 64; |
| self.reload_buffer_values(); |
| } |
| |
| Some(v) |
| } |
| |
| /// Reads a VLQ encoded (in little endian order) int from the stream. |
| /// The encoded int must start at the beginning of a byte. |
| /// |
| /// Returns `None` if there's not enough bytes in the stream. `Some` otherwise. |
| pub fn get_vlq_int(&mut self) -> Option<i64> { |
| let mut shift = 0; |
| let mut v: i64 = 0; |
| while let Some(byte) = self.get_aligned::<u8>(1) { |
| v |= ((byte & 0x7F) as i64) << shift; |
| shift += 7; |
| debug_assert!( |
| shift <= MAX_VLQ_BYTE_LEN * 7, |
| "Num of bytes exceed MAX_VLQ_BYTE_LEN ({})", |
| MAX_VLQ_BYTE_LEN |
| ); |
| if likely(byte & 0x80 == 0) { |
| return Some(v); |
| } |
| } |
| None |
| } |
| |
| /// Reads a zigzag-VLQ encoded (in little endian order) int from the stream |
| /// Zigzag-VLQ is a variant of VLQ encoding where negative and positive numbers are |
| /// encoded in a zigzag fashion. |
| /// See: https://developers.google.com/protocol-buffers/docs/encoding |
| /// |
| /// Note: the encoded int must start at the beginning of a byte. |
| /// |
| /// Returns `None` if the number of bytes there's not enough bytes in the stream. |
| /// `Some` otherwise. |
| #[inline] |
| pub fn get_zigzag_vlq_int(&mut self) -> Option<i64> { |
| self.get_vlq_int().map(|v| { |
| let u = v as u64; |
| (u >> 1) as i64 ^ -((u & 1) as i64) |
| }) |
| } |
| |
| fn reload_buffer_values(&mut self) { |
| let bytes_to_read = self.total_bytes - self.byte_offset; |
| self.buffered_values = if 8 > bytes_to_read { |
| read_num_bytes_u64(bytes_to_read, &self.buffer.as_slice()[self.byte_offset..]) |
| } else { |
| read_u64(&self.buffer.as_slice()[self.byte_offset..]) |
| }; |
| } |
| } |
| |
| impl From<Vec<u8>> for BitReader { |
| #[inline] |
| fn from(vec: Vec<u8>) -> Self { |
| let len = vec.len(); |
| BitReader::new(Buffer::from(vec.as_slice()), len) |
| } |
| } |
| |
| /// Returns the nearest multiple of `factor` that is `>=` than `num`. Here `factor` must |
| /// be a power of 2. |
| /// |
| /// Copied from the arrow crate to make arrow optional |
| pub fn round_upto_power_of_2(num: usize, factor: usize) -> usize { |
| debug_assert!(factor > 0 && (factor & (factor - 1)) == 0); |
| (num + (factor - 1)) & !(factor - 1) |
| } |
| |
| #[cfg(test)] |
| mod tests { |
| use super::*; |
| use crate::parquet::util::test_common::*; |
| |
| use rand::{ |
| distributions::{Distribution, Standard}, |
| Rng, |
| }; |
| use std::fmt::Debug; |
| |
| #[test] |
| fn test_read_num_bytes_u64() { |
| let buffer: Vec<u8> = vec![0, 1, 2, 3, 4, 5, 6, 7]; |
| for size in 0..buffer.len() { |
| assert_eq!( |
| read_num_bytes_u64(size, &buffer), |
| read_num_bytes!(u64, size, &buffer), |
| ); |
| } |
| } |
| |
| #[test] |
| fn test_read_u64() { |
| let buffer: Vec<u8> = vec![0, 1, 2, 3, 4, 5, 6, 7]; |
| assert_eq!(read_u64(&buffer), read_num_bytes!(u64, 8, &buffer),); |
| } |
| |
| #[test] |
| fn test_read_num_bytes_u32() { |
| let buffer: Vec<u8> = vec![0, 1, 2, 3]; |
| for size in 0..buffer.len() { |
| assert_eq!( |
| read_num_bytes_u32(size, &buffer), |
| read_num_bytes!(u32, size, &buffer), |
| ); |
| } |
| } |
| |
| #[test] |
| fn test_read_u32() { |
| let buffer: Vec<u8> = vec![0, 1, 2, 3]; |
| assert_eq!(read_u32(&buffer), read_num_bytes!(u32, 4, &buffer),); |
| } |
| |
| #[test] |
| fn test_ceil() { |
| assert_eq!(ceil(0, 1), 0); |
| assert_eq!(ceil(1, 1), 1); |
| assert_eq!(ceil(1, 2), 1); |
| assert_eq!(ceil(1, 8), 1); |
| assert_eq!(ceil(7, 8), 1); |
| assert_eq!(ceil(8, 8), 1); |
| assert_eq!(ceil(9, 8), 2); |
| assert_eq!(ceil(9, 9), 1); |
| assert_eq!(ceil(10000000000, 10), 1000000000); |
| assert_eq!(ceil(10, 10000000000), 1); |
| assert_eq!(ceil(10000000000, 1000000000), 10); |
| } |
| |
| #[test] |
| fn test_bit_reader_get_byte_offset() { |
| let buffer = vec![255; 10]; |
| let mut bit_reader = BitReader::from(buffer); |
| assert_eq!(bit_reader.get_byte_offset(), 0); // offset (0 bytes, 0 bits) |
| bit_reader.get_value::<i32>(6); |
| assert_eq!(bit_reader.get_byte_offset(), 1); // offset (0 bytes, 6 bits) |
| bit_reader.get_value::<i32>(10); |
| assert_eq!(bit_reader.get_byte_offset(), 2); // offset (0 bytes, 16 bits) |
| bit_reader.get_value::<i32>(20); |
| assert_eq!(bit_reader.get_byte_offset(), 5); // offset (0 bytes, 36 bits) |
| bit_reader.get_value::<i32>(30); |
| assert_eq!(bit_reader.get_byte_offset(), 9); // offset (8 bytes, 2 bits) |
| } |
| |
| #[test] |
| fn test_bit_reader_get_value() { |
| let buffer = vec![255, 0]; |
| let mut bit_reader = BitReader::from(buffer); |
| assert_eq!(bit_reader.get_value::<i32>(1), Some(1)); |
| assert_eq!(bit_reader.get_value::<i32>(2), Some(3)); |
| assert_eq!(bit_reader.get_value::<i32>(3), Some(7)); |
| assert_eq!(bit_reader.get_value::<i32>(4), Some(3)); |
| } |
| |
| #[test] |
| fn test_bit_reader_get_value_boundary() { |
| let buffer = vec![10, 0, 0, 0, 20, 0, 30, 0, 0, 0, 40, 0]; |
| let mut bit_reader = BitReader::from(buffer); |
| assert_eq!(bit_reader.get_value::<i64>(32), Some(10)); |
| assert_eq!(bit_reader.get_value::<i64>(16), Some(20)); |
| assert_eq!(bit_reader.get_value::<i64>(32), Some(30)); |
| assert_eq!(bit_reader.get_value::<i64>(16), Some(40)); |
| } |
| |
| #[test] |
| fn test_bit_reader_get_aligned() { |
| // 01110101 11001011 |
| let buffer = Buffer::from(&[0x75, 0xCB]); |
| let mut bit_reader = BitReader::new_all(buffer.clone()); |
| assert_eq!(bit_reader.get_value::<i32>(3), Some(5)); |
| assert_eq!(bit_reader.get_aligned::<i32>(1), Some(203)); |
| assert_eq!(bit_reader.get_value::<i32>(1), None); |
| bit_reader.reset(buffer); |
| assert_eq!(bit_reader.get_aligned::<i32>(3), None); |
| } |
| |
| #[test] |
| fn test_bit_reader_get_vlq_int() { |
| // 10001001 00000001 11110010 10110101 00000110 |
| let buffer: Vec<u8> = vec![0x89, 0x01, 0xF2, 0xB5, 0x06]; |
| let mut bit_reader = BitReader::from(buffer); |
| assert_eq!(bit_reader.get_vlq_int(), Some(137)); |
| assert_eq!(bit_reader.get_vlq_int(), Some(105202)); |
| } |
| |
| #[test] |
| fn test_bit_reader_get_zigzag_vlq_int() { |
| let buffer: Vec<u8> = vec![0, 1, 2, 3]; |
| let mut bit_reader = BitReader::from(buffer); |
| assert_eq!(bit_reader.get_zigzag_vlq_int(), Some(0)); |
| assert_eq!(bit_reader.get_zigzag_vlq_int(), Some(-1)); |
| assert_eq!(bit_reader.get_zigzag_vlq_int(), Some(1)); |
| assert_eq!(bit_reader.get_zigzag_vlq_int(), Some(-2)); |
| } |
| |
| #[test] |
| fn test_set_bit() { |
| let mut buffer = vec![0, 0, 0]; |
| set_bit(&mut buffer[..], 1); |
| assert_eq!(buffer, vec![2, 0, 0]); |
| set_bit(&mut buffer[..], 4); |
| assert_eq!(buffer, vec![18, 0, 0]); |
| unset_bit(&mut buffer[..], 1); |
| assert_eq!(buffer, vec![16, 0, 0]); |
| set_bit(&mut buffer[..], 10); |
| assert_eq!(buffer, vec![16, 4, 0]); |
| set_bit(&mut buffer[..], 10); |
| assert_eq!(buffer, vec![16, 4, 0]); |
| set_bit(&mut buffer[..], 11); |
| assert_eq!(buffer, vec![16, 12, 0]); |
| unset_bit(&mut buffer[..], 10); |
| assert_eq!(buffer, vec![16, 8, 0]); |
| } |
| |
| #[test] |
| fn test_set_bits() { |
| for offset in 0..=16 { |
| for length in 0..=16 { |
| let mut actual = vec![0, 0, 0, 0]; |
| set_bits(&mut actual[..], offset, length); |
| let mut expected = vec![0, 0, 0, 0]; |
| for i in 0..length { |
| set_bit(&mut expected, offset + i); |
| } |
| assert_eq!(actual, expected); |
| } |
| } |
| } |
| |
| #[test] |
| fn test_get_bit() { |
| // 00001101 |
| assert!(get_bit(&[0b00001101], 0)); |
| assert!(!get_bit(&[0b00001101], 1)); |
| assert!(get_bit(&[0b00001101], 2)); |
| assert!(get_bit(&[0b00001101], 3)); |
| |
| // 01001001 01010010 |
| assert!(get_bit(&[0b01001001, 0b01010010], 0)); |
| assert!(!get_bit(&[0b01001001, 0b01010010], 1)); |
| assert!(!get_bit(&[0b01001001, 0b01010010], 2)); |
| assert!(get_bit(&[0b01001001, 0b01010010], 3)); |
| assert!(!get_bit(&[0b01001001, 0b01010010], 4)); |
| assert!(!get_bit(&[0b01001001, 0b01010010], 5)); |
| assert!(get_bit(&[0b01001001, 0b01010010], 6)); |
| assert!(!get_bit(&[0b01001001, 0b01010010], 7)); |
| assert!(!get_bit(&[0b01001001, 0b01010010], 8)); |
| assert!(get_bit(&[0b01001001, 0b01010010], 9)); |
| assert!(!get_bit(&[0b01001001, 0b01010010], 10)); |
| assert!(!get_bit(&[0b01001001, 0b01010010], 11)); |
| assert!(get_bit(&[0b01001001, 0b01010010], 12)); |
| assert!(!get_bit(&[0b01001001, 0b01010010], 13)); |
| assert!(get_bit(&[0b01001001, 0b01010010], 14)); |
| assert!(!get_bit(&[0b01001001, 0b01010010], 15)); |
| } |
| |
| #[test] |
| fn test_log2() { |
| assert_eq!(log2(1), 0); |
| assert_eq!(log2(2), 1); |
| assert_eq!(log2(3), 2); |
| assert_eq!(log2(4), 2); |
| assert_eq!(log2(5), 3); |
| assert_eq!(log2(5), 3); |
| assert_eq!(log2(6), 3); |
| assert_eq!(log2(7), 3); |
| assert_eq!(log2(8), 3); |
| assert_eq!(log2(9), 4); |
| } |
| |
| #[test] |
| fn test_skip() { |
| let mut writer = BitWriter::new(5); |
| let old_offset = writer.skip(1).expect("skip() should return OK"); |
| writer.put_aligned(42, 4); |
| writer.put_aligned_offset(0x10, 1, old_offset); |
| let result = writer.consume(); |
| assert_eq!(result.as_ref(), [0x10, 42, 0, 0, 0]); |
| |
| writer = BitWriter::new(4); |
| let result = writer.skip(5); |
| assert!(result.is_err()); |
| } |
| |
| #[test] |
| fn test_get_next_byte_ptr() { |
| let mut writer = BitWriter::new(5); |
| { |
| let first_byte = writer |
| .get_next_byte_ptr(1) |
| .expect("get_next_byte_ptr() should return OK"); |
| first_byte[0] = 0x10; |
| } |
| writer.put_aligned(42, 4); |
| let result = writer.consume(); |
| assert_eq!(result.as_ref(), [0x10, 42, 0, 0, 0]); |
| } |
| |
| #[test] |
| fn test_consume_flush_buffer() { |
| let mut writer1 = BitWriter::new(3); |
| let mut writer2 = BitWriter::new(3); |
| for i in 1..10 { |
| writer1.put_value(i, 4); |
| writer2.put_value(i, 4); |
| } |
| let res1 = writer1.flush_buffer(); |
| let res2 = writer2.consume(); |
| assert_eq!(res1, &res2[..]); |
| } |
| |
| #[test] |
| fn test_put_get_bool() { |
| let len = 8; |
| let mut writer = BitWriter::new(len); |
| |
| for i in 0..8 { |
| let result = writer.put_value(i % 2, 1); |
| assert!(result); |
| } |
| |
| writer.flush(); |
| { |
| let buffer = writer.buffer(); |
| assert_eq!(buffer[0], 0b10101010); |
| } |
| |
| // Write 00110011 |
| for i in 0..8 { |
| let result = match i { |
| 0 | 1 | 4 | 5 => writer.put_value(false as u64, 1), |
| _ => writer.put_value(true as u64, 1), |
| }; |
| assert!(result); |
| } |
| writer.flush(); |
| { |
| let buffer = writer.buffer(); |
| assert_eq!(buffer[0], 0b10101010); |
| assert_eq!(buffer[1], 0b11001100); |
| } |
| |
| let mut reader = BitReader::from(writer.consume()); |
| |
| for i in 0..8 { |
| let val = reader |
| .get_value::<u8>(1) |
| .expect("get_value() should return OK"); |
| assert_eq!(val, i % 2); |
| } |
| |
| for i in 0..8 { |
| let val = reader |
| .get_value::<bool>(1) |
| .expect("get_value() should return OK"); |
| match i { |
| 0 | 1 | 4 | 5 => assert!(!val), |
| _ => assert!(val), |
| } |
| } |
| } |
| |
| #[test] |
| fn test_put_value_roundtrip() { |
| test_put_value_rand_numbers(32, 2); |
| test_put_value_rand_numbers(32, 3); |
| test_put_value_rand_numbers(32, 4); |
| test_put_value_rand_numbers(32, 5); |
| test_put_value_rand_numbers(32, 6); |
| test_put_value_rand_numbers(32, 7); |
| test_put_value_rand_numbers(32, 8); |
| test_put_value_rand_numbers(64, 16); |
| test_put_value_rand_numbers(64, 24); |
| test_put_value_rand_numbers(64, 32); |
| } |
| |
| fn test_put_value_rand_numbers(total: usize, num_bits: usize) { |
| assert!(num_bits < 64); |
| let num_bytes = ceil(num_bits, 8); |
| let mut writer = BitWriter::new(num_bytes * total); |
| let values: Vec<u64> = random_numbers::<u64>(total) |
| .iter() |
| .map(|v| v & ((1 << num_bits) - 1)) |
| .collect(); |
| (0..total).for_each(|i| { |
| assert!( |
| writer.put_value(values[i], num_bits), |
| "[{}]: put_value() failed", |
| i |
| ); |
| }); |
| |
| let mut reader = BitReader::from(writer.consume()); |
| (0..total).for_each(|i| { |
| let v = reader |
| .get_value::<u64>(num_bits) |
| .expect("get_value() should return OK"); |
| assert_eq!( |
| v, values[i], |
| "[{}]: expected {} but got {}", |
| i, values[i], v |
| ); |
| }); |
| } |
| |
| #[test] |
| fn test_get_bits() { |
| const NUM_BYTES: usize = 100; |
| |
| let mut vec = vec![0; NUM_BYTES]; |
| let total_num_bits = NUM_BYTES * 8; |
| let v = random_bools(total_num_bits); |
| (0..total_num_bits).for_each(|i| { |
| if v[i] { |
| set_bit(&mut vec, i); |
| } else { |
| unset_bit(&mut vec, i); |
| } |
| }); |
| |
| let expected = vec.clone(); |
| |
| // test reading the first time from a buffer |
| for &(offset, num_bits) in [(0, 10), (2, 10), (8, 16), (25, 40), (7, 64)].iter() { |
| let mut reader = BitReader::from(vec.clone()); |
| let mut buffer = vec![0; NUM_BYTES]; |
| |
| let actual_bits_read = reader.get_bits(&mut buffer, offset, num_bits); |
| let expected_bits_read = ::std::cmp::min(buffer.len() * 8 - offset, num_bits); |
| assert_eq!(expected_bits_read, actual_bits_read); |
| |
| for i in 0..actual_bits_read { |
| assert_eq!(get_bit(&expected, i), get_bit(&buffer, offset + i)); |
| } |
| } |
| |
| // test reading consecutively from a buffer |
| let mut reader = BitReader::from(vec); |
| let mut buffer = vec![0; NUM_BYTES]; |
| let mut rng = rand::thread_rng(); |
| let mut bits_read = 0; |
| |
| loop { |
| if bits_read >= total_num_bits { |
| break; |
| } |
| let n: usize = rng.gen(); |
| let num_bits = n % 20; |
| bits_read += reader.get_bits(&mut buffer, bits_read, num_bits); |
| } |
| |
| assert_eq!(total_num_bits, bits_read); |
| assert_eq!(&expected, &buffer); |
| } |
| |
| #[test] |
| fn test_skip_bits() { |
| const NUM_BYTES: usize = 100; |
| |
| let mut vec = vec![0; NUM_BYTES]; |
| let total_num_bits = NUM_BYTES * 8; |
| let v = random_bools(total_num_bits); |
| (0..total_num_bits).for_each(|i| { |
| if v[i] { |
| set_bit(&mut vec, i); |
| } else { |
| unset_bit(&mut vec, i); |
| } |
| }); |
| |
| let expected = vec.clone(); |
| |
| // test skipping and check the next value |
| let mut reader = BitReader::from(vec); |
| let mut bits_read = 0; |
| for &num_bits in [10, 60, 8].iter() { |
| let actual_bits_read = reader.skip_bits(num_bits); |
| assert_eq!(num_bits, actual_bits_read); |
| |
| bits_read += num_bits; |
| assert_eq!(Some(get_bit(&expected, bits_read)), reader.get_value(1)); |
| bits_read += 1; |
| } |
| |
| // test skipping consecutively |
| let mut rng = rand::thread_rng(); |
| loop { |
| if bits_read >= total_num_bits { |
| break; |
| } |
| let n: usize = rng.gen(); |
| let num_bits = n % 20; |
| bits_read += reader.skip_bits(num_bits); |
| } |
| |
| assert_eq!(total_num_bits, bits_read); |
| } |
| |
| #[test] |
| fn test_get_batch() { |
| const SIZE: &[usize] = &[1, 31, 32, 33, 128, 129]; |
| for s in SIZE { |
| for i in 0..33 { |
| match i { |
| 0..=8 => test_get_batch_helper::<u8>(*s, i), |
| 9..=16 => test_get_batch_helper::<u16>(*s, i), |
| _ => test_get_batch_helper::<u32>(*s, i), |
| } |
| } |
| } |
| } |
| |
| fn test_get_batch_helper<T>(total: usize, num_bits: usize) |
| where |
| T: FromBytes + Default + Clone + Debug + Eq, |
| { |
| assert!(num_bits <= 32); |
| let num_bytes = ceil(num_bits, 8); |
| let mut writer = BitWriter::new(num_bytes * total); |
| |
| let values: Vec<u32> = random_numbers::<u32>(total) |
| .iter() |
| .map(|v| v & ((1u64 << num_bits) - 1) as u32) |
| .collect(); |
| |
| // Generic values used to check against actual values read from `get_batch`. |
| let expected_values: Vec<T> = values.iter().map(|v| from_ne_slice(v.as_bytes())).collect(); |
| |
| (0..total).for_each(|i| { |
| assert!(writer.put_value(values[i] as u64, num_bits)); |
| }); |
| |
| let buf = writer.consume(); |
| let mut reader = BitReader::from(buf); |
| let mut batch = vec![T::default(); values.len()]; |
| let values_read = reader.get_batch::<T>(&mut batch, num_bits); |
| assert_eq!(values_read, values.len()); |
| for i in 0..batch.len() { |
| assert_eq!( |
| batch[i], expected_values[i], |
| "num_bits = {}, index = {}", |
| num_bits, i |
| ); |
| } |
| } |
| |
| #[test] |
| fn test_get_u32_batch() { |
| const SIZE: &[usize] = &[1, 31, 32, 33, 128, 129]; |
| for total in SIZE { |
| for num_bits in 1..33 { |
| let num_bytes = ceil(num_bits, 8); |
| let mut writer = BitWriter::new(num_bytes * total); |
| |
| let values: Vec<u32> = random_numbers::<u32>(*total) |
| .iter() |
| .map(|v| v & ((1u64 << num_bits) - 1) as u32) |
| .collect(); |
| |
| (0..*total).for_each(|i| { |
| assert!(writer.put_value(values[i] as u64, num_bits)); |
| }); |
| |
| let buf = writer.consume(); |
| let mut reader = BitReader::from(buf); |
| let mut batch = vec![0u32; values.len()]; |
| unsafe { |
| reader.get_u32_batch(batch.as_mut_ptr(), *total, num_bits); |
| } |
| for i in 0..batch.len() { |
| assert_eq!( |
| batch[i], values[i], |
| "num_bits = {}, index = {}", |
| num_bits, i |
| ); |
| } |
| } |
| } |
| } |
| |
| #[test] |
| fn test_put_aligned_roundtrip() { |
| test_put_aligned_rand_numbers::<u8>(4, 3); |
| test_put_aligned_rand_numbers::<u8>(16, 5); |
| test_put_aligned_rand_numbers::<i16>(32, 7); |
| test_put_aligned_rand_numbers::<i16>(32, 9); |
| test_put_aligned_rand_numbers::<i32>(32, 11); |
| test_put_aligned_rand_numbers::<i32>(32, 13); |
| test_put_aligned_rand_numbers::<i64>(32, 17); |
| test_put_aligned_rand_numbers::<i64>(32, 23); |
| } |
| |
| fn test_put_aligned_rand_numbers<T>(total: usize, num_bits: usize) |
| where |
| T: Copy + FromBytes + AsBytes + Debug + PartialEq, |
| Standard: Distribution<T>, |
| { |
| assert!(num_bits <= 32); |
| assert!(total % 2 == 0); |
| |
| let aligned_value_byte_width = std::mem::size_of::<T>(); |
| let value_byte_width = ceil(num_bits, 8); |
| let mut writer = |
| BitWriter::new((total / 2) * (aligned_value_byte_width + value_byte_width)); |
| let values: Vec<u32> = random_numbers::<u32>(total / 2) |
| .iter() |
| .map(|v| v & ((1 << num_bits) - 1)) |
| .collect(); |
| let aligned_values = random_numbers::<T>(total / 2); |
| |
| for i in 0..total { |
| let j = i / 2; |
| if i % 2 == 0 { |
| assert!( |
| writer.put_value(values[j] as u64, num_bits), |
| "[{}]: put_value() failed", |
| i |
| ); |
| } else { |
| assert!( |
| writer.put_aligned::<T>(aligned_values[j], aligned_value_byte_width), |
| "[{}]: put_aligned() failed", |
| i |
| ); |
| } |
| } |
| |
| let mut reader = BitReader::from(writer.consume()); |
| for i in 0..total { |
| let j = i / 2; |
| if i % 2 == 0 { |
| let v = reader |
| .get_value::<u64>(num_bits) |
| .expect("get_value() should return OK"); |
| assert_eq!( |
| v, values[j] as u64, |
| "[{}]: expected {} but got {}", |
| i, values[j], v |
| ); |
| } else { |
| let v = reader |
| .get_aligned::<T>(aligned_value_byte_width) |
| .expect("get_aligned() should return OK"); |
| assert_eq!( |
| v, aligned_values[j], |
| "[{}]: expected {:?} but got {:?}", |
| i, aligned_values[j], v |
| ); |
| } |
| } |
| } |
| |
| #[test] |
| fn test_put_vlq_int() { |
| let total = 64; |
| let mut writer = BitWriter::new(total * 32); |
| let values = random_numbers::<u32>(total); |
| (0..total).for_each(|i| { |
| assert!( |
| writer.put_vlq_int(values[i] as u64), |
| "[{}]; put_vlq_int() failed", |
| i |
| ); |
| }); |
| |
| let mut reader = BitReader::from(writer.consume()); |
| (0..total).for_each(|i| { |
| let v = reader |
| .get_vlq_int() |
| .expect("get_vlq_int() should return OK"); |
| assert_eq!( |
| v as u32, values[i], |
| "[{}]: expected {} but got {}", |
| i, values[i], v |
| ); |
| }); |
| } |
| |
| #[test] |
| fn test_put_zigzag_vlq_int() { |
| let total = 64; |
| let mut writer = BitWriter::new(total * 32); |
| let values = random_numbers::<i32>(total); |
| (0..total).for_each(|i| { |
| assert!( |
| writer.put_zigzag_vlq_int(values[i] as i64), |
| "[{}]; put_zigzag_vlq_int() failed", |
| i |
| ); |
| }); |
| |
| let mut reader = BitReader::from(writer.consume()); |
| (0..total).for_each(|i| { |
| let v = reader |
| .get_zigzag_vlq_int() |
| .expect("get_zigzag_vlq_int() should return OK"); |
| assert_eq!( |
| v as i32, values[i], |
| "[{}]: expected {} but got {}", |
| i, values[i], v |
| ); |
| }); |
| } |
| } |
