Skip to main content

vortex_buffer/bit/
buf_mut.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::ops::Not;
5
6use bitvec::view::BitView;
7
8use crate::BitBuffer;
9use crate::BufferMut;
10use crate::ByteBufferMut;
11use crate::bit::collect_bool_words;
12use crate::bit::get_bit_unchecked;
13use crate::bit::ops;
14use crate::bit::set_bit_unchecked;
15use crate::bit::unset_bit_unchecked;
16use crate::buffer_mut;
17
18/// Sets all bits in the bit-range `[start_bit, end_bit)` of `slice` to `value`.
19#[inline(always)]
20pub(crate) fn fill_bits(slice: &mut [u8], start_bit: usize, end_bit: usize, value: bool) {
21    if start_bit >= end_bit {
22        return;
23    }
24
25    let fill_byte: u8 = if value { 0xFF } else { 0x00 };
26
27    let start_byte = start_bit / 8;
28    let start_rem = start_bit % 8;
29    let end_byte = end_bit / 8;
30    let end_rem = end_bit % 8;
31
32    if start_byte == end_byte {
33        // All bits are in the same byte
34        let mask = ((1u8 << (end_rem - start_rem)) - 1) << start_rem;
35        if value {
36            slice[start_byte] |= mask;
37        } else {
38            slice[start_byte] &= !mask;
39        }
40    } else {
41        // First partial byte
42        if start_rem != 0 {
43            let mask = !((1u8 << start_rem) - 1);
44            if value {
45                slice[start_byte] |= mask;
46            } else {
47                slice[start_byte] &= !mask;
48            }
49        }
50
51        // Middle bytes
52        let fill_start = if start_rem != 0 {
53            start_byte + 1
54        } else {
55            start_byte
56        };
57        if fill_start < end_byte {
58            slice[fill_start..end_byte].fill(fill_byte);
59        }
60
61        // Last partial byte
62        if end_rem != 0 {
63            let mask = (1u8 << end_rem) - 1;
64            if value {
65                slice[end_byte] |= mask;
66            } else {
67                slice[end_byte] &= !mask;
68            }
69        }
70    }
71}
72
73/// A mutable bitset buffer that allows random access to individual bits for set and get.
74///
75///
76/// # Example
77/// ```
78/// use vortex_buffer::BitBufferMut;
79///
80/// let mut bools = BitBufferMut::new_unset(10);
81/// bools.set_to(9, true);
82/// for i in 0..9 {
83///    assert!(!bools.value(i));
84/// }
85/// assert!(bools.value(9));
86///
87/// // Freeze into a new bools vector.
88/// let bools = bools.freeze();
89/// ```
90///
91/// See also: [`BitBuffer`].
92#[derive(Debug, Clone)]
93pub struct BitBufferMut {
94    buffer: ByteBufferMut,
95    /// Represents the offset of the bit buffer into the first byte.
96    ///
97    /// This is always less than 8 (for when the bit buffer is not aligned to a byte).
98    offset: usize,
99    len: usize,
100}
101
102impl BitBufferMut {
103    /// Create new bit buffer from given byte buffer and logical bit length
104    pub fn from_buffer(buffer: ByteBufferMut, offset: usize, len: usize) -> Self {
105        assert!(
106            len <= buffer.len() * 8,
107            "Buffer len {} is too short for the given length {len}",
108            buffer.len()
109        );
110        Self {
111            buffer,
112            offset,
113            len,
114        }
115    }
116
117    /// Creates a `BitBufferMut` from a [`BitBuffer`] by copying all of the data over.
118    pub fn copy_from(bit_buffer: &BitBuffer) -> Self {
119        Self {
120            buffer: ByteBufferMut::copy_from(bit_buffer.inner()),
121            offset: bit_buffer.offset(),
122            len: bit_buffer.len(),
123        }
124    }
125
126    /// Create a new empty mutable bit buffer with requested capacity (in bits).
127    pub fn with_capacity(capacity: usize) -> Self {
128        Self {
129            buffer: BufferMut::with_capacity(capacity.div_ceil(8)),
130            offset: 0,
131            len: 0,
132        }
133    }
134
135    /// Create a new mutable buffer with requested `len` and all bits set to `true`.
136    pub fn new_set(len: usize) -> Self {
137        Self {
138            buffer: buffer_mut![0xFF; len.div_ceil(8)],
139            offset: 0,
140            len,
141        }
142    }
143
144    /// Create a new mutable buffer with requested `len` and all bits set to `false`.
145    pub fn new_unset(len: usize) -> Self {
146        Self {
147            buffer: BufferMut::zeroed(len.div_ceil(8)),
148            offset: 0,
149            len,
150        }
151    }
152
153    /// Create a new empty `BitBufferMut`.
154    #[inline(always)]
155    pub fn empty() -> Self {
156        Self::with_capacity(0)
157    }
158
159    /// Create a new mutable buffer with requested `len` and all bits set to `value`.
160    pub fn full(value: bool, len: usize) -> Self {
161        if value {
162            Self::new_set(len)
163        } else {
164            Self::new_unset(len)
165        }
166    }
167
168    /// Create a bit buffer of `len` with `indices` set as true.
169    pub fn from_indices(len: usize, indices: impl IntoIterator<Item = usize>) -> BitBufferMut {
170        let mut buffer = BufferMut::<u64>::zeroed(len.div_ceil(64));
171        for idx in indices {
172            assert!(idx < len, "index {idx} exceeds len {len}");
173            buffer.as_mut_slice()[idx / 64] |= 1 << (idx % 64);
174        }
175
176        let mut buffer = buffer.into_byte_buffer();
177        buffer.truncate(len.div_ceil(8));
178
179        Self {
180            buffer,
181            offset: 0,
182            len,
183        }
184    }
185
186    /// Invokes `f` with indexes `0..len` collecting the boolean results into a new `BitBufferMut`
187    #[inline]
188    pub fn collect_bool<F: FnMut(usize) -> bool>(len: usize, f: F) -> Self {
189        let num_words = len.div_ceil(64);
190        let mut buffer: BufferMut<u64> = BufferMut::with_capacity(num_words);
191        // SAFETY: `collect_bool_words` writes every word in `0..num_words` below
192        // before any read; `u64` has no invalid bit patterns and the assignments
193        // inside `collect_bool_words` are pure writes.
194        unsafe { buffer.set_len(num_words) };
195        collect_bool_words(buffer.as_mut_slice(), len, f);
196
197        let mut bytes = buffer.into_byte_buffer();
198        bytes.truncate(len.div_ceil(8));
199
200        Self {
201            buffer: bytes,
202            offset: 0,
203            len,
204        }
205    }
206
207    /// Return the underlying byte buffer.
208    pub fn inner(&self) -> &ByteBufferMut {
209        &self.buffer
210    }
211
212    /// Consumes the buffer and return the underlying byte buffer.
213    pub fn into_inner(self) -> ByteBufferMut {
214        self.buffer
215    }
216
217    /// Get the current populated length of the buffer.
218    #[inline(always)]
219    pub fn len(&self) -> usize {
220        self.len
221    }
222
223    /// True if the buffer has length 0.
224    #[inline(always)]
225    pub fn is_empty(&self) -> bool {
226        self.len == 0
227    }
228
229    /// Get the current bit offset of the buffer.
230    #[inline(always)]
231    pub fn offset(&self) -> usize {
232        self.offset
233    }
234
235    /// Get the value at the requested index.
236    #[inline(always)]
237    pub fn value(&self, index: usize) -> bool {
238        assert!(index < self.len);
239        // SAFETY: checked by assertion
240        unsafe { self.value_unchecked(index) }
241    }
242
243    /// Get the value at the requested index without bounds checking.
244    ///
245    /// # Safety
246    ///
247    /// The caller must ensure that `index` is less than the length of the buffer.
248    #[inline(always)]
249    pub unsafe fn value_unchecked(&self, index: usize) -> bool {
250        unsafe { get_bit_unchecked(self.buffer.as_ptr(), self.offset + index) }
251    }
252
253    /// Get the bit capacity of the buffer.
254    #[inline(always)]
255    pub fn capacity(&self) -> usize {
256        (self.buffer.capacity() * 8) - self.offset
257    }
258
259    /// Reserve additional bit capacity for the buffer.
260    pub fn reserve(&mut self, additional: usize) {
261        let required_bits = self.offset + self.len + additional;
262        let required_bytes = required_bits.div_ceil(8); // Rounds up.
263
264        let additional_bytes = required_bytes.saturating_sub(self.buffer.len());
265        self.buffer.reserve(additional_bytes);
266    }
267
268    /// Clears the bit buffer (but keeps any allocated memory).
269    pub fn clear(&mut self) {
270        // Since there are no items we need to drop, we simply set the length to 0.
271        self.len = 0;
272        self.offset = 0;
273    }
274
275    /// Set the bit at `index` to the given boolean value.
276    ///
277    /// This operation is checked so if `index` exceeds the buffer length, this will panic.
278    #[inline]
279    pub fn set_to(&mut self, index: usize, value: bool) {
280        if value {
281            self.set(index);
282        } else {
283            self.unset(index);
284        }
285    }
286
287    /// Set the bit at `index` to the given boolean value without checking bounds.
288    ///
289    /// # Safety
290    ///
291    /// The caller must ensure that `index` does not exceed the largest bit index in the backing buffer.
292    #[inline]
293    pub unsafe fn set_to_unchecked(&mut self, index: usize, value: bool) {
294        if value {
295            // SAFETY: checked by caller
296            unsafe { self.set_unchecked(index) }
297        } else {
298            // SAFETY: checked by caller
299            unsafe { self.unset_unchecked(index) }
300        }
301    }
302
303    /// Set a position to `true`.
304    ///
305    /// This operation is checked so if `index` exceeds the buffer length, this will panic.
306    #[inline]
307    pub fn set(&mut self, index: usize) {
308        assert!(index < self.len, "index {index} exceeds len {}", self.len);
309
310        // SAFETY: checked by assertion
311        unsafe { self.set_unchecked(index) };
312    }
313
314    /// Set a position to `false`.
315    ///
316    /// This operation is checked so if `index` exceeds the buffer length, this will panic.
317    #[inline]
318    pub fn unset(&mut self, index: usize) {
319        assert!(index < self.len, "index {index} exceeds len {}", self.len);
320
321        // SAFETY: checked by assertion
322        unsafe { self.unset_unchecked(index) };
323    }
324
325    /// Set the bit at `index` to `true` without checking bounds.
326    ///
327    /// Note: Do not call this in a tight loop. Prefer to use [`set_bit_unchecked`].
328    ///
329    /// # Safety
330    ///
331    /// The caller must ensure that `index` does not exceed the largest bit index in the backing buffer.
332    #[inline]
333    pub unsafe fn set_unchecked(&mut self, index: usize) {
334        // SAFETY: checked by caller
335        unsafe { set_bit_unchecked(self.buffer.as_mut_ptr(), self.offset + index) }
336    }
337
338    /// Unset the bit at `index` without checking bounds.
339    ///
340    /// Note: Do not call this in a tight loop. Prefer to use [`unset_bit_unchecked`].
341    ///
342    /// # Safety
343    ///
344    /// The caller must ensure that `index` does not exceed the largest bit index in the backing buffer.
345    #[inline]
346    pub unsafe fn unset_unchecked(&mut self, index: usize) {
347        // SAFETY: checked by caller
348        unsafe { unset_bit_unchecked(self.buffer.as_mut_ptr(), self.offset + index) }
349    }
350
351    /// Foces the length of the `BitBufferMut` to `new_len`.
352    ///
353    /// # Safety
354    ///
355    /// - `new_len` must be less than or equal to [`capacity()`](Self::capacity)
356    /// - The elements at `old_len..new_len` must be initialized
357    #[inline(always)]
358    pub unsafe fn set_len(&mut self, new_len: usize) {
359        debug_assert!(
360            new_len <= self.capacity(),
361            "`set_len` requires that new_len <= capacity()"
362        );
363
364        // Calculate the new byte length required to hold the bits
365        let bytes_len = (self.offset + new_len).div_ceil(8);
366        unsafe { self.buffer.set_len(bytes_len) };
367
368        self.len = new_len;
369    }
370
371    /// Truncate the buffer to the given length.
372    ///
373    /// If the given length is greater than the current length, this is a no-op.
374    pub fn truncate(&mut self, len: usize) {
375        if len > self.len {
376            return;
377        }
378
379        let end_bit = self.offset + len;
380        let new_len_bytes = end_bit.div_ceil(8);
381        self.buffer.truncate(new_len_bytes);
382        self.len = len;
383
384        // Clear stale bits in the final partial byte so the "bits beyond len are zero" invariant
385        // holds. `append_false` (and `append_buffer`) rely on it to avoid a read-modify-write.
386        if !end_bit.is_multiple_of(8) {
387            let keep = (1u8 << (end_bit % 8)) - 1;
388            self.buffer.as_mut_slice()[new_len_bytes - 1] &= keep;
389        }
390    }
391
392    /// Append a new boolean into the bit buffer, incrementing the length.
393    #[inline]
394    pub fn append(&mut self, value: bool) {
395        if value {
396            self.append_true()
397        } else {
398            self.append_false()
399        }
400    }
401
402    /// Append a new true value to the buffer.
403    #[inline]
404    pub fn append_true(&mut self) {
405        let bit_pos = self.offset + self.len;
406        let byte_pos = bit_pos / 8;
407        let bit_in_byte = bit_pos % 8;
408
409        // Ensure buffer has enough bytes
410        if byte_pos >= self.buffer.len() {
411            self.buffer.push(0u8);
412        }
413
414        // Set the bit
415        self.buffer.as_mut_slice()[byte_pos] |= 1 << bit_in_byte;
416        self.len += 1;
417    }
418
419    /// Append a new false value to the buffer.
420    #[inline]
421    pub fn append_false(&mut self) {
422        let bit_pos = self.offset + self.len;
423        let byte_pos = bit_pos / 8;
424
425        // Ensure buffer has enough bytes (pushed as 0x00, so bit is already unset).
426        if byte_pos >= self.buffer.len() {
427            self.buffer.push(0u8);
428        }
429
430        // The bit is guaranteed to be 0: new bytes are zero-initialized, and
431        // existing bytes have this bit unset (it's beyond the current length).
432        self.len += 1;
433    }
434
435    /// Append several boolean values into the bit buffer. After this operation,
436    /// the length will be incremented by `n`.
437    ///
438    /// Panics if the buffer does not have `n` slots left.
439    #[inline]
440    pub fn append_n(&mut self, value: bool, n: usize) {
441        if n == 0 {
442            return;
443        }
444
445        let end_bit_pos = self.offset + self.len + n;
446        let required_bytes = end_bit_pos.div_ceil(8);
447
448        // Ensure buffer has enough bytes
449        if required_bytes > self.buffer.len() {
450            self.buffer.push_n(0x00, required_bytes - self.buffer.len());
451        }
452
453        let start = self.len;
454        self.len += n;
455        self.fill_range(start, self.len, value);
456    }
457
458    /// Sets all bits in the range `[start, end)` to `value`.
459    ///
460    /// This operates on an arbitrary range within the existing length of the buffer.
461    /// Panics if `end > self.len` or `start > end`.
462    #[inline(always)]
463    pub fn fill_range(&mut self, start: usize, end: usize, value: bool) {
464        assert!(end <= self.len, "end {end} exceeds len {}", self.len);
465        assert!(start <= end, "start {start} exceeds end {end}");
466
467        // SAFETY: assertions above guarantee start <= end <= self.len,
468        // so offset + end fits within the buffer.
469        unsafe { self.fill_range_unchecked(start, end, value) }
470    }
471
472    /// Sets all bits in the range `[start, end)` to `value` without bounds checking.
473    ///
474    /// # Safety
475    ///
476    /// The caller must ensure that `start <= end <= self.len`.
477    #[inline(always)]
478    pub unsafe fn fill_range_unchecked(&mut self, start: usize, end: usize, value: bool) {
479        fill_bits(
480            self.buffer.as_mut_slice(),
481            self.offset + start,
482            self.offset + end,
483            value,
484        );
485    }
486
487    /// Append a [`BitBuffer`] to this [`BitBufferMut`]
488    ///
489    /// This efficiently copies all bits from the source buffer to the end of this buffer.
490    pub fn append_buffer(&mut self, buffer: &BitBuffer) {
491        let bit_len = buffer.len();
492        if bit_len == 0 {
493            return;
494        }
495
496        let start_bit_pos = self.offset + self.len;
497        let end_bit_pos = start_bit_pos + bit_len;
498        let required_bytes = end_bit_pos.div_ceil(8);
499
500        // Ensure buffer has enough bytes, zero-initialized for OR-based writes.
501        if required_bytes > self.buffer.len() {
502            self.buffer.push_n(0x00, required_bytes - self.buffer.len());
503        }
504
505        let dst_bit_offset = start_bit_pos % 8;
506        let src_bit_offset = buffer.offset();
507
508        if dst_bit_offset == 0 && src_bit_offset == 0 {
509            // Both byte-aligned: use memcpy for full bytes, then mask the tail.
510            let dst_byte = start_bit_pos / 8;
511            let src_bytes = buffer.inner().as_slice();
512            let full_bytes = bit_len / 8;
513            self.buffer.as_mut_slice()[dst_byte..dst_byte + full_bytes]
514                .copy_from_slice(&src_bytes[..full_bytes]);
515            let rem = bit_len % 8;
516            if rem != 0 {
517                let mask = (1u8 << rem) - 1;
518                self.buffer.as_mut_slice()[dst_byte + full_bytes] |= src_bytes[full_bytes] & mask;
519            }
520        } else {
521            // Use bitvec for unaligned bit copying.
522            let self_slice = self
523                .buffer
524                .as_mut_slice()
525                .view_bits_mut::<bitvec::prelude::Lsb0>();
526            let other_slice = buffer
527                .inner()
528                .as_slice()
529                .view_bits::<bitvec::prelude::Lsb0>();
530            let source_range = src_bit_offset..src_bit_offset + bit_len;
531            self_slice[start_bit_pos..end_bit_pos].copy_from_bitslice(&other_slice[source_range]);
532        }
533
534        self.len += bit_len;
535    }
536
537    /// Absorbs a mutable buffer that was previously split off.
538    ///
539    /// If the two buffers were previously contiguous and not mutated in a way that causes
540    /// re-allocation i.e., if other was created by calling split_off on this buffer, then this is
541    /// an O(1) operation that just decreases a reference count and sets a few indices.
542    ///
543    /// Otherwise, this method degenerates to self.append_buffer(&other).
544    pub fn unsplit(&mut self, other: Self) {
545        if (self.offset + self.len).is_multiple_of(8) && other.offset == 0 {
546            // We are aligned and can just append the buffers
547            self.buffer.unsplit(other.buffer);
548            self.len += other.len;
549            return;
550        }
551
552        // Otherwise, we need to append the bits one by one
553        self.append_buffer(&other.freeze())
554    }
555
556    /// Freeze the buffer in its current state into an immutable `BoolBuffer`.
557    pub fn freeze(self) -> BitBuffer {
558        BitBuffer::new_with_offset(self.buffer.freeze(), self.len, self.offset)
559    }
560
561    /// Get the underlying bytes as a slice
562    pub fn as_slice(&self) -> &[u8] {
563        self.buffer.as_slice()
564    }
565
566    /// Get the underlying bytes as a mutable slice
567    pub fn as_mut_slice(&mut self) -> &mut [u8] {
568        self.buffer.as_mut_slice()
569    }
570}
571
572impl Default for BitBufferMut {
573    fn default() -> Self {
574        Self::with_capacity(0)
575    }
576}
577
578// Mutate-in-place implementation of bitwise NOT.
579impl Not for BitBufferMut {
580    type Output = BitBufferMut;
581
582    #[inline]
583    fn not(mut self) -> Self::Output {
584        ops::bitwise_unary_op_mut(&mut self, |b| !b);
585        self
586    }
587}
588
589impl From<&[bool]> for BitBufferMut {
590    fn from(value: &[bool]) -> Self {
591        BitBufferMut::collect_bool(value.len(), |i| value[i])
592    }
593}
594
595// allow building a buffer from a set of truthy byte values.
596impl From<&[u8]> for BitBufferMut {
597    fn from(value: &[u8]) -> Self {
598        BitBufferMut::collect_bool(value.len(), |i| value[i] > 0)
599    }
600}
601
602impl From<Vec<bool>> for BitBufferMut {
603    fn from(value: Vec<bool>) -> Self {
604        value.as_slice().into()
605    }
606}
607
608impl FromIterator<bool> for BitBufferMut {
609    #[inline]
610    fn from_iter<T: IntoIterator<Item = bool>>(iter: T) -> Self {
611        let mut iter = iter.into_iter();
612
613        // Since we do not know the length of the iterator, we can only guess how much memory we
614        // need to reserve. Note that these hints may be inaccurate.
615        let (lower_bound, _) = iter.size_hint();
616
617        // We choose not to use the optional upper bound size hint to match the standard library.
618
619        // Initialize all bits to 0 with the given length. By doing this, we only need to set bits
620        // that are true (and this is faster from benchmarks).
621        let mut buf = BitBufferMut::new_unset(lower_bound);
622        assert_eq!(buf.offset, 0);
623
624        // Directly write within our known capacity.
625        let ptr = buf.buffer.as_mut_ptr();
626        for i in 0..lower_bound {
627            let Some(v) = iter.next() else {
628                // SAFETY: We are definitely under the capacity and all values are already
629                // initialized from `new_unset`.
630                unsafe { buf.set_len(i) };
631                return buf;
632            };
633
634            if v {
635                // SAFETY: We have ensured that we are within the capacity.
636                unsafe { set_bit_unchecked(ptr, i) }
637            }
638        }
639
640        // Append any remaining items one at a time, as we do not know how many more there are.
641        // (`append` is already a single branch + bit set, see `append_true`/`append_false`.)
642        for v in iter {
643            buf.append(v);
644        }
645
646        buf
647    }
648}
649
650#[cfg(test)]
651mod tests {
652    use crate::BufferMut;
653    use crate::bit::buf_mut::BitBufferMut;
654    use crate::bitbuffer;
655    use crate::bitbuffer_mut;
656    use crate::buffer_mut;
657
658    #[test]
659    fn test_bits_mut() {
660        let mut bools = bitbuffer_mut![false; 10];
661        bools.set_to(0, true);
662        bools.set_to(9, true);
663
664        let bools = bools.freeze();
665        assert!(bools.value(0));
666        for i in 1..=8 {
667            assert!(!bools.value(i));
668        }
669        assert!(bools.value(9));
670    }
671
672    #[test]
673    fn test_append_n() {
674        let mut bools = BitBufferMut::with_capacity(10);
675        assert_eq!(bools.len(), 0);
676        assert!(bools.is_empty());
677
678        bools.append(true);
679        bools.append_n(false, 8);
680        bools.append_n(true, 1);
681
682        let bools = bools.freeze();
683
684        assert_eq!(bools.true_count(), 2);
685        assert!(bools.value(0));
686        assert!(bools.value(9));
687    }
688
689    #[test]
690    fn append_false_after_truncate_reads_back_false() {
691        // `truncate` leaves stale bits in the final partial byte; a subsequent `append_false`
692        // must still read back as false. Regression test for the `append_false` fast path.
693        let mut bools = BitBufferMut::new_set(16);
694        bools.truncate(12);
695        bools.append_false();
696        bools.append_true();
697
698        let bools = bools.freeze();
699        assert_eq!(bools.len(), 14);
700        assert!(
701            !bools.value(12),
702            "appended false must read back false after truncate"
703        );
704        assert!(bools.value(13));
705    }
706
707    #[test]
708    fn test_reserve_ensures_len_plus_additional() {
709        // This test documents the fix for the bug where reserve was incorrectly
710        // calculating additional bytes from capacity instead of len.
711
712        let mut bits = BitBufferMut::with_capacity(10);
713        assert_eq!(bits.len(), 0);
714
715        bits.reserve(100);
716
717        // Should have capacity for at least len + 100 = 0 + 100 = 100 bits.
718        assert!(bits.capacity() >= 100);
719
720        bits.append_n(true, 50);
721        assert_eq!(bits.len(), 50);
722
723        bits.reserve(100);
724
725        // Should have capacity for at least len + 100 = 50 + 100 = 150 bits.
726        assert!(bits.capacity() >= 150);
727    }
728
729    #[test]
730    fn test_with_offset_zero() {
731        // Test basic operations when offset is 0
732        let buf = BufferMut::zeroed(2);
733        let mut bit_buf = BitBufferMut::from_buffer(buf, 0, 16);
734
735        // Set some bits
736        bit_buf.set(0);
737        bit_buf.set(7);
738        bit_buf.set(8);
739        bit_buf.set(15);
740
741        // Verify values
742        assert!(bit_buf.value(0));
743        assert!(bit_buf.value(7));
744        assert!(bit_buf.value(8));
745        assert!(bit_buf.value(15));
746        assert!(!bit_buf.value(1));
747        assert!(!bit_buf.value(9));
748
749        // Verify underlying bytes
750        assert_eq!(bit_buf.as_slice()[0], 0b10000001);
751        assert_eq!(bit_buf.as_slice()[1], 0b10000001);
752    }
753
754    #[test]
755    fn test_with_offset_within_byte() {
756        // Test operations with offset=3 (within first byte)
757        let buf = buffer_mut![0b11111111, 0b00000000, 0b00000000];
758        let mut bit_buf = BitBufferMut::from_buffer(buf, 3, 10);
759
760        // Initially, bits 3-7 from first byte are set (5 bits)
761        // and bits 0-4 from second byte are unset (5 bits more)
762        assert!(bit_buf.value(0)); // bit 3 of byte 0
763        assert!(bit_buf.value(4)); // bit 7 of byte 0
764        assert!(!bit_buf.value(5)); // bit 0 of byte 1
765
766        // Set a bit in the second byte's range
767        bit_buf.set(7);
768        assert!(bit_buf.value(7));
769
770        // Unset a bit in the first byte's range
771        bit_buf.unset(0);
772        assert!(!bit_buf.value(0));
773    }
774
775    #[test]
776    fn test_with_offset_byte_boundary() {
777        // Test operations with offset=8 (exactly one byte)
778        let buf = buffer_mut![0xFF, 0x00, 0xFF];
779        let mut bit_buf = BitBufferMut::from_buffer(buf, 8, 16);
780
781        // Buffer starts at byte 1, so all bits should be unset initially
782        for i in 0..8 {
783            assert!(!bit_buf.value(i));
784        }
785        // Next byte has all bits set
786        for i in 8..16 {
787            assert!(bit_buf.value(i));
788        }
789
790        // Set some bits
791        bit_buf.set(0);
792        bit_buf.set(3);
793        assert!(bit_buf.value(0));
794        assert!(bit_buf.value(3));
795    }
796
797    #[test]
798    fn test_with_large_offset() {
799        // Test with offset=13 (one byte + 5 bits)
800        let buf = buffer_mut![0xFF, 0xFF, 0xFF, 0xFF];
801        let mut bit_buf = BitBufferMut::from_buffer(buf, 13, 10);
802
803        // All bits should initially be set
804        for i in 0..10 {
805            assert!(bit_buf.value(i));
806        }
807
808        // Unset some bits
809        bit_buf.unset(0);
810        bit_buf.unset(5);
811        bit_buf.unset(9);
812
813        assert!(!bit_buf.value(0));
814        assert!(bit_buf.value(1));
815        assert!(!bit_buf.value(5));
816        assert!(!bit_buf.value(9));
817    }
818
819    #[test]
820    fn test_append_with_offset() {
821        // Create buffer with offset
822        let buf = buffer_mut![0b11100000]; // First 3 bits unset, last 5 set
823        let mut bit_buf = BitBufferMut::from_buffer(buf, 3, 0); // Start at bit 3, len=0
824
825        // Append some bits
826        bit_buf.append(false); // Should use bit 3
827        bit_buf.append(true); // Should use bit 4
828        bit_buf.append(true); // Should use bit 5
829
830        assert_eq!(bit_buf.len(), 3);
831        assert!(!bit_buf.value(0));
832        assert!(bit_buf.value(1));
833        assert!(bit_buf.value(2));
834    }
835
836    #[test]
837    fn test_append_n_with_offset_crossing_boundary() {
838        // Create buffer with offset that will cross byte boundary when appending
839        let buf = BufferMut::zeroed(4);
840        let mut bit_buf = BitBufferMut::from_buffer(buf, 5, 0);
841
842        // Append enough bits to cross into next byte
843        bit_buf.append_n(true, 10); // 5 bits left in first byte, then 5 in second
844
845        assert_eq!(bit_buf.len(), 10);
846        for i in 0..10 {
847            assert!(bit_buf.value(i));
848        }
849
850        // Verify the underlying bytes
851        // Bits 5-7 of byte 0 should be set (3 bits)
852        // Bits 0-6 of byte 1 should be set (7 bits)
853        assert_eq!(bit_buf.as_slice()[0], 0b11100000);
854        assert_eq!(bit_buf.as_slice()[1], 0b01111111);
855    }
856
857    #[test]
858    fn test_truncate_with_offset() {
859        let buf = buffer_mut![0xFF, 0xFF];
860        let mut bit_buf = BitBufferMut::from_buffer(buf, 4, 12);
861
862        assert_eq!(bit_buf.len(), 12);
863
864        // Truncate to 8 bits
865        bit_buf.truncate(8);
866        assert_eq!(bit_buf.len(), 8);
867
868        // Truncate to 3 bits
869        bit_buf.truncate(3);
870        assert_eq!(bit_buf.len(), 3);
871
872        // Truncating to larger length should be no-op
873        bit_buf.truncate(10);
874        assert_eq!(bit_buf.len(), 3);
875    }
876
877    #[test]
878    fn test_capacity_with_offset() {
879        // Use exact buffer size to test capacity calculation
880        let buf = buffer_mut![0, 0, 0, 0, 0, 0, 0, 0, 0, 0]; // Exactly 10 bytes = 80 bits
881        let bit_buf = BitBufferMut::from_buffer(buf, 5, 0);
882
883        // Capacity should be at least buffer length minus offset
884        // (may be more due to allocator rounding)
885        assert!(bit_buf.capacity() >= 75);
886        // And should account for offset
887        assert_eq!(bit_buf.capacity() % 8, (80 - 5) % 8);
888    }
889
890    #[test]
891    fn test_reserve_with_offset() {
892        // Use exact buffer to test reserve
893        let buf = buffer_mut![0, 0]; // Exactly 2 bytes = 16 bits
894        let mut bit_buf = BitBufferMut::from_buffer(buf, 3, 0);
895
896        // Current capacity should be at least 13 bits (16 - 3)
897        let initial_capacity = bit_buf.capacity();
898        assert!(initial_capacity >= 13);
899
900        // Reserve 20 more bits (need total of offset 3 + len 0 + additional 20 = 23 bits)
901        bit_buf.reserve(20);
902
903        // Should now have at least 20 bits of capacity
904        assert!(bit_buf.capacity() >= 20);
905    }
906
907    #[test]
908    fn test_freeze_with_offset() {
909        let buf = buffer_mut![0b11110000, 0b00001111];
910        let mut bit_buf = BitBufferMut::from_buffer(buf, 4, 8);
911
912        // Set some bits
913        bit_buf.set(0);
914        bit_buf.set(7);
915
916        // Freeze and verify offset is preserved
917        let frozen = bit_buf.freeze();
918        assert_eq!(frozen.offset(), 4);
919        assert_eq!(frozen.len(), 8);
920
921        // Verify values through frozen buffer
922        assert!(frozen.value(0));
923        assert!(frozen.value(7));
924    }
925
926    #[cfg_attr(miri, ignore)] // bitvec crate uses a ptr cast that Miri doesn't support
927    #[test]
928    fn test_append_buffer_with_offsets() {
929        // Create source buffer with offset
930        let source = bitbuffer![false, false, true, true, false, true];
931
932        // Create destination buffer with offset
933        let buf = BufferMut::zeroed(4);
934        let mut dest = BitBufferMut::from_buffer(buf, 3, 0);
935
936        // Append 2 initial bits
937        dest.append(true);
938        dest.append(false);
939
940        // Append the source buffer
941        dest.append_buffer(&source);
942
943        assert_eq!(dest.len(), 8);
944        assert!(dest.value(0)); // Our first append
945        assert!(!dest.value(1)); // Our second append
946        assert!(!dest.value(2)); // From source[0]
947        assert!(!dest.value(3)); // From source[1]
948        assert!(dest.value(4)); // From source[2]
949        assert!(dest.value(5)); // From source[3]
950        assert!(!dest.value(6)); // From source[4]
951        assert!(dest.value(7)); // From source[5]
952    }
953
954    #[test]
955    fn test_set_unset_unchecked_with_offset() {
956        let buf = BufferMut::zeroed(3);
957        let mut bit_buf = BitBufferMut::from_buffer(buf, 7, 10);
958
959        unsafe {
960            bit_buf.set_unchecked(0);
961            bit_buf.set_unchecked(5);
962            bit_buf.set_unchecked(9);
963        }
964
965        assert!(bit_buf.value(0));
966        assert!(bit_buf.value(5));
967        assert!(bit_buf.value(9));
968
969        unsafe {
970            bit_buf.unset_unchecked(5);
971        }
972
973        assert!(!bit_buf.value(5));
974    }
975
976    #[test]
977    fn test_value_unchecked_with_offset() {
978        let buf = buffer_mut![0b11110000, 0b00001111];
979        let bit_buf = BitBufferMut::from_buffer(buf, 4, 8);
980
981        unsafe {
982            // First 4 bits of logical buffer come from bits 4-7 of first byte (all 1s)
983            assert!(bit_buf.value_unchecked(0));
984            assert!(bit_buf.value_unchecked(3));
985
986            // Next 4 bits come from bits 0-3 of second byte (all 1s)
987            assert!(bit_buf.value_unchecked(4));
988            assert!(bit_buf.value_unchecked(7));
989        }
990    }
991
992    #[test]
993    fn test_append_alternating_with_offset() {
994        let buf = BufferMut::zeroed(4);
995        let mut bit_buf = BitBufferMut::from_buffer(buf, 2, 0);
996
997        // Append alternating pattern across byte boundaries
998        for i in 0..20 {
999            bit_buf.append(i % 2 == 0);
1000        }
1001
1002        assert_eq!(bit_buf.len(), 20);
1003        for i in 0..20 {
1004            assert_eq!(bit_buf.value(i), i % 2 == 0);
1005        }
1006    }
1007
1008    #[test]
1009    fn test_new_set_new_unset() {
1010        let set_buf = bitbuffer_mut![true; 10];
1011        let unset_buf = bitbuffer_mut![false; 10];
1012
1013        for i in 0..10 {
1014            assert!(set_buf.value(i));
1015            assert!(!unset_buf.value(i));
1016        }
1017
1018        assert_eq!(set_buf.len(), 10);
1019        assert_eq!(unset_buf.len(), 10);
1020    }
1021
1022    #[test]
1023    fn test_append_n_false_with_offset() {
1024        let buf = BufferMut::zeroed(4);
1025        let mut bit_buf = BitBufferMut::from_buffer(buf, 5, 0);
1026
1027        bit_buf.append_n(false, 15);
1028
1029        assert_eq!(bit_buf.len(), 15);
1030        for i in 0..15 {
1031            assert!(!bit_buf.value(i));
1032        }
1033    }
1034
1035    #[test]
1036    fn test_append_n_true_with_offset() {
1037        let buf = BufferMut::zeroed(4);
1038        let mut bit_buf = BitBufferMut::from_buffer(buf, 5, 0);
1039
1040        bit_buf.append_n(true, 15);
1041
1042        assert_eq!(bit_buf.len(), 15);
1043        for i in 0..15 {
1044            assert!(bit_buf.value(i));
1045        }
1046    }
1047
1048    #[test]
1049    fn test_mixed_operations_with_offset() {
1050        // Complex test combining multiple operations with offset
1051        let buf = BufferMut::zeroed(5);
1052        let mut bit_buf = BitBufferMut::from_buffer(buf, 3, 0);
1053
1054        // Append some bits
1055        bit_buf.append_n(true, 5);
1056        bit_buf.append_n(false, 3);
1057        bit_buf.append(true);
1058
1059        assert_eq!(bit_buf.len(), 9);
1060
1061        // Set and unset
1062        bit_buf.set(6); // Was false, now true
1063        bit_buf.unset(2); // Was true, now false
1064
1065        // Verify
1066        assert!(bit_buf.value(0));
1067        assert!(bit_buf.value(1));
1068        assert!(!bit_buf.value(2)); // Unset
1069        assert!(bit_buf.value(3));
1070        assert!(bit_buf.value(4));
1071        assert!(!bit_buf.value(5));
1072        assert!(bit_buf.value(6)); // Set
1073        assert!(!bit_buf.value(7));
1074        assert!(bit_buf.value(8));
1075
1076        // Truncate
1077        bit_buf.truncate(6);
1078        assert_eq!(bit_buf.len(), 6);
1079
1080        // Freeze and verify offset preserved
1081        let frozen = bit_buf.freeze();
1082        assert_eq!(frozen.offset(), 3);
1083        assert_eq!(frozen.len(), 6);
1084    }
1085
1086    #[test]
1087    fn test_from_iterator_with_incorrect_size_hint() {
1088        // This test catches a bug where FromIterator assumed the upper bound
1089        // from size_hint was accurate. The iterator contract allows the actual
1090        // count to exceed the upper bound, which could cause UB if we used
1091        // append_unchecked beyond the allocated capacity.
1092
1093        // Custom iterator that lies about its size hint.
1094        struct LyingIterator {
1095            values: Vec<bool>,
1096            index: usize,
1097        }
1098
1099        impl Iterator for LyingIterator {
1100            type Item = bool;
1101
1102            fn next(&mut self) -> Option<Self::Item> {
1103                (self.index < self.values.len()).then(|| {
1104                    let val = self.values[self.index];
1105                    self.index += 1;
1106                    val
1107                })
1108            }
1109
1110            fn size_hint(&self) -> (usize, Option<usize>) {
1111                // Deliberately return an incorrect upper bound that's smaller
1112                // than the actual number of elements we'll yield.
1113                let remaining = self.values.len() - self.index;
1114                let lower = remaining.min(5); // Correct lower bound (but capped).
1115                let upper = Some(5); // Incorrect upper bound - we actually have more!
1116                (lower, upper)
1117            }
1118        }
1119
1120        // Create an iterator that claims to have at most 5 elements but actually has 10.
1121        let lying_iter = LyingIterator {
1122            values: vec![
1123                true, false, true, false, true, false, true, false, true, false,
1124            ],
1125            index: 0,
1126        };
1127
1128        // Collect the iterator. This would cause UB in the old implementation
1129        // if it trusted the upper bound and used append_unchecked beyond capacity.
1130        let bit_buf: BitBufferMut = lying_iter.collect();
1131
1132        // Verify all 10 elements were collected correctly.
1133        assert_eq!(bit_buf.len(), 10);
1134        for i in 0..10 {
1135            assert_eq!(bit_buf.value(i), i % 2 == 0);
1136        }
1137    }
1138}