1
0
Fork 0
mirror of https://codeberg.org/icewind/bitbuffer.git synced 2026-06-03 16:44:06 +02:00

remove full reserve

this makes code a lot simpler
This commit is contained in:
Robin Appelman 2021-07-24 16:25:25 +02:00
commit 2a1718b5f6
4 changed files with 40 additions and 449 deletions

View file

@ -1,6 +1,6 @@
[package] [package]
name = "bitbuffer" name = "bitbuffer"
version = "0.9.9" version = "0.10.0"
authors = ["Robin Appelman <robin@icewind.nl>"] authors = ["Robin Appelman <robin@icewind.nl>"]
edition = "2018" edition = "2018"
description = "Reading bit sequences from a byte slice" description = "Reading bit sequences from a byte slice"
@ -10,7 +10,7 @@ repository = "https://github.com/icewind1991/bitbuffer"
[dependencies] [dependencies]
num-traits = "0.2" num-traits = "0.2"
err-derive = "0.3" err-derive = "0.3"
bitbuffer_derive = { version = "0.9", path = "bitbuffer_derive" } bitbuffer_derive = { version = "0.10", path = "bitbuffer_derive" }
memchr = "2" memchr = "2"
[dev-dependencies] [dev-dependencies]

View file

@ -1,6 +1,6 @@
[package] [package]
name = "bitbuffer_derive" name = "bitbuffer_derive"
version = "0.9.0" version = "0.10.0"
authors = ["Robin Appelman <robin@icewind.nl>"] authors = ["Robin Appelman <robin@icewind.nl>"]
edition = "2018" edition = "2018"
description = "Reading bit sequences from a byte slice" description = "Reading bit sequences from a byte slice"
@ -18,4 +18,4 @@ proc-macro2 = "1.0"
syn_util = "0.4" syn_util = "0.4"
[dev-dependencies] [dev-dependencies]
bitbuffer = { version = "0.9", path = ".." } bitbuffer = { version = "0.10", path = ".." }

View file

@ -3,22 +3,27 @@ use std::cmp::min;
use std::iter::{once, repeat}; use std::iter::{once, repeat};
use std::marker::PhantomData; use std::marker::PhantomData;
use std::mem::size_of; use std::mem::size_of;
use std::ops::Range;
const USIZE_BITS: usize = size_of::<usize>() * 8; const USIZE_BITS: usize = size_of::<usize>() * 8;
pub struct WriteBuffer<'a, E: Endianness>(CowWriteBuffer<'a, E>); pub struct WriteBuffer<'a, E: Endianness> {
bit_len: usize,
bytes: &'a mut Vec<u8>,
endianness: PhantomData<E>,
}
impl<'a, E: Endianness> WriteBuffer<'a, E> { impl<'a, E: Endianness> WriteBuffer<'a, E> {
pub fn new(bytes: &'a mut Vec<u8>, endianness: E) -> Self { pub fn new(bytes: &'a mut Vec<u8>, _endianness: E) -> Self {
WriteBuffer(CowWriteBuffer::ExpandBorrowed(ExpandWriteBuffer::new( WriteBuffer {
bytes, endianness, bit_len: 0,
))) bytes,
endianness: PhantomData,
}
} }
/// The number of written bits in the buffer /// The number of written bits in the buffer
pub fn bit_len(&self) -> usize { pub fn bit_len(&self) -> usize {
self.0.bit_len() self.bit_len
} }
pub fn push_non_fit_bits<I>(&mut self, bits: I, count: usize) pub fn push_non_fit_bits<I>(&mut self, bits: I, count: usize)
@ -44,104 +49,17 @@ impl<'a, E: Endianness> WriteBuffer<'a, E> {
/// Push up to an usize worth of bits /// Push up to an usize worth of bits
pub fn push_bits(&mut self, bits: usize, count: usize) { pub fn push_bits(&mut self, bits: usize, count: usize) {
if count > 0 { if count == 0 {
self.0.push_bits(bits, count) return;
} }
}
pub fn reserve(&mut self, length: usize) -> (WriteBuffer<E>, WriteBuffer<E>) {
let (head, tail) = self.0.reserve(length);
(WriteBuffer(head), WriteBuffer(tail))
}
}
enum CowWriteBuffer<'a, E: Endianness> {
FixedBorrowed(FixedWriteBuffer<'a, E>),
ExpandBorrowed(ExpandWriteBuffer<'a, E>),
}
impl<'a, E: Endianness> CowWriteBuffer<'a, E> {
/// The number of written bits in the buffer
fn bit_len(&self) -> usize {
match self {
CowWriteBuffer::FixedBorrowed(buffer) => buffer.bit_len(),
CowWriteBuffer::ExpandBorrowed(buffer) => buffer.bit_len(),
}
}
/// Push up to an usize worth of bits
fn push_bits(&mut self, bits: usize, count: usize) {
match self {
CowWriteBuffer::FixedBorrowed(buffer) => buffer.push_bits(bits, count),
CowWriteBuffer::ExpandBorrowed(buffer) => buffer.push_bits(bits, count),
}
}
/// Reserve some bits to be written later by splitting of two parts
fn reserve(&mut self, length: usize) -> (CowWriteBuffer<E>, CowWriteBuffer<E>) {
match self {
CowWriteBuffer::FixedBorrowed(buffer) => {
let (head, tail) = buffer.reserve(length);
(
CowWriteBuffer::FixedBorrowed(head),
CowWriteBuffer::FixedBorrowed(tail),
)
}
CowWriteBuffer::ExpandBorrowed(buffer) => {
let (head, tail) = buffer.reserve(length);
(
CowWriteBuffer::FixedBorrowed(head),
CowWriteBuffer::ExpandBorrowed(tail),
)
}
}
}
}
struct ExpandWriteBuffer<'a, E: Endianness> {
bit_len: usize,
bytes: *mut Vec<u8>,
endianness: PhantomData<E>,
lifetime: PhantomData<&'a u8>,
parent_bit_len: Option<&'a mut usize>,
}
impl<'a, E: Endianness> Drop for ExpandWriteBuffer<'a, E> {
fn drop(&mut self) {
if let Some(parent_bit_len) = self.parent_bit_len.take() {
*parent_bit_len = self.bit_len;
}
}
}
impl<'a, E: Endianness> ExpandWriteBuffer<'a, E> {
fn new(bytes: &'a mut Vec<u8>, _endianness: E) -> Self {
ExpandWriteBuffer {
bit_len: 0,
bytes: bytes as *mut _,
endianness: PhantomData,
lifetime: PhantomData,
parent_bit_len: None,
}
}
/// The number of written bits in the buffer
fn bit_len(&self) -> usize {
self.bit_len
}
/// Push up to an usize worth of bits
fn push_bits(&mut self, bits: usize, count: usize) {
debug_assert!(count < USIZE_BITS - 8); debug_assert!(count < USIZE_BITS - 8);
let bytes = unsafe { self.bytes.as_mut().unwrap() };
// ensure there are no stray bits // ensure there are no stray bits
let bits = bits & (usize::MAX >> (USIZE_BITS - count)); let bits = bits & (usize::MAX >> (USIZE_BITS - count));
let bit_offset = self.bit_len & 7; let bit_offset = self.bit_len & 7;
let last_written_byte = if bit_offset > 0 { let last_written_byte = if bit_offset > 0 {
bytes.pop().unwrap_or(0) self.bytes.pop().unwrap_or(0)
} else { } else {
0 0
}; };
@ -149,340 +67,30 @@ impl<'a, E: Endianness> ExpandWriteBuffer<'a, E> {
if E::is_le() { if E::is_le() {
let merged = last_written_byte as usize | bits << bit_offset; let merged = last_written_byte as usize | bits << bit_offset;
bytes.extend_from_slice(&merged.to_le_bytes()[0..merged_byte_count]); self.bytes
.extend_from_slice(&merged.to_le_bytes()[0..merged_byte_count]);
} else { } else {
let merged = ((last_written_byte as usize) << (USIZE_BITS - 8)) let merged = ((last_written_byte as usize) << (USIZE_BITS - 8))
| (bits << (USIZE_BITS - bit_offset - count)); | (bits << (USIZE_BITS - bit_offset - count));
bytes.extend_from_slice(&merged.to_be_bytes()[0..merged_byte_count]); self.bytes
.extend_from_slice(&merged.to_be_bytes()[0..merged_byte_count]);
} }
self.bit_len += count; self.bit_len += count;
} }
/// Reserve some bits to be written later by splitting of two parts pub fn set_at(&mut self, pos: usize, bits: u64, count: usize) {
/// debug_assert!(count < 64 - 8);
/// One fixed size part and one expanding part
fn reserve(&mut self, length: usize) -> (FixedWriteBuffer<E>, ExpandWriteBuffer<E>) {
let bit_offset = self.bit_len & 7;
let byte_index = self.bit_len / 8;
let end_byte = (self.bit_len + length + 7) / 8; let bit_offset = pos & 7;
let byte_pos = pos / 8;
let byte_count = (count + bit_offset + 7) / 8;
{ let mut old = [0; 8];
let bytes = unsafe { self.bytes.as_mut().unwrap() }; old[0..byte_count].copy_from_slice(&self.bytes[byte_pos..byte_pos + byte_count]);
bytes.resize(end_byte, 0);
} let old = u64::from_le_bytes(old);
self.bit_len += length; let merged = old | (bits << bit_offset);
( let merged = merged.to_le_bytes();
unsafe { self.bytes[byte_pos..byte_pos + byte_count].copy_from_slice(&merged[0..byte_count]);
FixedWriteBuffer::new(
self.bytes,
byte_index..end_byte,
bit_offset,
length + bit_offset,
E::endianness(),
)
},
ExpandWriteBuffer {
bit_len: self.bit_len,
bytes: self.bytes,
endianness: PhantomData,
lifetime: PhantomData,
parent_bit_len: Some(&mut self.bit_len),
},
)
} }
} }
#[test]
fn test_push_expand_be() {
use crate::BigEndian;
let mut buffer = vec![];
{
let mut write = ExpandWriteBuffer::new(&mut buffer, BigEndian);
write.push_bits(0b1101, 4);
write.push_bits(0b1, 1);
write.push_bits(0b0, 1);
write.push_bits(0b101_01010, 8);
}
assert_eq!(vec![0b1101_1_0_10, 0b101010_00], buffer)
}
#[test]
fn test_push_expand_le() {
use crate::LittleEndian;
let mut buffer = vec![];
{
let mut write = ExpandWriteBuffer::new(&mut buffer, LittleEndian);
write.push_bits(0b1101, 4);
write.push_bits(0b1, 1);
write.push_bits(0b0, 1);
write.push_bits(0b101_01010, 8);
}
assert_eq!(vec![0b10_0_1_1101, 0b00101010], buffer)
}
#[test]
fn test_push_expand_reserve_be() {
use crate::BigEndian;
let mut buffer = vec![];
{
let mut write = ExpandWriteBuffer::new(&mut buffer, BigEndian);
write.push_bits(0b1101, 4);
let (mut reserved, mut rest) = write.reserve(2);
rest.push_bits(0b101_01010, 8);
reserved.push_bits(0b1, 1);
reserved.push_bits(0b0, 1);
}
assert_eq!(vec![0b1101_1_0_10, 0b101010_00], buffer)
}
#[test]
fn test_push_expand_reserve_le() {
use crate::LittleEndian;
let mut buffer = vec![];
{
let mut write = ExpandWriteBuffer::new(&mut buffer, LittleEndian);
write.push_bits(0b1101, 4);
let (mut reserved, mut rest) = write.reserve(2);
rest.push_bits(0b101_01010, 8);
reserved.push_bits(0b1, 1);
reserved.push_bits(0b0, 1);
}
assert_eq!(vec![0b10_0_1_1101, 0b00101010], buffer)
}
#[test]
fn test_push_expand_reserve_byte_wrap_le() {
use crate::LittleEndian;
let mut buffer = vec![];
{
let mut write = ExpandWriteBuffer::new(&mut buffer, LittleEndian);
write.push_bits(0b0101101, 7);
let (mut reserved, mut rest) = write.reserve(20);
rest.push_bits(0b101_01010, 5);
reserved.push_bits(40, 20);
}
assert_eq!(vec![0b0_0101101, 0b0001010_0, 0, 80], buffer)
}
#[test]
fn test_push_expand_reserve_resize() {
use crate::LittleEndian;
let mut buffer = Vec::with_capacity(1);
{
let mut write = ExpandWriteBuffer::new(&mut buffer, LittleEndian);
write.push_bits(0b1101, 4);
let (mut reserved, mut rest) = write.reserve(2);
rest.push_bits(0b10, 2);
// trigger a resize/reallocation
for _ in 0..128 {
rest.push_bits(0x55, 8);
}
reserved.push_bits(0b1, 1);
reserved.push_bits(0b0, 1);
}
assert_eq!([0b10_0_1_1101, 0x55], buffer[0..2])
}
struct FixedWriteBuffer<'a, E: Endianness> {
bit_start: usize,
bit_len: usize,
bytes: *mut Vec<u8>,
byte_range: Range<usize>,
endianness: PhantomData<E>,
lifetime: PhantomData<&'a u8>,
bit_size: usize,
}
impl<'a, E: Endianness> FixedWriteBuffer<'a, E> {
/// Safety: ensure that nobody else writes the bits
///
/// from byte_range.start * 8 + bit_start
/// to byte_range.start * 8 + bit_size
///
/// or reads from those bits before this is dropped
unsafe fn new(
bytes: *mut Vec<u8>,
byte_range: Range<usize>,
bit_start: usize,
bit_size: usize,
_endianness: E,
) -> Self {
FixedWriteBuffer {
bit_start,
bit_len: bit_start,
bytes,
byte_range,
endianness: PhantomData,
lifetime: PhantomData,
bit_size,
}
}
/// The number of written bits in the buffer
fn bit_len(&self) -> usize {
self.bit_len - self.bit_start
}
/// Push up to an usize worth of bits
fn push_bits(&mut self, bits: usize, count: usize) {
if count == 0 {
return;
}
debug_assert!(count < USIZE_BITS - 8);
assert!(self.bit_len + count <= self.bit_size);
// ensure there are no stray bits
let bits = bits & (usize::MAX >> (USIZE_BITS - count));
let bit_offset = self.bit_len & 7;
let byte_index = self.bit_len / 8;
let merged_byte_count = (count + bit_offset + 7) / 8;
// get a mut slice from our ptr
// this is safe because
//
// 1. the other half of the `reserve` output can't write to our reserved bits
// 2. the underlying vec can only be used again after both parts have been dropped
let bytes = unsafe {
let bytes = self.bytes.as_mut().unwrap();
let ptr = bytes[self.byte_range.clone()].as_ptr() as *mut u8;
std::slice::from_raw_parts_mut(ptr, self.byte_range.len())
};
let mut source = [0; USIZE_BITS / 8];
source[0..merged_byte_count]
.copy_from_slice(&bytes[byte_index..byte_index + merged_byte_count]);
let last_written = usize::from_le_bytes(source);
if E::is_le() {
let merged = last_written as usize | bits << bit_offset;
bytes[byte_index..byte_index + merged_byte_count]
.copy_from_slice(&merged.to_le_bytes()[0..merged_byte_count]);
} else {
let merged = ((last_written as usize) << (USIZE_BITS - 8))
| (bits << (USIZE_BITS - bit_offset - count));
bytes[byte_index..byte_index + merged_byte_count]
.copy_from_slice(&merged.to_be_bytes()[0..merged_byte_count]);
}
self.bit_len += count;
}
fn reserve(&mut self, length: usize) -> (FixedWriteBuffer<E>, FixedWriteBuffer<E>) {
assert!(self.bit_len + length <= self.bit_size);
let byte_count = (length + 7) / 8;
let bit_offset = self.bit_len & 7;
let byte_index = self.bit_len / 8;
self.bit_len += length;
unsafe {
(
FixedWriteBuffer::new(
self.bytes,
self.byte_range.start + byte_index
..self.byte_range.start + byte_index + byte_count,
bit_offset,
length + bit_offset,
E::endianness(),
),
FixedWriteBuffer::new(
self.bytes,
self.byte_range.clone(),
self.bit_len,
self.bit_size,
E::endianness(),
),
)
}
}
}
#[test]
fn test_push_fixed_be() {
use crate::BigEndian;
let mut buffer = vec![0; 2];
let mut write = unsafe { FixedWriteBuffer::new(&mut buffer, 0..2, 0, 16, BigEndian) };
write.push_bits(0b1101, 4);
assert_eq!(4, write.bit_len());
write.push_bits(0b1, 1);
assert_eq!(5, write.bit_len());
write.push_bits(0b0, 1);
assert_eq!(6, write.bit_len());
write.push_bits(0b101_01010, 8);
assert_eq!(vec![0b1101_1_0_10, 0b101010_00], buffer)
}
#[test]
fn test_push_fixed_le() {
use crate::LittleEndian;
let mut buffer = vec![0; 2];
let mut write = unsafe { FixedWriteBuffer::new(&mut buffer, 0..2, 0, 16, LittleEndian) };
write.push_bits(0b1101, 4);
assert_eq!(4, write.bit_len());
write.push_bits(0b1, 1);
assert_eq!(5, write.bit_len());
write.push_bits(0b0, 1);
assert_eq!(6, write.bit_len());
write.push_bits(0b101_01010, 8);
assert_eq!(vec![0b10_0_1_1101, 0b00101010], buffer)
}
#[test]
fn test_push_fixed_reserve_be() {
use crate::BigEndian;
let mut buffer = vec![0; 2];
let mut write = unsafe { FixedWriteBuffer::new(&mut buffer, 0..2, 0, 16, BigEndian) };
write.push_bits(0b1101, 4);
let (mut reserved, mut rest) = write.reserve(2);
rest.push_bits(0b101_01010, 8);
reserved.push_bits(0b1, 1);
reserved.push_bits(0b0, 1);
assert_eq!(vec![0b1101_1_0_10, 0b101010_00], buffer)
}
#[test]
fn test_push_fixed_reserve_le() {
use crate::LittleEndian;
let mut buffer = vec![0; 2];
let mut write = unsafe { FixedWriteBuffer::new(&mut buffer, 0..2, 0, 16, LittleEndian) };
write.push_bits(0b1101, 4);
let (mut reserved, mut rest) = write.reserve(2);
rest.push_bits(0b101_01010, 8);
reserved.push_bits(0b1, 1);
reserved.push_bits(0b0, 1);
assert_eq!(vec![0b10_0_1_1101, 0b00101010], buffer)
}

View file

@ -288,23 +288,6 @@ where
value.write_sized(self, length) value.write_sized(self, length)
} }
/// Reserve some bits to be written later by splitting of two parts
///
/// This allows skipping a few bits to write later
pub fn reserve<
Err,
F: Fn(&mut BitWriteStream<E>, &mut BitWriteStream<E>) -> Result<(), Err>,
>(
&mut self,
count: usize,
body_fn: F,
) -> Result<(), Err> {
let (head, tail) = self.buffer.reserve(count);
let mut head = BitWriteStream { buffer: head };
let mut tail = BitWriteStream { buffer: tail };
body_fn(&mut head, &mut tail)
}
/// Write the length of a section before the section /// Write the length of a section before the section
pub fn reserve_length<Err: From<BitError>, F: Fn(&mut BitWriteStream<E>) -> Result<(), Err>>( pub fn reserve_length<Err: From<BitError>, F: Fn(&mut BitWriteStream<E>) -> Result<(), Err>>(
&mut self, &mut self,
@ -345,15 +328,15 @@ where
/// Reserve the length to write an integer /// Reserve the length to write an integer
pub fn reserve_int<Err: From<BitError>, F: Fn(&mut BitWriteStream<E>) -> Result<u64, Err>>( pub fn reserve_int<Err: From<BitError>, F: Fn(&mut BitWriteStream<E>) -> Result<u64, Err>>(
&mut self, &mut self,
length_bit_size: usize, count: usize,
body_fn: F, body_fn: F,
) -> Result<(), Err> { ) -> Result<(), Err> {
let (head, tail) = self.buffer.reserve(length_bit_size); let start = self.bit_len();
let mut head = BitWriteStream { buffer: head }; self.write_int(0u64, count)?;
let mut tail = BitWriteStream { buffer: tail };
let head_int = body_fn(self)?;
self.buffer.set_at(start, head_int, count);
let head_int = body_fn(&mut tail)?;
head.write_sized(&head_int, length_bit_size)?;
Ok(()) Ok(())
} }
} }